在圖片中繪制標記框，並標注標簽
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)