VOCデータのターゲットボックスを元の図にマッピングする

# 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)