脚本 yolov5数据集转换 txt_xml xml_txt

脚本 yolov5数据集转换 txt_xml xml_txt

txt_xml

import os
import xml.etree.ElementTree as ET
from PIL import Image
import numpy as np

# 图片文件夹，后面的/不能省
img_path = 'G:/Zhiwei/test/line_data/train/images/'

# txt文件夹，后面的/不能省
labels_path = 'G:/Zhiwei/test/line_data/train/labels/'

# xml存放的文件夹，后面的/不能省
annotations_path = 'G:/Zhiwei/test/line_data/train/labels_xml/'

labels = os.listdir(labels_path)

# 类别
classes = ["fault1","fault2","fault3"]  #类别名

# 图片的高度、宽度、深度
sh = sw = sd = 0

def write_xml(imgname, sw, sh, sd, filepath, labeldicts):
    '''
    imgname: 没有扩展名的图片名称
    '''

    # 创建Annotation根节点
    root = ET.Element('Annotation')

    # 创建filename子节点，无扩展名                 
    ET.SubElement(root, 'filename').text = str(imgname)        

    # 创建size子节点 
    sizes = ET.SubElement(root,'size')                                      
    ET.SubElement(sizes, 'width').text = str(sw)
    ET.SubElement(sizes, 'height').text = str(sh)
    ET.SubElement(sizes, 'depth').text = str(sd) 

    for labeldict in labeldicts:
        objects = ET.SubElement(root, 'object')                 
        ET.SubElement(objects, 'name').text = labeldict['name']
        ET.SubElement(objects, 'pose').text = 'Unspecified'
        ET.SubElement(objects, 'truncated').text = '0'
        ET.SubElement(objects, 'difficult').text = '0'
        bndbox = ET.SubElement(objects,'bndbox')
        ET.SubElement(bndbox, 'xmin').text = str(int(labeldict['xmin']))
        ET.SubElement(bndbox, 'ymin').text = str(int(labeldict['ymin']))
        ET.SubElement(bndbox, 'xmax').text = str(int(labeldict['xmax']))
        ET.SubElement(bndbox, 'ymax').text = str(int(labeldict['ymax']))
    tree = ET.ElementTree(root)
    tree.write(filepath, encoding='utf-8')


for label in labels:
    with open(labels_path + label, 'r') as f:
        img_id = os.path.splitext(label)[0]
        contents = f.readlines()
        labeldicts = []
        for content in contents:
            # !!!这里要看你的图片格式了，我这里是png，注意修改
            img = np.array(Image.open(img_path + label.strip('.txt') + '.jpg'))

            # 图片的高度和宽度
            sh, sw, sd = img.shape[0], img.shape[1], img.shape[2]
            content = content.strip('\n').split()
            x = float(content[1])*sw
            y = float(content[2])*sh
            w = float(content[3])*sw
            h = float(content[4])*sh

            # 坐标的转换，x_center y_center width height -> xmin ymin xmax ymax
            new_dict = {'name': classes[int(content[0])],
                        'difficult': '0',
                        'xmin': x+1-w/2,                     
                        'ymin': y+1-h/2,
                        'xmax': x+1+w/2,
                        'ymax': y+1+h/2
                        }
            labeldicts.append(new_dict)
        write_xml(img_id, sw, sh, sd, annotations_path + label.strip('.txt') + '.xml', labeldicts)

xml_txt

import os
import xml.etree.ElementTree as ET
from decimal import Decimal
 
dirpath = '/home/*/'  # 原来存放xml文件的目录

newdir = '/home/*/'  # 修改label后形成的txt目录
 
if not os.path.exists(newdir):
    os.makedirs(newdir)
 
for fp in os.listdir(dirpath):
 
    root = ET.parse(os.path.join(dirpath, fp)).getroot()
 
    xmin, ymin, xmax, ymax = 0, 0, 0, 0
    sz = root.find('size')
    width = float(sz[0].text)
    height = float(sz[1].text)
    filename = root.find('filename').text
    print(fp)
    with open(os.path.join(newdir, fp.split('.')[0] + '.txt'), 'a+') as f:
        for child in root.findall('object'):  # 找到图片中的所有框
 
            sub = child.find('bndbox')  # 找到框的标注值并进行读取
            sub_label = child.find('name')
            xmin = float(sub[0].text)
            ymin = float(sub[1].text)
            xmax = float(sub[2].text)
            ymax = float(sub[3].text)
            try:  # 转换成yolov的标签格式，需要归一化到（0-1）的范围内
                x_center = Decimal(str(round(float((xmin + xmax) / (2 * width)),6))).quantize(Decimal('0.000000'))
                y_center = Decimal(str(round(float((ymin + ymax) / (2 * height)),6))).quantize(Decimal('0.000000'))
                w = Decimal(str(round(float((xmax - xmin) / width),6))).quantize(Decimal('0.000000'))
                h = Decimal(str(round(float((ymax - ymin) / height),6))).quantize(Decimal('0.000000'))
                print(str(x_center) + ' ' + str(y_center)+ ' '+str(w)+ ' '+str(h))
                #读取需要的标签
                #if sub_label.text == 'armor':
                f.write(' '.join([str(0), str(x_center), str(y_center), str(w), str(h) + '\n']))

            except ZeroDivisionError:
                print(' width有问题')
            '''有其他标签选用
                            if sub_label.text == 'xxx':
                                f.write(' '.join([str(1), str(x_center), str(y_center), str(w), str(h) + '\n']))
                            if sub_label.text == 'xxx':
                                f.write(' '.join([str(2), str(x_center), str(y_center), str(w), str(h) + '\n']))'''
            # with open(os.path.join(newdir, fp.split('.')[0] + '.txt'), 'a+') as f:
            #     f.write(' '.join([str(2), str(x_center), str(y_center), str(w), str(h) + '\n']))