【Python 实操】labelImg标注的xml格式转换为yolo的txt文件

【Python 实操】labelImg标注的xml格式转换为yolo的txt文件

背景

YOLO家族开枝散叶(尤其是YOLOv5)的广泛应用就涉及到图像的标注与文件格式转换，常用的标注工具有 LabelMe， labelImg 等。
用 labelImg 标注完后，需要将 voc 格式的 xml 转换为 yolo 格式的 txt ; 废话不多说，直接上菜。

上菜

• vocxml2yolotxt.py

# filename: vocxml2yolotxt.py
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

# 根据自己情况修改
classes = [
    "Head",
    "Helmet",
    "WorkCloth",
    "NormalCloth"]

imgspath = "."
imgsuffixlist = ['jpg', 'png', 'jpeg', 'bmp']    

# 检查图像列表
def checkimgslist(imgspath):
    imgslist = os.listdir(imgspath)
    newimgslist = []
    for imname in imgslist:
        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)

# 转换标注信息
def convert_annotation(image_id):
    # 如果找不到对应的xml，那么就直接新建一个文件即可
    print(image_id)
    
    # 直接创建文件
    out_file = open('labels/%s.txt' % image_id, 'w')
    # 找不到对应的xml文件，就直接新建一个txt文件，并退出
    if not os.path.exists('Annotations/%s.xml' % image_id):
        pass
    else:
        # UnicodeDecodeError: 'gbk' codec can't decode byte 0x80, 所以用 with open UTF-8 编码格式打开
        with open('Annotations/%s.xml' % image_id, 'r', encoding='UTF-8') as xml_file:
            tree = ET.parse(xml_file)
            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 print_classesname(classes):
    classesname = 'classes_name.txt'
    out_file = open(classesname, 'w')
    for cls in classes:
        out_file.write(str(cls)+"\n")
    out_file.close()
    

if __name__ == '__main__':
    print("In main")
    i = 0
    
    # 找到图像文件列表
    imgslist = checkimgslist(imgspath=imgspath)
    
    # 打印类别信息到文件
    print_classesname(classes=classes)
    
    # 遍历图像文件列表并转换数据
    for imname in imgslist:
        # 调试用
        # i += 1
        # if i > 5:
        #     break
        # print(imname)
        
        ## 找到文件后缀最后一个位置
        suffix = imname.split(".")[-1]
        idx = imname.index(suffix)        
        imnameid = imname[0:idx-1]
        
        # 转换生成文件
        convert_annotation(imnameid)