已经由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打开看坐标对不对,然后训练就行了