UA数据集：XML_遮蔽忽略区域，不遮蔽标注框，还原（无行人）

!pip install xmltodict
!pip install pybaseutils
import os
import cv2
import xmltodict
from tqdm import tqdm
from pybaseutils.maker import maker_voc
from pybaseutils import file_utils, image_utils

def read_xml2json(xml_file):
    with open(xml_file, encoding='utf-8') as fd:  # 将XML文件装载到dict里面
        content = xmltodict.parse(fd.read())
    return content

def get_objects_info(objects):
    """
    解析标注信息
    """
    rects = []
    labels = []
    targets = objects['target_list']['target']  # collections.OrderedDict
    if not isinstance(targets, list): targets = [targets]
    for data in targets:
        box = data['box']
        attribute = data['attribute']
        label = attribute['@vehicle_type']
        rect = [box['@left'], box['@top'], box['@width'], box['@height']]
        rect = [float(r) for r in rect]
        rects.append(rect)
        labels.append(label)
        #print(rect)
    bboxes = image_utils.rects2bboxes(rects)
    num = objects['@num']
    image_id = "img{:0=5d}.jpg".format(int(num))  # img00002.jpg
    return image_id, bboxes, labels


def get_ignored_region(objects):
    """获得ignored区域"""
    rects = []
    labels = []
    #print(objects)
    # 无ignore区域
    if not objects:
        return  rects, labels
    if not "box" in objects: 
        return rects, labels

    #print(objects)
    #print(len(objects))
    '''
    # 一个ignore区域 特殊处理
    if len(objects) == 1:
        label = "ignored_region"
        print(objects)
        data = objects['box']
        print(data)
        rect = [data['@left'], data['@top'], data['@width'], data['@height']]
        rect = [float(r) for r in rect]
        rects.append(rect)
        labels.append(label)
    bboxes = image_utils.rects2bboxes(rects)
    return bboxes, labels 
    '''
    for data in objects['box']:
        label = "ignored_region"
        #print(data)
        rect = [data['@left'], data['@top'], data['@width'], data['@height']]
        rect = [float(r) for r in rect]
        rects.append(rect)
        labels.append(label)
    bboxes = image_utils.rects2bboxes(rects)
    return bboxes, labels

def show_ua_detrac_dataset(image_dir, annot_dir, out_draw="", vis=False):
    print("image_dir:{}".format(image_dir))
    print("annot_dir:{}".format(annot_dir))
    xml_list = file_utils.get_files_list(annot_dir, postfix=["*.xml"])
    class_set = []
    is_in = 0
    # 进度条
    for annot_file in tqdm(xml_list):
        #print(annot_file)
        # 将xml转换为OrderedDict格式，方便解析
        annotations = read_xml2json(annot_file)
        subname = annotations['sequence']['@name']  # UA-DETRAC子目录
        # 被忽略的区域
        ignore_bboxes, ignore_labels = get_ignored_region(annotations['sequence']['ignored_region'])
        #print(ignore_bboxes, ignore_labels)
        # 遍一帧图像，获得车辆目标框
        frame_info = annotations['sequence']['frame']
        for i in range(len(frame_info)):
            image_name, bboxes, labels = get_objects_info(frame_info[i])
            #image_name = annot_file.split("\\")[-1]
            image_id = image_name.split(".")[0]
            image_file = os.path.join(image_dir, subname, image_name)
            class_set = labels + class_set
            class_set = list(set(class_set))
            if not os.path.exists(image_file):
                print("not exist:{}".format(image_file))
                continue
            image = cv2.imread(image_file)
            #print(image)
            #print(bboxes)
            #print(labels)
            # 还原ignore区域
            #image = image_utils.draw_image_bboxes_text(image, ignore_bboxes, ignore_labels,
            #                                           color=(10, 10, 10), thickness=2, fontScale=1.0)
            # 还原标注区域
            #image = image_utils.draw_image_bboxes_text(image, bboxes, labels,
            #                                           color=(255, 0, 0), thickness=1, fontScale=0.5)

            # 
            for i in range(image.shape[0]):#540
                for j in range(image.shape[1]):#960
                    is_in_ignore = 0
                    is_in_label = 0
                    for box in ignore_bboxes:
                        if i>=(box[1]-20) and i<=(box[3]+20) and j>=(box[0]) and j<=(box[2]):
                            is_in_ignore = 1
                            break
                            #continue

                    for box in bboxes:
                        if i>=(box[1]-5) and i<=(box[3]+5) and j>=(box[0]-5) and j<=(box[2]+5):
                            is_in_label = 1
                            break
                            #continue

                    if is_in_ignore == 1 and is_in_label == 0:
                        image[i][j] = 0 

            if out_draw:
                dst_file = file_utils.create_dir(out_draw, None, "{}_{}.jpg".format(subname, image_id))
                #print(dst_file)
                cv2.imwrite(dst_file, image)
            if vis:
                image_utils.cv_show_image("det", image, use_rgb=False)
    #print("class_set:{}".format(class_set))

#主程序
# 图片路径
image_dir = "UA_NoPeople/UA_MVT_Train_nopeople"
# XML标注路径
annot_dir = "UA_NoPeople/DETRAC-Train-Annotations-XML/DETRAC-Train-Annotations-XML"
print(os.path.dirname(image_dir))
#out_draw = os.path.join(os.path.dirname(image_dir), "result")
# 图片处理保存路径
out_draw = os.path.dirname(image_dir)+"/result63554"
show_ua_detrac_dataset(image_dir, annot_dir, out_draw=out_draw, vis=False)