vocxml2yolotxt.py

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import argparse
import shutil

imgsuffixlist = ['jpg', 'png', 'jpeg', 'bmp']


# 检查图像列表
def checkimgslist(imgspath):
    imgslist = os.listdir(imgspath)

    print(imgslist)
    newimgslist = []
    for imname in imgslist:

        if not os.path.exists(os.getcwd() + '/img'):
            os.makedirs(os.getcwd() + '/img')

        shutil.copy(imgspath + '/' + imname, os.getcwd() + '/img/' + imname)
        suffix = imname.split(".")[-1]
        if suffix in imgsuffixlist:
            newimgslist.append(imname)

    return newimgslist


# 转换尺寸和矩形框
def convert(size, box):
    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)

#得到class信息



# 转换标注信息
def convert_annotation(labelspath, image_id):
    # 如果找不到对应的xml，那么就直接新建一个文件即可
    # 直接创建文件
    if not os.path.exists(os.getcwd() + '/labels'):
        os.makedirs(os.getcwd() + '/labels')

    out_file = open(os.getcwd() + f'/labels/{image_id}.txt', 'w+')

    # 找不到对应的xml文件，就直接新建一个txt文件，并退出
    if not os.path.exists(f'{labelspath}/{image_id}.xml'):
        pass
    else:
        # UnicodeDecodeError: 'gbk' codec can't decode byte 0x80, 所以用 with open UTF-8 编码格式打开
        with open(f'{labelspath}/{image_id}.xml', 'r', encoding='UTF-8') as xml_file:
            tree = ET.parse(xml_file)
            print(tree)
            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'):
                cls = obj.find('name').text
                if cls not in classes:
                    continue
                cls_id = classes.index(cls)
                bbox = obj.find('bndbox')
                b = (float(bbox.find('xmin').text), float(bbox.find('xmax').text),
                     float(bbox.find('ymin').text), float(bbox.find('ymax').text))
                bb = convert((w, h), b)
                out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

    # 写完就关闭该文件
    out_file.close()



def get_all_classes(labelspath):
    name_list = []
    lll = os.listdir(labelspath)
    print(lll)
    for llllll in lll:
        # UnicodeDecodeError: 'gbk' codec can't decode byte 0x80, 所以用 with open UTF-8 编码格式打开
        with open(f'{labelspath}/{llllll}', 'r', encoding='UTF-8') as xml_file:
            tree = ET.parse(xml_file)
            print(tree)
            root = tree.getroot()
            name = [j.text for i in root.findall('object') for j in i.findall('name')]
            xml_file.close()
        for n in name:
            name_list.append(n)
    name_list = list(set(name_list))
    with open(os.getcwd() + '/classes.names', 'w+') as f:
        for ii in name_list:
            f.write(ii + "\n")
        f.close()
    return name_list


# 将类别信息写入到文件中
def print_classesname(classes):
    classesname = os.getcwd() + '/classes_name.txt'
    out_file = open(classesname, 'w+')
    for cls in classes:
        out_file.write(str(cls) + "\n")
    out_file.close()

def create_list_imgs_labels(imgspath, labelspath):
    i = os.listdir(imgspath)
    l = os.listdir(labelspath)
    with open(os.getcwd() + '/train.txt', 'w+') as f:
        for ii in i:
            f.write('data/custom/images/' + ii + "\n")
        f.close()
    with open(os.getcwd() + '/valid.txt', 'w+') as f:
        for ii in i:
            f.write('data/custom/images/' + ii + "\n")
        f.close()

if __name__ == '__main__':


    # arg_parse = argparse.ArgumentParser()
    # arg_parse.add_argument('imgspath', help='img path', type=str)
    # arg_parse.add_argument('labelspath', help='label path', type=str)
    #
    # args = vars(arg_parse.parse_args())
    # imgspath = args['imgspath']
    # labelspath = args['labelspath']

    # 根据自己情况修改
    imgspath =r'C:\Users\Administrator\PycharmProjects\pythonProject\PyTorch-YOLOv3\data\custom\images'
    labelspath =r'C:\Users\Administrator\PycharmProjects\pythonProject\PyTorch-YOLOv3\data\custom\labels'
    create_list_imgs_labels(imgspath, labelspath)
    # classes = ['person', 'bicycle', 'car', 'motorbike', 'aeroplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'sofa', 'pottedplant', 'bed', 'diningtable', 'toilet', 'tvmonitor', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
    classes = get_all_classes(labelspath)


    i = 0

    # 找到图像文件列表
    imgslist = checkimgslist(imgspath=imgspath)
    # 打印类别信息到文件
    # print_classesname(classes=classes)


    # 遍历图像文件列表并转换数据
    for imname in imgslist:
        ## 找到文件后缀最后一个位置
        suffix = imname.split(".")[-1]
        idx = imname.index(suffix)
        imnameid = imname[0:idx - 1]

        # 转换生成文件
        convert_annotation(labelspath, imnameid)