前言

本文章实现yolov5的txt数据格式转xml格式,一方面共享读者直接使用,另一方面便于我能快速复制使用,以便查看或互换yolov5数据。为此,我将写下本篇文章。主要内容包含yolov5训练数据格式介绍,xml生成代码说明及txt转xml格式逻辑说明,并将所有内容附属整个源码。


一、yolov5训练数据格式介绍

我以coco数据集转yolov5的txt格式说明。

1、txt的类别对应说明

coco数据集共有80个类别,yolov5将80个类别标签为0到79,其中0表示coco的第一个类别,即person,79表示coco对应的第80个类别,即toothbrush。

coco类别如下:

['person', 'bicycle', 'car', 'motorcycle', 'airplane', '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', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']

2、txt的文件说明

一张图对应一个txt文件,其图片与txt名称相同,后缀不同。
假如一张图有7个目标需要检测,对应的txt文件有7行数据记录目标相应的信息。

3、txt文件格式

txt文件每行表示顺序为目标类别、目标中心点x、目标中心点y、目标宽、目标高,其中目标中心点与宽高需要归一化,也就是对应数据除以对应宽高。
假如一张图有2个目标,假设图像宽=1920,高=1080;
格式:类别,box中心点x坐标,中心点y坐标,box宽w,box高h;
目标信息如下:

bicycle,6080,20,40;
motorcycle,60100,80,240;

以coco类别对应转txt为:

1,60/192080/1080,20/1920,40/1080;
3,60/1920100/1080,80/1920,240/1080;

3、yolov5训练文件形式

一般是images和labels,当然你也可以自己取名字,但必须保证images与labels里面的文件名相同。
在这里插入图片描述
yolov5调用形式:

path: ../datasets/coco128  # dataset root dir
train: images/train  # train images (relative to 'path') 128 images
val: images/val  # val images (relative to 'path') 128 images
test:  # test images (optional)

以一张图转txt文件格式展示,如下:
在这里插入图片描述

以上内容来源点击这里

二、生成xml文件代码说明

生成xml文件主要三个来源内容,其一类别,其二box框,其三图像尺寸。

1、yolov5的txt读取代码

过于通常,直接上代码:

def read_txt(path):
    txt_info_lst = []
    with open(path, "r", encoding='utf-8') as f:
        for line in f:
            txt_info_lst.append(list(line.strip('\n').split()))
    txt_info_lst = np.array(txt_info_lst)
    return txt_info_lst

2、生成xml代码

product_xml函数是一张图目标生成一个对应的xml文件函数,其中boxes与codes分别为一一对应框与类别,详细 代码如下:


def product_xml(name_img, boxes, codes, img=None, wh=None):
    '''
    :param img: 以读好的图片
    :param name_img: 图片名字
    :param boxes: box为列表
    :param codes: 为列表
    :return:
    '''
    if img is not None:
        width = img.shape[0]
        height = img.shape[1]
    else:
        assert wh is not None
        width = wh[0]
        height = wh[1]
    # print('xml w:{} h:{}'.format(width,height))

    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'VOC2007'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = name_img  # 图片名字

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(width)

    node_height = SubElement(node_size, 'height')
    node_height.text = str(height)

    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '3'

    for i, code in enumerate(codes):
        box = [boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]]
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = code
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = str(int(box[0]))
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = str(int(box[1]))
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = str(int(box[2]))
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = str(int(box[3]))

    xml = tostring(node_root, pretty_print=True)  # 格式化显示,该换行的换行
    dom = parseString(xml)

    name = name_img[:-4] + '.xml'

    tree = ElementTree(node_root)

    # print('name:{},dom:{}'.format(name, dom))
    return tree, name


三、yolov5的txt文件转xml文件步骤

步骤:
1、获得该txt对应图像的高与宽,读图获得
2、提取txt文件信息
3、提取txt文件信息进行类别转换,txt数字表示转对应标签名
4、提取box信息中心点x、y与宽高w、h转图像实际尺寸
5、生成xml文件并保存

大致读取一个txt文件转对应xml文件代码如下:

# 通过图像获得图像高与宽
img = cv2.imread(img_root)
height, width = img.shape[:2]
# 读取对应txt的信息
txt_info = read_txt(txt_root)
# 以下获得txt信息,并保存labels_lst与boxes_lst中,且一一对应
labels_lst, boxes_lst = [], []
for info in txt_info:
    label_str = str(info[0])
    x, y, w, h = float(info[1]) * width, float(info[2]) * height, float(info[3]) * width, float(
        info[4]) * height
    xmin, ymin, xmax, ymax = int(x - w / 2), int(y - h / 2), int(x + w / 2), int(y + h / 2)
    labels_lst.append(label_str)
    boxes_lst.append([xmin, ymin, xmax, ymax])
# 是否转换信息
if gt_labels: # gt_labels需要和txt类别对应
    labels_lst=[gt_labels[int(lb)] for lb in labels_lst]
# 构建xml文件
if len(labels_lst) > 0:
    tree, xml_name = product_xml('img_name.jpg', boxes_lst, labels_lst, wh=[w, h])
    tree.write(os.path.join('', xml_name))

四、完整代码

以下代码将txt文件夹中的*.txt转对应*.jpg格式,但需要提供对应图像文件夹,代码会自动匹配,其完整代码如下:

import os
import cv2
from tqdm import tqdm
from lxml.etree import Element, SubElement, tostring, ElementTree
from xml.dom.minidom import parseString
import numpy as np


def build_dir(out_dir):
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    return out_dir

def get_root_lst(root, suffix='jpg', suffix_n=3):
    root_lst, name_lst = [], []

    for dir, file, names in os.walk(root):
        root_lst = root_lst + [os.path.join(dir, name) for name in names if name[-suffix_n:] == suffix]
        name_lst = name_lst + [name for name in names if name[-suffix_n:] == suffix]

    return root_lst, name_lst
def read_txt(path):
    txt_info_lst = []
    with open(path, "r", encoding='utf-8') as f:
        for line in f:
            txt_info_lst.append(list(line.strip('\n').split()))
    txt_info_lst = np.array(txt_info_lst)
    return txt_info_lst

def product_xml(name_img, boxes, codes, img=None, wh=None):
    '''
    :param img: 以读好的图片
    :param name_img: 图片名字
    :param boxes: box为列表
    :param codes: 为列表
    :return:
    '''
    if img is not None:
        width = img.shape[0]
        height = img.shape[1]
    else:
        assert wh is not None
        width = wh[0]
        height = wh[1]
    # print('xml w:{} h:{}'.format(width,height))

    node_root = Element('annotation')
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = 'VOC2007'

    node_filename = SubElement(node_root, 'filename')
    node_filename.text = name_img  # 图片名字

    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = str(width)

    node_height = SubElement(node_size, 'height')
    node_height.text = str(height)

    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '3'

    for i, code in enumerate(codes):
        box = [boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]]
        node_object = SubElement(node_root, 'object')
        node_name = SubElement(node_object, 'name')
        node_name.text = code
        node_difficult = SubElement(node_object, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_object, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = str(int(box[0]))
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = str(int(box[1]))
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = str(int(box[2]))
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = str(int(box[3]))

    xml = tostring(node_root, pretty_print=True)  # 格式化显示,该换行的换行
    dom = parseString(xml)

    name = name_img[:-4] + '.xml'

    tree = ElementTree(node_root)

    # print('name:{},dom:{}'.format(name, dom))
    return tree, name

def yolov5txt2xml(root_data, txt_root, gt_labels=None,out_dir=None):

    # 获得图像与txt的路径与名称的列表
    img_roots_lst, img_names_lst = get_root_lst(root_data, suffix='jpg', suffix_n=3)
    txt_roots_lst, txt_names_lst = get_root_lst(txt_root, suffix='txt', suffix_n=3)
    # 创建保存xml的文件
    out_dir = build_dir(out_dir) if out_dir is not None else build_dir(os.path.join(txt_root, 'out_dir_xml'))

    label_str_lst = []
    # 通过图像遍历
    for i, img_root in tqdm(enumerate(img_roots_lst)):
        # 获得图像名称,并得到对应txt名称
        img_name = img_names_lst[i]
        txt_name = img_name[:-3] + 'txt'

        if txt_name in txt_names_lst:  # 通过图像获得txt名称是否存在,存在则继续,否则不继续
            txt_index = list(txt_names_lst).index(str(txt_name))  # 获得列表txt对应索引,以便后续获得路径
            # 通过图像获得图像高与宽
            img = cv2.imread(img_root)
            height, width = img.shape[:2]
            # 读取对应txt的信息
            txt_info = read_txt(txt_roots_lst[txt_index])

            # 以下获得txt信息,并保存labels_lst与boxes_lst中,且一一对应
            labels_lst, boxes_lst = [], []
            for info in txt_info:
                label_str = str(info[0])
                if label_str not in label_str_lst:
                    label_str_lst.append(label_str)

                x, y, w, h = float(info[1]) * width, float(info[2]) * height, float(info[3]) * width, float(
                    info[4]) * height
                xmin, ymin, xmax, ymax = int(x - w / 2), int(y - h / 2), int(x + w / 2), int(y + h / 2)
                labels_lst.append(label_str)
                boxes_lst.append([xmin, ymin, xmax, ymax])
            # 是否转换信息
            if gt_labels: # gt_labels需要和txt类别对应
                labels_lst=[gt_labels[int(lb)] for lb in labels_lst]
            # 构建xml文件
            if len(labels_lst) > 0:
                tree, xml_name = product_xml(img_name, boxes_lst, labels_lst, wh=[w, h])
                tree.write(os.path.join(out_dir, xml_name))

    print('gt label:', gt_labels)
    print('txt label:', label_str_lst)
    print('save root:',out_dir)

if __name__ == '__main__':
    root_path = r'E:\project\data\voc2012_verity\images\val'
    txt_root = r'E:\project\data\voc2012_verity\labels\val'
    gt_labels =['cow','horse',  'sheep','bird', 'dog', 'cat', 'person',
               'bicycle', 'car', 'train','aeroplane','bus', 'boat','motorbike',
               'bottle', 'sofa',  'chair',  'tvmonitor',  'pottedplant',    'diningtable']

    yolov5txt2xml(root_path, txt_root,gt_labels=gt_labels)



Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐