数据集 VOC转YOLO格式-编程知识

数据集 VOC转YOLO格式

一、xml转换为txt

import os.path
import xml.etree.ElementTree as ET
import os
import random
# class_names = ['palm', 'stone', 'scissor', 'awesome', 'heartB', 'OK', 'ROCK', 'one', 'swear', 'thanks', 'heartA',
#                'heartC', 'good', 'bad', 'pray', 'call', 'take_picture', 'salute']
class_names = ['menopause', 'hairball', 'broken yarn', 'hole','stains']
xmlpath = 'F:/Project_code/yolov7-main/VOCdevkit/VOC2007/Annotations/'  # 原xml路径
txtpath = 'F:/Project_code/yolov7-main/VOCdevkit/VOC2007/labels_copy/'  # 转换后txt文件存放路径
if not os.path.exists(txtpath):os.makedirs(txtpath)
files = []for root, dirs, files in os.walk(xmlpath):Nonenumber = len(files)
print(number)
i = 0
while i < number:name = files[i][0:-4]xml_name = name + ".xml"txt_name = name + ".txt"xml_file_name = xmlpath + xml_nametxt_file_name = txtpath + txt_namexml_file = open(xml_file_name)tree = ET.parse(xml_file)root = tree.getroot()filename = root.find('filename').textimage_name = root.find('filename').textw = int(root.find('size').find('width').text)h = int(root.find('size').find('height').text)f_txt = open(txt_file_name, 'w+')content = ""first = Truefor obj in root.iter('object'):name = obj.find('name').textclass_num = class_names.index(name)xmlbox = obj.find('bndbox')x1 = int(xmlbox.find('xmin').text)x2 = int(xmlbox.find('xmax').text)y1 = int(xmlbox.find('ymin').text)y2 = int(xmlbox.find('ymax').text)if first:content += str(class_num) + " " + \str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \str((x2 - x1) / w) + " " + str((y2 - y1) / h)first = Falseelse:content += "\n" + \str(class_num) + " " + \str((x1 + x2) / 2 / w) + " " + str((y1 + y2) / 2 / h) + " " + \str((x2 - x1) / w) + " " + str((y2 - y1) / h)# print(str(i / (number - 1) * 100) + "%\n")print(content)f_txt.write(content)f_txt.close()xml_file.close()i += 1print("done!")

二、数据集划分

划分为训练集、验证集和测试集

import os, shutil
from sklearn.model_selection import train_test_splitval_size = 0.1
test_size = 0.2
postfix = 'jpg'
imgpath = 'E:/pycharm/yolov7/yolov7/data/hongwai/images/train'
txtpath = 'E:/pycharm/yolov7/yolov7/data/hongwai/labels/'os.makedirs('images/train', exist_ok=True)
os.makedirs('images/val', exist_ok=True)
os.makedirs('images/test', exist_ok=True)
os.makedirs('labels/train', exist_ok=True)
os.makedirs('labels/val', exist_ok=True)
os.makedirs('labels/test', exist_ok=True)listdir = os.listdir(txtpath)
train, test = train_test_split(listdir, test_size=test_size, shuffle=True, random_state=0)
train, val = train_test_split(train, test_size=val_size, shuffle=True, random_state=0)for i in train:shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'Images/train/{}.{}'.format(i[:-4], postfix))shutil.copy('{}/{}'.format(txtpath, i), 'labels/train/{}'.format(i))for i in val:shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'Images/val/{}.{}'.format(i[:-4], postfix))shutil.copy('{}/{}'.format(txtpath, i), 'labels/val/{}'.format(i))for i in test:shutil.copy('{}/{}.{}'.format(imgpath, i[:-4], postfix), 'Images/test/{}.{}'.format(i[:-4], postfix))shutil.copy('{}/{}'.format(txtpath, i), 'labels/test/{}'.format(i))