一、复现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
2.运行train.py
建议:worker默认值修改为0。
利用Pycharm运行时出现闪退情况:IDE error occured。
解决方案:IDEA 性能优化设置。
打开vmoptions,做出如下修改。