YOLO带标签的数据增强,txt转xml,xml转txt

已经由txt格式下标好的图片和labels,需要进行数据增强

首先将txt转xml

参考链接:yolov5的txt文件转xml文件格式(详细解释与完整代码供应)_yolo转xml_tangjunjun-owen的博客-CSDN博客

python代码如下:需要修改图片路径与labels路径,还有种类名称会在txt标签同目录下生成xml标签的文件夹。

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)
将带有xml标签的数据进行数据增强

参考链接:目标检测数据集数据增强python脚本_目标检测数据增强的方法资源-CSDN文库

修改代码中图片地址,xml文件地址,建立两个新文件夹来保存增强后的图片与xml文件。

(代码自行下载,不能发布他人付费代码)本人已跑通

数据增强后,xml转txt

参考链接:YOLO目标检测数据集xml格式转txt格式_xml转yolotxt资源-CSDN文库

修改代码中的字典名,也就是按顺序改成自己的类别名,然后修改xml所在文件夹地址与要保存txt的地址。

(代码自行下载,不能发布他人付费代码)本人已跑通

转完就可以了,可以使用labelimg打开看坐标对不对,然后训练就行了

  • 1
    点赞
  • 23
    收藏
    觉得还不错? 一键收藏
  • 3
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值