统计xml文件包含的标注信息 parse_dataset_annotation

雨点打透心脏的1/2处 2023-02-14 13:01 23阅读 0赞

## **YOLO**算法，统计标注的`xml`文件中包含的**标记框**信息 ##

创建`parse_dataset_annotation.py`

import os
    import sys
    filedir = os.path.dirname(sys.argv[0])      #获取脚本所在目录
    os.chdir(filedir)       #将脚本所在的目录设置为工作目录
    wdir = os.getcwd()
    print('当前工作目录：{}\n'.format(wdir))      #打印当前工作目录
    
    from xml.dom.minidom import parse
    
    def xml_parser( xml_file ):
        '''
        Parse an xml file and return the annotation info in the file
    
        :param xml_file: the xml file name to be parsed
        :return: file_name, width, height, objects.
            file_name, filename of the xml file (without extension)
            width, width of the annotated image
            height, height of the annotated image
            objects, annotated objects in the image
            object, (object_name, xmin, ymin, xmax, ymax)
                object_name, name of the annotated object
                xmin, ymin, xmax, ymax, coordinate of the bounding box of the object
    
        '''
        DOMTree = parse( xml_file )
        collection = DOMTree.documentElement #得到xml文件的根节点
        file_name_xml = collection.getElementsByTagName( 'filename' )[0]
        objects_xml = collection.getElementsByTagName( 'object' )
        size_xml = collection.getElementsByTagName( 'size' )
    
        file_name = file_name_xml.childNodes[0].data
    
        for size in size_xml:
            width = size.getElementsByTagName( 'width' )[0]
            height = size.getElementsByTagName( 'height' )[0]
    
            width = width.childNodes[0].data
            height = height.childNodes[0].data
    
        objects = []
        for object_xml in objects_xml:
            object_name = object_xml.getElementsByTagName( 'name' )[0]
            bdbox = object_xml.getElementsByTagName( 'bndbox' )[0]
            xmin = bdbox.getElementsByTagName( 'xmin' )[0]
            ymin = bdbox.getElementsByTagName( 'ymin' )[0]
            xmax = bdbox.getElementsByTagName( 'xmax' )[0]
            ymax = bdbox.getElementsByTagName( 'ymax' )[0]
    
            object = [ object_name.childNodes[0].data,
                       float(xmin.childNodes[0].data),
                       float(ymin.childNodes[0].data),
                       float(xmax.childNodes[0].data),
                       float(ymax.childNodes[0].data) ]
    
            objects.append( object )
    
        return file_name, int(width), int(height), objects
        
    image_dir = 'images'    
    xml_dir = 'labels'
    xml_files = os.listdir(xml_dir)
    image_files = os.listdir(image_dir)
    image_ext = image_files[0].split('.')[-1] #图片文件的扩展名
    print(image_ext)
    if len(image_files) == len(xml_files):
        print('共有{:d}个xml文件。'.format(len(xml_files)))
    else:
        print('图片数量和xml文件数量不一致。')
    obj_dict = {}
    for xml_file in xml_files:
        annotation = xml_parser(os.path.join(xml_dir, xml_file))
        name_1 = xml_file.split('.')[0] + '.' + image_ext.lower()
        name_2 = xml_file.split('.')[0] + '.' + image_ext.upper()
        if  name_1 not in image_files and name_2 not in image_files:
            print('{:s}没有对应的图片。'.format(xml_file))
        for obj in annotation[-1]:
            key = obj[0]
            x = (obj[1] + obj[3])/2
            y = (obj[2] + obj[4])/2
            width = obj[3] - obj[1]
            height = obj[4] - obj[2]
            box = [x,y,width,height]
            if key in obj_dict:
                obj_dict[key][0] += 1
                n = obj_dict[key][0]
                obj_dict[key][1:5] = [ (i*(n-1)+j)/(n) for i,j in zip(obj_dict[key][1:5] , box)]
                #obj_dict[key][5:9] = [ i if i>=j else j for i,j in zip(obj_dict[key][5:9] , box)]
                #obj_dict[key][9:] = [ i if i>=j else j for i,j in zip(obj_dict[key][9:] , box)]
            else:
                obj_dict[key] = []
                obj_dict[key].append(1)  # 0,个数
                obj_dict[key] += box     # 1-4， 平均坐标
                #obj_dict[key] += box     # 5-8， 最大值
                #obj_dict[key] += box     # 9-12，最小值
    for key,value in obj_dict.items():
        print('一共有 {:4d} 个 {:20s}，其边框平均位置为{:4.0f} *{:4.0f}；平均尺寸为{:3.0f} *{:3.0f}。'.format(value[0],key,*value[1:]))