YOLOv5老鼠识别关键代码解读

接上篇 使用YOLOv5训练自己的数据集 --- 老鼠识别-CSDN博客

本篇我们对上篇使用到的python脚本进行解读

split_train_val.py 代码，解读注释标注在代码中

这段Python 脚本，用于将一个包含 XML 文件的数据集划分成训练集（train）、验证集（val）、测试集（test），以及一个包含训练和验证集的组合（trainval）。

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 目录下，并且会生成三个文件列表分别对应训练集、验证集和测试集，这些列表中包含了图片的完整路径。在使用该脚本之前，请确保指定的路径正确无误，且对应的文件夹和文件存在。

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()  # 关闭文件列表