目标检测标签格式转换 txt

目标检测标签格式转换 txt

目标检测标签格式转换 txt->xml

将YOLO格式的标签转换成VOC格式的标签，用于适应多种算法的目标检测

# _*_ coding:utf-8 _*_
ElementTree as ET
import os
import cv2
"""
标签格式转换：txt->xml
时间：2022.04.15
作者：余小晖
"""# 增加换行符
def __indent(elem, level=0):i = "n" + level*"t"if len(elem):if  or strip(): = i + "t"if not elem.tail or not elem.tail.strip():elem.tail = ifor elem in elem:__indent(elem, level+1)if not elem.tail or not elem.tail.strip():elem.tail = ielse:if level and (not elem.tail or not elem.tail.strip()):elem.tail = idef txt_xml(img_size, txt_file_path):# object-class> <x_center> <y_center> <width> <height>bboxs = []  # cls_name, xmin, ymin, xmax, ymaxwidth, 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 bboxsdef 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、depthimage = cv2.imread(img_path+ img_name)img_size = image.shapetxt_file_path = txt_path+ txt_filebboxs = txt_xml(img_size, txt_file_path)root = ET.Element('annotation')tree = ET.ElementTree(root)# 添加root子节点filenamefilename = ET.Element('filename') = img_nameroot.append(filename)#添加root子节点sizesize = ET.Element('size')img_width = ET.Element('width') = str(img_size[1])size.append(img_width)img_height = ET.Element('height') = str(img_size[0])size.append(img_height)img_depth = ET.Element('depth') = str(img_size[-1])size.append(img_depth)root.append(size)# 添加root子节点objectfor i in range(len(bboxs)):#添加所有的object子节点object = ET.Element('object')cls_name = ET.Element('name') = cls_dict[str(bboxs[i][0])]object.append(cls_name)difficult = ET.Element('difficult') = 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( = str(bboxs[i][j+1])bndbox.append(ele)object.append(bndbox)root.append(object)__indent(root)  # 增加换行符tree.write(xml_save_path+ xml_name, encoding='utf-8', xml_declaration=True)print('num------------------------',num+1,'==',txt_file)if __name__ == '__main__':txt_path= './face_mask-1/ann_txt/'  #txt标签路径xml_save_path = './face_mask-1/ann_xml-1/'  #xml文件保存路径img_path = './face_mask-1/images/'  #图像路径if not ists(xml_save_path):os.makedirs(xml_save_path)cls_dict = {'0': 'face', '1': 'face_mask'}  #ID以及对应类别write_xml(img_path, txt_path, xml_save_path)