plt_在图片中绘制标记框

在图片中绘制标记框,并标注标签

  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)

 

posted @ 2021-12-10 11:43  Bro_Li  阅读(656)  评论(0编辑  收藏  举报