【YOLOV3】目标检测中xml文件的预测生成和内容读取

最新推荐文章于 2023-03-27 09:41:11 发布

leSerein_

最新推荐文章于 2023-03-27 09:41:11 发布

阅读量1.3k

点赞数

文章标签： python 深度学习 opencv

本文链接：https://blog.csdn.net/a_piece_of_ppx/article/details/118670179

版权

源代码：https://github.com/ultralytics/yolov3

1. 通过YOLO预测结果生成xml文件

训练好的模型能够对输入进行预测。有时候我们想保留预测的结果，在YOLOV3源代码中，预测结果图片是可以保留的，那么预测的bounding box，name等信息，像训练样本一样保存在xml文件中该怎么做呢。

我在detect.py中def detect(opt)这个函数中加入了save_xml这个变量，如果是True就保存预测的xml文件。

首先定义两函数，用于写xml文件中的内容

# 定义一个创建一级分支object的函数
def create_object(root, xi, yi, xa, ya, obj_name):  # 参数依次，树根，xmin，ymin，xmax，ymax
    # 创建一级分支object
    _object = ET.SubElement(root, 'object')
    # 创建二级分支
    name = ET.SubElement(_object, 'name')
    print(obj_name)
    name.text = str(obj_name)
    pose = ET.SubElement(_object, 'pose')
    pose.text = 'Unspecified'
    truncated = ET.SubElement(_object, 'truncated')
    truncated.text = '0'
    difficult = ET.SubElement(_object, 'difficult')
    difficult.text = '0'
    # 创建bndbox
    bndbox = ET.SubElement(_object, 'bndbox')
    xmin = ET.SubElement(bndbox, 'xmin')
    xmin.text = '%s' % xi
    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s' % yi
    xmax = ET.SubElement(bndbox, 'xmax')
    xmax.text = '%s' % xa
    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s' % ya


# 创建xml文件的函数
def create_tree(image_name, h, w):
    global annotation
    # 创建树根annotation
    annotation = ET.Element('annotation')
    # 创建一级分支folder
    folder = ET.SubElement(annotation, 'folder')
    # 添加folder标签内容
    imgdir = 'images'
    folder.text = (imgdir)

    # 创建一级分支filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = image_name

    # 创建一级分支path
    path = ET.SubElement(annotation, 'path')

    path.text = getcwd() + '\{}\{}'.format(imgdir, image_name)  # 用于返回当前工作目录

    # 创建一级分支source
    source = ET.SubElement(annotation, 'source')
    # 创建source下的二级分支database
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'

    # 创建一级分支size
    size = ET.SubElement(annotation, 'size')
    # 创建size下的二级分支图像的宽、高及depth
    width = ET.SubElement(size, 'width')
    width.text = str(w)
    height = ET.SubElement(size, 'height')
    height.text = str(h)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    # 创建一级分支segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'

之后在预测结果变量detprint后，加入生成xml文件的代码

                # Print results #这是原来的
                for c in det[:, -1].unique():
                    n = (det[:, -1] == c).sum()  # detections per class
                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                #save xml   #这是自己加的
                if save_xml:
                    with open(opt.data) as f:
                        data_dict = yaml.safe_load(f)  # data dict
                    LABELS = data_dict['names']  # class names
                    (h, w) = im0.shape[:2]
                    img_name = path[path.rfind('/')+1:]
                    create_tree(img_name, h, w)

                    for coordinate in det:
                        label_id = int(coordinate[5].cpu().item())
                        create_object(annotation,
                                      int(coordinate[0].cpu().item()),
                                      int(coordinate[1].cpu().item()),
                                      int(coordinate[2].cpu().item()),
                                      int(coordinate[3].cpu().item()),
                                      LABELS[label_id])
                        # if coordinates_list==[]:
                        #     break

                    # 将树模型写入xml文件
                    tree = ET.ElementTree(annotation)
                    tree.write('/data/data_hx/yolov3/pred_annotations/{}.xml'.format(img_name.strip('.jpg')))

就可以在指定位置保存预测生成的xml文件了。

2. 读取xml文件，统计其内部信息

自己做了一个数据集，想要统计其中出现的大、中、小目标的个数。
在这里插入图片描述
也可以按bounding box对目标进行裁剪crop

def countAera(xml_fold, savecrop = False):
    '''
    xml_fold: xml存放文件夹
    savecrop: 是否用bounding box对目标进行裁剪
    '''
    files = os.listdir(xml_fold)
    dict={}
    countNum = 0
    countSmall = 0
    countMiddle = 0
    countLarge = 0
    for xmlFile in files:
        file_path = os.path.join(xml_fold, xmlFile)
        if xmlFile.endswith("jpg"):
            continue
        DOMTree = xml.dom.minidom.parse(file_path)
        collection = DOMTree.documentElement
        objects = collection.getElementsByTagName("object")
        for object in objects:
            countNum += 1
            bndbox = object.getElementsByTagName('bndbox')[0]
            xmin = bndbox.getElementsByTagName('xmin')[0]
            xmin_data = xmin.childNodes[0].data
            ymin = bndbox.getElementsByTagName('ymin')[0]
            ymin_data = ymin.childNodes[0].data
            xmax = bndbox.getElementsByTagName('xmax')[0]
            xmax_data = xmax.childNodes[0].data
            ymax = bndbox.getElementsByTagName('ymax')[0]
            ymax_data = ymax.childNodes[0].data
            xmin = int(xmin_data)
            xmax = int(xmax_data)
            ymin = int(ymin_data)
            ymax = int(ymax_data)
            if savecrop is True: #为Ture时才进行crop的操作
                # 使用imdecode读入
                img = cv2.imdecode(np.fromfile(file_path[:-4] + '.jpg', dtype=np.uint8), -1)
                #imgtmp = cv2.imdecode(np.fromfile(file_path[:-4] + '.jpg', dtype=np.uint8), cv2.IMREAD_COLOR) 和-1的效果一样
                img_cut = img[ymin:ymax, xmin:xmax, :]
                # 使用imdecode保存
                cv2.imencode('.jpg', img_cut)[1].tofile("D:/Peppa/711_2152.jpg")
            tempaera = (xmax - xmin) * (ymax - ymin)
            if tempaera <= 32 * 32:
                countSmall += 1
            elif tempaera <= 96 * 96:
                countMiddle += 1
            else:
                countLarge += 1
    print("一共%d个bounding box" % countNum)
    print("一共%d个small objects" % countSmall)
    print("一共%d个medium objects" % countMiddle)
    print("一共%d个large objects" % countLarge)