Generate yolo tag file and name file from the xml file of pascalvoc

Generate the yolo tag file and name file from the xml file of pascalvoc. How to
use it, copy the python file and run it in the upper directory of the xml folder

import os
import xml.etree.ElementTree as ET
import sys
'''
1.制作name.txt ok
2.制作train.txt
3.制作test.txt
'''


def convert(size, box):  # 将PascalVoc数据转化为yolo
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return x, y, w, h


def convert_annotation(image_id):
    # in_file = open(image_id.replace('\\', '/'), 'rb')
    in_file = image_id.replace('\\', '/')
    out_file = open(image_id.split('.')[0] + '.txt', 'w')

    tree = ET.parse(in_file)  # 解析xml树
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            classes.append(cls)
            # continue
        cls_id = classes.index(cls)

        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
    # in_file.close()
    out_file.close()


if __name__ == "__main__":
    classes = []

    wd = os.getcwd()  # 获取当前文件目录
    dir_list = os.listdir(wd)
    dir_xml = []
    for item in dir_list:
        if os.path.isdir(item):
            # dir.append(item)
            dir_path = os.path.join(wd, item)
            dir_list_temp = os.listdir(dir_path)
            for item2 in dir_list_temp:
                if os.path.isfile(dir_path.replace('\\', '/') + '/' + item2):
                    try:
                        if item2.split('.')[1] == 'xml':
                            dir_xml.append(dir_path + '/' + item2)  # 此处搜集完所有xml文件路径名称放在dir中
                    except IndexError:
                        pass

    for item in dir_xml:
        # *****hahaha, 这一小段是整来耍的
        sys.stdout.write('\r' + '/')
        sys.stdout.flush()
        sys.stdout.write('\r' + '|')
        sys.stdout.flush()
        sys.stdout.write('\r' + '\\')
        sys.stdout.flush()
        # print('\r'+'你的输出详情', flush=True)
        try:
            convert_annotation(item)
        except:
            pass

    with open('object.names', 'w') as f:  # 制作name文件
        for item in classes:
            f.write(item)
            f.write('\n')

Guess you like

Origin blog.csdn.net/ohhardtoname/article/details/115298471