当前位置：首页 > backend >正文

复现YOLOV5+训练指定数据集

backend 2025/8/6 3:27:54

一、复现YOLOV5代码

1.github下载：https://github.com/MIPIT-Team/SSA-YOLO

2.配置环境：创建虚拟环境yolo5

conda create -n yolo5 python=3.9
#对应文件夹下pip install -r requirements.txt

报错：ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
imbalanced-learn 0.12.0 requires scikit-learn>=1.0.2, which is not installed.
imbalanced-learn 0.12.0 requires threadpoolctl>=2.0.0, which is not installed.
moviepy 1.0.3 requires imageio<3.0,>=2.5; python_version >= "3.4", which is not installed.

解决方案：按照要求下载对应库文件。

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple  " scikit-learn>=1.0.2"
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple  " threadpoolctl>=2.0.0"
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple  " imageio<3.0"&&">=2.5"

3.运行train.py，增加代码：

os.environ["GIT_PYTHON_REFRESH"] = "quiet"

4.成功运行coco128数据集

二、训练指定数据集--DeepPCB

1.数据集归一化

a)明确数据集尺寸、标签信息。本文使用DeepPCB数据集，该数据集图片尺寸为640*640，标签信息为左上右下两个坐标点（绝对位置）及缺陷类型。需要将标签信息修改为YOLOV5代码格式（坐标点位置为相对位置）：缺陷类型、中心点坐标、宽、高。--缺陷类型：1-open, 2-short, 3-mousebite, 4-spur, 5-copper, 6-pin-hole。

DeepPCB数据集：tangsanli5201/DeepPCB：一个PCB缺陷数据集。（具体阅读对应的readme.txt）

根据自建数据集要求，利用trainval.txt文件生成图片images和标签labels文件夹。--配置1000张训练图。

import os
import shutildef move_files_from_txt(txt_path, img_target="./images", label_target="./labels"):"""从指定txt文件读取路径，移动对应图片和标签文件到目标文件夹参数:txt_path: 包含文件路径的txt文件img_target: 图片目标文件夹label_target: 标签目标文件夹"""# 初始化统计变量success_count = 0fail_records = []# 创建目标文件夹（不存在则创建）os.makedirs(img_target, exist_ok=True)os.makedirs(label_target, exist_ok=True)print(f"图片目标路径: {os.path.abspath(img_target)}")print(f"标签目标路径: {os.path.abspath(label_target)}")try:# 读取文件内容并处理每一行with open(txt_path, 'r', encoding='utf-8') as file:for line_num, line in enumerate(file, 1):  # 记录行号，方便定位错误line = line.strip()  # 去除首尾空白和换行符if not line:  # 跳过空行continuetry:# 分割行内容（处理可能的多空格情况）parts = line.split()if len(parts) < 2:raise ValueError("行格式不正确，至少需要两个部分")# 构建图片和标签路径img_name = f"{parts[0].split('.')[0]}_test.jpg"img_path = os.path.abspath(img_name)  # 转为绝对路径，避免相对路径问题label_path = os.path.abspath(parts[-1])# 检查文件是否存在if not os.path.exists(img_path):raise FileNotFoundError(f"图片文件不存在: {img_path}")if not os.path.exists(label_path):raise FileNotFoundError(f"标签文件不存在: {label_path}")# 复制文件shutil.copy2(img_path, os.path.join(img_target, os.path.basename(img_path).split("_")[0]+".jpg"))shutil.copy2(label_path, os.path.join(label_target, os.path.basename(label_path)))success_count += 1print(f"已移动 [{line_num}行]: {os.path.basename(img_path)} 和 {os.path.basename(label_path)}")except Exception as e:# 记录单行处理错误fail_info = f"第{line_num}行处理失败: {str(e)}"fail_records.append(fail_info)print(fail_info)except FileNotFoundError:print(f"错误: 找不到txt文件 {txt_path}")returnexcept Exception as e:print(f"读取txt文件时发生错误: {str(e)}")return# 输出最终统计print("\n" + "=" * 60)print(f"处理完成: 成功移动 {success_count} 组文件")if fail_records:print(f"处理失败 {len(fail_records)} 组文件，详情如下:")for info in fail_records:print(f"- {info}")if __name__ == "__main__":# 配置路径txt_file_path = "./trainval.txt"  # 源txt文件路径images_folder = "./images"  # 图片目标文件夹labels_folder = "./labels"  # 标签目标文件夹# 执行移动操作move_files_from_txt(txt_file_path, images_folder, labels_folder)

将标签信息更改为标准形式：类别、中心点坐标、宽、高。--相对位置

import osdef convert_to_yolo_format(input_dir, output_dir=None, image_width=640, image_height=640):"""批量转换标注文件格式为YOLO格式参数:input_dir: 包含原始txt标注文件的目录output_dir: 转换后文件的保存目录，默认与输入目录相同image_width: 图像宽度，默认640image_height: 图像高度，默认640"""# 如果未指定输出目录，则使用输入目录if output_dir is None:output_dir = input_direlse:# 创建输出目录（如果不存在）os.makedirs(output_dir, exist_ok=True)# 统计变量total_files = 0total_lines = 0error_files = []# 遍历输入目录中的所有txt文件for filename in os.listdir(input_dir):if filename.endswith('.txt'):total_files += 1input_path = os.path.join(input_dir, filename)output_path = os.path.join(output_dir, filename)try:with open(input_path, 'r', encoding='utf-8') as infile, \open(output_path, 'w', encoding='utf-8') as outfile:for line_num, line in enumerate(infile, 1):line = line.strip()if not line:  # 跳过空行continuetotal_lines += 1# 分割行内容（格式为：x1 y1 x2 y2 class）parts = line.split()if len(parts) != 5:raise ValueError(f"格式错误，预期5个值，实际{len(parts)}个值")# 解析坐标和类别x1, y1, x2, y2, cls = parts# 转换为浮点数try:x1 = float(x1)y1 = float(y1)x2 = float(x2)y2 = float(y2)except ValueError:raise ValueError("坐标值必须为数字")# 计算YOLO格式所需的参数center_x = (x1 + x2) / 2.0 / image_width  # 中心点x坐标（归一化）center_y = (y1 + y2) / 2.0 / image_height  # 中心点y坐标（归一化）width = (x2 - x1) / image_width  # 宽度（归一化）height = (y2 - y1) / image_height  # 高度（归一化）# 确保值在[0, 1]范围内center_x = max(0.0, min(1.0, center_x))center_y = max(0.0, min(1.0, center_y))width = max(0.0, min(1.0, width))height = max(0.0, min(1.0, height))# 写入转换后的数据（保留6位小数）outfile.write(f"{cls} {center_x:.6f} {center_y:.6f} {width:.6f} {height:.6f}\n")print(f"已转换: {filename}")except Exception as e:error_msg = f"处理文件 {filename} 时出错: {str(e)}"print(error_msg)error_files.append(error_msg)# 输出统计信息print("\n" + "=" * 60)print(f"转换完成: 共处理 {total_files} 个文件，{total_lines} 行标注")if error_files:print(f"处理失败 {len(error_files)} 个文件:")for err in error_files:print(f"- {err}")if __name__ == "__main__":# 配置输入输出目录input_directory = "./labels"  # 包含原始标注文件的目录output_directory = "./labels-c"  # 转换后文件的保存目录# 执行转换convert_to_yolo_format(input_dir=input_directory,output_dir=output_directory,image_width=640,image_height=640)

b)将数据集文件放在与YOLOV5文件同一级的datasets中（./datasets/DeepPCB），生成PCBData.yaml(./data)文件。

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/DeepPCB # dataset root dir
train: images/train2025 # train images (relative to 'path') 128 images
val: images/train2025 # val images (relative to 'path') 128 images
test: # test images (optional)# Classes
names:0: background1: open2: short3: mousebite4: spur5: copper6: pin-hole