将VOC数据中目标框映射要原图上

将标签文件放入annotations文件夹中，图片放入到image文件夹中，接着读取xml文件中的标签框，并通过cv2库读取图片并画出目标框。
代码如下：

# coding: utf-8
# author:csy
# 2021-04-23


"""
将VOC数据标签映射到原图上
"""

import cv2
import os
import time
import xml.etree.cElementTree as ET


def get_bbox(xml_path):
    tree = ET.ElementTree(file=xml_path)
    root = tree.getroot()
    object_set = root.findall('object')
    object_bbox = list()

    for Object in object_set:
        bbox = Object.find('bndbox')
        x1 = int(bbox.find('xmin').text.split('.')[0])
        y1 = int(bbox.find('ymin').text.split('.')[0])
        x2 = int(bbox.find('xmax').text.split('.')[0])
        y2 = int(bbox.find('ymax').text.split('.')[0])
        obj_bbox = [x1, y1, x2, y2]   # 若要加上标签，则这里加上物体名字即可
        object_bbox.append(obj_bbox)
    return object_bbox


def drow_object(img_file, bndboxes):

    img = cv2.imread(img_file)

    for i in range(len(bndboxes)):
        xmin = bndboxes[i][0]
        ymin = bndboxes[i][1]
        xmax = bndboxes[i][2]
        ymax = bndboxes[i][3]

        cv2.rectangle(img, (xmin, ymax), (xmax, ymin), (0, 0, 255), 2)

    cv2.imwrite('result' + str(time.time()) + '.jpg', img)


if __name__ == '__main__':
    xml_dir = 'anntations'
    img_dir = 'images'
    for file in os.listdir(xml_dir):
        file_name = file.split('.')[0]
        xml = file_name + '.xml'
        pic = file_name + '.jpg'
        bndboxes = get_bbox(xml_path=os.path.join(xml_dir, xml))

        drow_object(img_file=os.path.join(img_dir, pic), bndboxes=bndboxes)