目标检测批量 xml文件手动可视化，设置指定颜色

1. 设置原始图片路径
2. 设置xml文件路径
3. 设置保存文件路径
4. 设置不同类别对应的颜色
5. run

# -*- encoding: utf-8 -*-
# Author  : haitong
# Time    : 2023/5/25 16:47
# File    : XmlVisualization.py
# Software: PyCharm

import cv2
import xml.etree.ElementTree as ET
from PIL import Image
import numpy
import os


def XmlVisualization(data_file_dir, xml_file_dir, save_file_dir, bbox_color_dict):
    for img_name in os.listdir(data_file_dir):
        if img_name.endswith('.jpg') or img_name.endswith('.png'):
            print(img_name)
            img_path = os.path.join(data_file_dir, img_name)
            xml_name = img_name[:-3] + "xml"
            xml_path = os.path.join(xml_file_dir, xml_name)

            # img_bgr = cv2.imread(img_path)
            img_bgr = Image.open(img_path)
            img_bgr = cv2.cvtColor(numpy.asarray(img_bgr), cv2.COLOR_RGB2BGR)

            xml_inf = open(xml_path, encoding='utf-8')
            tree = ET.parse(xml_inf)
            root = tree.getroot()

            for obj in root.iter('object'):
                bbox_label = obj.find('name').text

                bbox_top_left_x = int(obj.find('bndbox').find('xmin').text)
                bbox_top_left_y = int(obj.find('bndbox').find('ymin').text)
                bbox_bottom_right_x = int(obj.find('bndbox').find('xmax').text)
                bbox_bottom_right_y = int(obj.find('bndbox').find('ymax').text)

                img_bgr = cv2.rectangle(img_bgr, (bbox_top_left_x, bbox_top_left_y), (bbox_bottom_right_x, bbox_bottom_right_y), bbox_color_dict[bbox_label], 3)
                img_bgr = cv2.putText(img_bgr, bbox_label, (bbox_top_left_x + 2, bbox_top_left_y + -20), cv2.FONT_HERSHEY_SIMPLEX, 2, bbox_color_dict[bbox_label], 2)

            # cv2.imwrite(os.path.join(save_file_dir, img_name), img_bgr)
            image = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
            image.save(os.path.join(save_file_dir, img_name))

    return 1


if __name__ == "__main__":
    dir_path = os.path.dirname(os.path.abspath(__file__))
    data_file_dir = os.path.join(dir_path, 'img')  					        # 原始图片路径
    xml_file_dir = os.path.join(dir_path, 'xml')  							# xml路径
    save_file_dir =  os.path.join(dir_path, 'save')						    # 保存路径
    bbox_color_dict = {'breakage': (0, 0, 255), 'corrosion': (0, 255, 0)}   # 设置对应检测目标的颜色

    _ = XmlVisualization(data_file_dir, xml_file_dir, save_file_dir, bbox_color_dict)