【yolo格式标签转VOC格式】

yolo格式标签转VOC格式

yolo格式标签转VOC格式

txt–> xml

# _*_ coding:utf-8 _*_
import xml.etree.ElementTree as ET
import os
import cv2
"""
标签格式转换：txt->xml
时间：2022.04.15
作者：yuxiaohui
"""

# 增加换行符
def __indent(elem, level=0):
    i = "\n" + level*"\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            __indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


def txt_xml(img_size, txt_file_path):
    # object-class> <x_center> <y_center> <width> <height>
    bboxs = []  # cls_name, xmin, ymin, xmax, ymax
    width, height= img_size[1], img_size[0]     # 获取对应图像的宽高
    with open(txt_file_path, 'r', encoding='UTF-8')as f:
        lines = f.readlines()   # 可能含有空的txt文件,空txt文件将直接返回空bboxs列表
        # 读取txt文件中每一行数据，每一行数据包含有一个目标
        for line in lines :
            line= line.split()
            line = [float(i) for i in line ]
            cls_name= int(line[0])
            xmin = int((line[1]- line[3]/2)* width)
            xmax = int((line[1]+ line[3]/2)* width)
            ymin = int((line[2]- line[4]/2)* height)
            ymax = int((line[2]+ line[4]/2)* height)
            bbox = [cls_name, xmin, ymin, xmax, ymax]
            bboxs.append(bbox)
    return bboxs


def write_xml(img_path, txt_path, xml_save_path):
    txt_files = os.listdir(txt_path)
    # txt_files.sort(key=lambda x: int(x.split('_')[-1][:-4]))  # 按照数字排序
    for num, txt_file in enumerate(txt_files):
        txt_name_list = txt_file.split('.')
        xml_name = txt_name_list[0]+ '.xml'
        img_name = txt_name_list[0]+ '.jpg' # 图片后缀为jpg时
        # 获取图像尺寸width、height、depth
        image = cv2.imread(os.path.join(img_path, img_name))
        img_size = image.shape
        txt_file_path = os.path.join(txt_path, txt_file)
        bboxs = txt_xml(img_size, txt_file_path)
        root = ET.Element('annotation')
        tree = ET.ElementTree(root)
        # 添加root子节点filename
        filename = ET.Element('filename')
        filename.text = img_name
        root.append(filename)
        #添加root子节点size
        size = ET.Element('size')
        img_width = ET.Element('width')
        img_width.text = str(img_size[1])
        size.append(img_width)
        img_height = ET.Element('height')
        img_height.text = str(img_size[0])
        size.append(img_height)
        img_depth = ET.Element('depth')
        img_depth.text = str(img_size[-1])
        size.append(img_depth)
        root.append(size)
        # 添加root子节点object
        for i in range(len(bboxs)):
            #添加所有的object子节点
            object = ET.Element('object')
            cls_name = ET.Element('name')
            cls_name.text = cls_dict[str(bboxs[i][0])]
            object.append(cls_name)
            difficult = ET.Element('difficult')
            difficult.text = str(0)
            object.append(difficult)
            bndbox = ET.Element('bndbox')
            element_list = ['xmin','ymin', 'xmax', 'ymax']
            for j, element in enumerate(element_list):
                ele = ET.Element(element)
                ele.text = str(bboxs[i][j+1])
                bndbox.append(ele)
            object.append(bndbox)
            root.append(object)
        __indent(root)  # 增加换行符
        tree.write(os.path.join(xml_save_path+ xml_name), encoding='utf-8', xml_declaration=True)
        print('num------------------------',num+1,'==',txt_file)


if __name__ == '__main__':
    txt_path= r'E:\gz_download\ultralytics-main-yolov8-smoking_dataset\ultralytics-main-yolov8-smoking_dataset\smoking_dataset/valid/labels/'  #txt标签路径
    xml_save_path = r'E:\gz_download\ultralytics-main-yolov8-smoking_dataset\ultralytics-main-yolov8-smoking_dataset\smoking_dataset/valid/labels_xml/'  #xml文件保存路径
    img_path = r'E:\gz_dataset_download\ultralytics-main-yolov8-smoking_dataset\ultralytics-main-yolov8-smoking_dataset\smoking_dataset/valid/images/'  #图像路径
    if not os.path.exists(xml_save_path):
        os.makedirs(xml_save_path)
    cls_dict = {'0': 'smoking',}  #ID以及对应类别
    write_xml(img_path, txt_path, xml_save_path)