在图片中绘制标记框，并标注标签
1 # 2 import os 3 import xml.dom.minidom 4 import matplotlib.pyplot as plt 5 from matplotlib.image import imread 6 import matplotlib.patches as patches 7 8 # 定义画矩形框的函数 9 def draw_rectangle(currentAxis, bbox, edgecolor='y', facecolor='r', fill=False, linestyle='-'): 10 # 坐标格式为 x y w h 11 rect = patches.Rectangle((bbox[0], bbox[1]), bbox[2], bbox[3], linewidth=1, 12 edgecolor=edgecolor, facecolor=facecolor, fill=fill, linestyle=linestyle) 13 currentAxis.add_patch(rect) 14 15 # 自定义函数，输入图像和 gtbox 16 def draw_bbox(img_path, bboxes, img_save_path): 17 img = imread(img_path) 18 plt.figure(num=1) # 使用同一张画布，最终只会展示最后一张图片 19 plt.axis('off') 20 plt.imshow(img) 21 currentAxis = plt.gca() 22 # 绘制矩形框 23 for box in bboxes: 24 # print(img_path.split('\\')[-1], box) 25 box_name = box[0] 26 my_box = box[1] 27 temp_box = [int( my_box[0] ), 28 int( my_box[1] ), 29 int( my_box[2]-my_box[0] ), 30 int( my_box[3]-my_box[1] ) 31 ] 32 draw_rectangle(currentAxis, temp_box) 33 34 ### 标注标记框对应的标签 35 plt.text(my_box[0]+3, my_box[1]+13, box_name, fontsize=3, color='yellow') 36 37 # 挨个画，最终在图片上一起展示 38 plt.savefig(img_save_path, bbox_inches='tight', pad_inches=0, dpi=500) 39 # plt.show() 40 plt.close() 41 42 # 返回 xml文件中的标签名，和标记框数据 43 def get_bbox(xml_path): 44 box_list = [] 45 # 打开xml文档 46 DOMTree = xml.dom.minidom.parse(xml_path) 47 # 得到文档元素对象 48 collection = DOMTree.documentElement 49 ### 获取文件名 50 filenamelist = collection.getElementsByTagName("filename") 51 filename = filenamelist[0].childNodes[0].data 52 # print('\n', len(filenamelist), filename) 53 ### 得到标签名为object的信息 54 objectlist = collection.getElementsByTagName("object") 55 for objects in objectlist: 56 ### 每个 object 中得到子标签名为 name 的信息 57 namelist = objects.getElementsByTagName('name') 58 ### 获得标记框的标签名 59 objectname = namelist[0].childNodes[0].data 60 # print('类别名为： ', objectname) ########### 索要统计的信息 61 62 bndbox = objects.getElementsByTagName('bndbox') 63 for box in bndbox: 64 x1_list = box.getElementsByTagName('xmin') 65 x1 = int(x1_list[0].childNodes[0].data) 66 y1_list = box.getElementsByTagName('ymin') 67 y1 = int(y1_list[0].childNodes[0].data) 68 x2_list = box.getElementsByTagName('xmax') # 注意坐标，看是否需要转换 69 x2 = int(x2_list[0].childNodes[0].data) 70 y2_list = box.getElementsByTagName('ymax') 71 y2 = int(y2_list[0].childNodes[0].data) 72 bbox = [x1, y1, x2, y2] 73 74 # print('文件名：', filename, ' ，标签名：', objectname, ' ，标记框：', bbox, ' ', len(bndbox)) 75 box_list.append([objectname, bbox]) 76 return box_list 77 78 79 def xml_label_names(xml_root_path, data_root_path, img_save_root_path): 80 xml_files = os.listdir(xml_root_path) 81 file_have_box_count = 0 82 file_count = len(xml_files) 83 for xml_file in xml_files: 84 xml_path = os.path.join(xml_root_path, xml_file) 85 box_list = get_bbox(xml_path) 86 img_name = xml_file.split('.')[0] + '.jpg' 87 print('图片名称：', img_name, ' , 标记框数量： ', len(box_list)) 88 image_path = os.path.join(data_root_path, img_name) 89 ### 出入图片路径和 bbox 信息，在图片中绘制 标记框 90 if len(box_list)>0: 91 print() 92 file_have_box_count = file_have_box_count + 1 93 94 ### 传入图像和标记框信息，在图像中绘制标记框 95 img_save_path = os.path.join(img_save_root_path, img_name) 96 draw_bbox(image_path, box_list, img_save_path) 97 # print(xml_file, len(box_list)) 98 99 100 101 print('拥有标记框的图片数量：', file_have_box_count, ' ; 总文件数量：', file_count) 102 103 ### 路径构成 104 folder_name = ['Czech','India','Japan'] 105 root_path = r' 你存储数据集的根路径 \8_JapanRoad_RDD200\train' # 8_JapanRoad_RDD200 的根路径 106 107 for temp_name in folder_name: 108 xml_root_path = os.path.join(root_path, temp_name, r'annotations\xmls') 109 data_root_path = os.path.join(root_path, temp_name, r'images') 110 img_save_root_path = os.path.join(root_path, 'temp_labeled', temp_name) 111 xml_label_names(xml_root_path, data_root_path, img_save_root_path)