接上篇 使用YOLOv5训练自己的数据集 --- 老鼠识别-CSDN博客
本篇我们对上篇使用到的python脚本进行解读
split_train_val.py 代码，解读注释标注在代码中
这段Python 脚本，用于将一个包含 XML 文件的数据集划分成训练集（train）、验证集（val）、测试集（test），以及一个包含训练和验证集的组合（trainval）。文章来源：https://www.toymoban.com/news/detail-781524.html
import os
import random
import argparse
# 创建一个参数解析器
parser = argparse.ArgumentParser()
# 添加XML文件夹路径参数，默认值是你的数据集的Annotations文件夹路径
parser.add_argument('--xml_path', default='D:/yolov5-master/Y1/Annotations', type=str, help='input xml file path')
# 添加文本文件夹路径参数，默认值是你的数据集的ImageSets/Main文件夹路径
parser.add_argument('--txt_path', default='D:/yolov5-master/Y1/ImageSets/Main', type=str, help='output txt file path')
opt = parser.parse_args()
# 设置训练和验证数据集比例
trainval_percent = 0.9
# 设置训练数据集在训练和验证数据集中的比例
train_percent = 8/9
# 获取XML文件路径
xmlfilepath = opt.xml_path
# 获取TXT文件保存路径
txtsavepath = opt.txt_path
# 列出所有XML文件
total_xml = os.listdir(xmlfilepath)
# 如果TXT文件保存路径不存在，创建它
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
# 计算文件总数
num = len(total_xml)
# 计算训练和验证数据集数量
tv = int(num * trainval_percent)
# 计算训练数据集数量
tr = int(tv * train_percent)
# 随机选择训练验证数据集
trainval = random.sample(range(num), tv)
# 在训练验证数据集中随机选择训练数据集
train = random.sample(trainval, tr)
# 创建文件以保存训练验证集、测试集、训练集和验证集的文件名
file_trainval = open(os.path.join(txtsavepath, 'trainval.txt'), 'w')
file_test = open(os.path.join(txtsavepath, 'test.txt'), 'w')
file_train = open(os.path.join(txtsavepath, 'train.txt'), 'w')
file_val = open(os.path.join(txtsavepath, 'val.txt'), 'w')
# 遍历每个文件，决定它应该属于哪个集合
for i in range(num):
name = total_xml[i][:-4] + '\n' # 去掉.xml扩展名
if i in trainval:
file_trainval.write(name) # 写入trainval.txt
if i in train:
file_train.write(name) # 写入train.txt
else:
file_val.write(name) # 写入val.txt
else:
file_test.write(name) # 写入test.txt
# 关闭所有打开的文件
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
voc_label.py 代码，解读注释标注在代码中
这段代码用于将VOC格式的标注文件（XML格式）转换成YOLO格式的标注文件（TXT格式）。并生成用于训练、验证和测试的图片路径列表。代码假设你的图片和XML标注文件都已经准备好，并存放在指定的目录下。执行这个脚本后，YOLO格式的标注文件将会被创建在指定的 labels 目录下，并且会生成三个文件列表分别对应训练集、验证集和测试集，这些列表中包含了图片的完整路径。在使用该脚本之前，请确保指定的路径正确无误，且对应的文件夹和文件存在。文章来源地址https://www.toymoban.com/news/detail-781524.html
import xml.etree.ElementTree as ET # 用于解析XML文件
import os
from os import getcwd # 获取当前工作目录的函数
sets = ['train', 'val', 'test'] # 定义数据集的子集
classes = ["mouse"] # 定义类别列表，应根据实际情况修改
abs_path = os.getcwd() # 获取并打印当前工作目录的绝对路径
print(abs_path)
# 定义函数转换坐标
def convert(size, box):
dw = 1. / (size[0])
dh = 1. / (size[1])
x = (box[0] + box[1]) / 2.0 - 1
y = (box[2] + box[3]) / 2.0 - 1
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):
in_file = open('D:/yolov5-master/Y1/Annotations/%s.xml' % (image_id), encoding='UTF-8') # 打开XML文件
out_file = open('D:/yolov5-master/Y1/labels/%s.txt' % (image_id), 'w') # 创建TXT文件
tree = ET.parse(in_file) # 解析XML文件
root = tree.getroot() # 获取XML的根节点
filename = root.find('filename').text # 获取文件名
filenameFormat = filename.split(".")[1] # 提取文件格式
size = root.find('size') # 获取图片尺寸
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'): # 遍历所有对象
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1: # 忽略不在类别列表中或标记为难以识别的对象
continue
cls_id = classes.index(cls) # 获取类别索引
xmlbox = obj.find('bndbox') # 获取边界框
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text),
float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
# 对越界的坐标进行修正
b = (max(min(b[0], w), 0), max(min(b[1], w), 0),
max(min(b[2], h), 0), max(min(b[3], h), 0))
bb = convert((w, h), b) # 转换坐标
# 写入类别索引和归一化的边界框坐标
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
in_file.close() # 关闭XML文件
out_file.close() # 关闭TXT文件
return filenameFormat # 返回文件格式
wd = getcwd() # 获取当前工作目录
# 为每个子集创建标注文件和文件列表
for image_set in sets:
if not os.path.exists('D:/yolov5-master/Y1/labels/'): # 如果labels目录不存在，创建它
os.makedirs('D:/yolov5-master/Y1/labels/')
image_ids = open('D:/yolov5-master/Y1/ImageSets/Main/%s.txt' % (image_set)).read().strip().split() # 读取图片ID列表
list_file = open('D:/yolov5-master/Y1/%s.txt' % (image_set), 'w') # 创建对应的文件列表
for image_id in image_ids:
filenameFormat = convert_annotation(image_id) # 转换标注文件
list_file.write('D:/yolov5-master/Y1/images/%s.%s\n' % (image_id, filenameFormat)) # 将图片路径写入文件列表
list_file.close() # 关闭文件列表
到了这里，关于YOLOv5老鼠识别关键代码解读的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
