本篇文章将带你从数据标注开始,经过数据集转换和划分,最后训练 YOLO 模型并进行检测。包括必要的代码示例,以及路径和文件的详细说明,以帮助你完成整个流程。
1. 数据标注
首先,我们需要对目标检测的数据进行标注。这里我们使用 Labelme
进行标注,生成每张图片对应的 .json
标注文件。你可以在每个目标上绘制矩形框并标注类别。
使用 Labelme 进行标注
- 安装 Labelme:
pip install labelme
- 启动 Labelme:
labelme
- 打开包含待标注图片的目录,逐张标注目标区域,保存为
.json
文件。
标注完成后,每个 .json
文件会包含图片中目标的位置、类别等信息。例如,一个典型的 JSON 文件可能包含以下内容:
{"version": "5.5.0","flags": {},"shapes": [{"label": "In_Chip","points": [[958.7826, 1375.1304],[1069.6521, 1612.087]],"group_id": null,"description": "","shape_type": "rectangle","flags": {}}],"imagePath": "5ACL470268A1FF_C_TOP_In_Chip[38,121].jpg","imageData": "已忽略...","imageHeight": 3072,"imageWidth": 2048
}
上面的内容描述了一个类别为 In_Chip
的缺陷,其矩形框由两个点定义。
2. 数据格式转换
YOLO 模型需要特定的标注格式。我们需要将 Labelme
的 JSON 标注文件转换为 YOLO 格式的 .txt
文件。
使用 labelme2yolo 工具转换
- 安装 labelme2yolo:
pip install labelme2yolo
- 使用以下命令将 JSON 文件转换为 YOLO 格式:
这将把所有的 JSON 文件转换为 YOLO 格式的labelme2yolo --json_dir D:\Temp\dataset\TEST_LABELS --out_dir D:\Temp\dataset\YOLO_FORMAT
.txt
文件,并保存到输出目录中。
代码示例:转换标注文件
如果不使用工具,可以使用以下 Python 脚本将 JSON 文件转换为 YOLO 格式的文本文件:
import json
import os# 定义类别映射
class_name_to_id = {'Chip': 0,'In_Chip': 1,'In_Burr': 2,# 根据需要添加其他类别
}# 定义输入和输出路径
input_dir = r"D:\Temp\dataset\TEST_LABELS" # JSON 标注文件所在目录
output_dir = r"D:\Temp\dataset\YOLO_FORMAT\labels" # YOLO 格式的输出目录
os.makedirs(output_dir, exist_ok=True)# 遍历输入目录中的 JSON 文件
for filename in os.listdir(input_dir):if filename.endswith('.json'):input_path = os.path.join(input_dir, filename)with open(input_path, 'r', encoding='utf-8') as f:data = json.load(f)# 获取图片的宽度和高度image_width = data['imageWidth']image_height = data['imageHeight']# 准备写入 YOLO 格式的标注内容yolo_lines = []# 解析标注信息for shape in data['shapes']:label = shape['label']if label not in class_name_to_id:print(f"跳过未知类别: {label}")continueclass_id = class_name_to_id[label]# 获取标注框的坐标点(左上角和右下角)x1, y1 = shape['points'][0]x2, y2 = shape['points'][1]# 计算 YOLO 格式的中心点坐标和宽高x_center = (x1 + x2) / 2 / image_widthy_center = (y1 + y2) / 2 / image_heightbox_width = abs(x2 - x1) / image_widthbox_height = abs(y2 - y1) / image_height# 生成 YOLO 格式的标注yolo_lines.append(f"{class_id} {x_center:.6f} {y_center:.6f} {box_width:.6f} {box_height:.6f}")# 将 YOLO 标注保存到 .txt 文件output_filename = os.path.splitext(filename)[0] + '.txt'output_path = os.path.join(output_dir, output_filename)with open(output_path, 'w') as f:f.write("\n".join(yolo_lines))print("转换完成,YOLO 格式的标注文件已保存到:", output_dir)
上述脚本会读取 JSON 文件中的目标信息,将其转换为 YOLO 格式并保存到指定目录。
YOLO__txt__127">转换后的 YOLO 格式 .txt
文件
每个 .txt
文件对应一张图片,内容格式如下:
0 0.526006 0.454710 0.062627 0.067935
1 0.412428 0.608271 0.045644 0.039629
2 0.471871 0.537506 0.100841 0.086334
每一行代表一个标注框,内容包含 class_id
(类别编号)、x_center
、y_center
、width
、height
。这些值都是相对于图片宽高的归一化坐标。
class_id
:类别编号,对应于dataset.yaml
中的类别映射。x_center
、y_center
:标注框中心的坐标。width
、height
:标注框的宽和高。
3. 数据集划分
在训练之前,我们需要将数据集划分为训练集和验证集,通常按照 8:2 的比例划分。
代码示例:划分数据集
以下代码将原始图片和标签文件划分为训练集和验证集:
import os
import shutil
import random# 定义输入目录(原始图片和标注文件的路径)
images_dir = r'D:\Temp\dataset\YOLO_FORMAT\images' # 存放所有原始图片的目录
labels_dir = r'D:\Temp\dataset\YOLO_FORMAT\labels' # 存放所有标注文件(标签)的目录# 定义输出目录(用于划分训练集和验证集)
train_images_dir = r'D:\Temp\dataset\YOLO_FORMAT\images\train' # 存放训练集图片的目录
val_images_dir = r'D:\Temp\dataset\YOLO_FORMAT\images\val' # 存放验证集图片的目录
train_labels_dir = r'D:\Temp\dataset\YOLO_FORMAT\labels\train' # 存放训练集标签文件的目录
val_labels_dir = r'D:\Temp\dataset\YOLO_FORMAT\labels\val' # 存放验证集标签文件的目录# 创建输出目录(如果不存在则创建)
os.makedirs(train_images_dir, exist_ok=True)
os.makedirs(val_images_dir, exist_ok=True)
os.makedirs(train_labels_dir, exist_ok=True)
os.makedirs(val_labels_dir, exist_ok=True)# 获取所有图片文件名
all_images = [f for f in os.listdir(images_dir) if f.endswith('.jpg')]
random.shuffle(all_images)# 划分训练集和验证集
train_ratio = 0.8 # 训练集比例为 80%
train_size = int(len(all_images) * train_ratio)
train_images = all_images[:train_size]
val_images = all_images[train_size:]# 将文件移动到相应的目录
for image_name in train_images:label_name = os.path.splitext(image_name)[0] + '.txt'shutil.move(os.path.join(images_dir, image_name), os.path.join(train_images_dir, image_name))shutil.move(os.path.join(labels_dir, label_name), os.path.join(train_labels_dir, label_name))for image_name in val_images:label_name = os.path.splitext(image_name)[0] + '.txt'shutil.move(os.path.join(images_dir, image_name), os.path.join(val_images_dir, image_name))shutil.move(os.path.join(labels_dir, label_name), os.path.join(val_labels_dir, label_name))print("数据集划分完成!")
目录结构说明
D:/
└── Temp/└── dataset/├── YOLO_FORMAT/│ ├── images/│ │ ├── train/ # 训练集图片│ │ └── val/ # 验证集图片│ ├── labels/│ │ ├── train/ # 训练集标签│ │ └── val/ # 验证集标签└── dataset.yaml # 数据集描述文件
images_dir
:存放所有原始图片的目录,这些图片是你在标注过程中使用的原始图像,转换到 YOLO 格式后仍然存放在这个目录中。labels_dir
:存放所有标注文件(标签)的目录,这些.txt
文件是你通过convert_labelme_to_yolo.py
转换得到的,包含了图片中的缺陷的位置信息和类别信息。train_images_dir
:存放训练集图片的目录。划分数据集时,脚本会将 80% 的原始图片移到这个目录中,供模型在训练阶段使用。val_images_dir
:存放验证集图片的目录。剩下 20% 的图片将被移动到这个目录,用于模型在训练过程中验证性能。train_labels_dir
:存放训练集标签文件的目录。对应于train_images_dir
中的每一张图片,它们的标注文件(即.txt
文件)会被移动到这个目录中。val_labels_dir
:存放验证集标签文件的目录。对应于val_images_dir
中的每一张图片,它们的标注文件会被移动到这个目录中。
划分数据集的目的是确保模型能够在训练时使用大部分数据进行学习(训练集),同时用一部分数据评估模型的泛化性能(验证集)。这种方式可以有效地判断模型是否有过拟合现象。
4. 数据集配置文件 dataset.yaml
为了训练 YOLO 模型,我们需要一个数据集配置文件 dataset.yaml
,它定义了训练集和验证集的路径,以及类别数量和名称:
train: D:/Temp/dataset/YOLO_FORMAT/images/train
val: D:/Temp/dataset/YOLO_FORMAT/images/valnc: 3 # 类别数量,例如 3 个缺陷类型names: ['Chip', 'In_Chip', 'In_Burr'] # 类别名称
将这个配置文件保存到 YOLO 格式数据集的根目录中。
5. 训练模型
接下来,我们使用 Ultralytics YOLO
库进行模型训练。
代码示例:训练模型
from ultralytics import YOLO# 加载预训练模型
model = YOLO('yolov8n.pt')# 开始训练
model.train(data='D:/Temp/dataset/YOLO_FORMAT/dataset.yaml', epochs=100, imgsz=640)
该代码会加载预训练的 YOLOv8 模型,并使用我们准备好的数据进行训练。
命令行训练模型
也可以使用以下命令在命令行中训练模型:
yolo detect train data=D:/Temp/dataset/YOLO_FORMAT/dataset.yaml model=yolov8n.pt epochs=100 imgsz=640
6. 使用训练好的模型进行检测
训练完成后,你会得到一个最优权重文件(如 best.pt
),然后可以使用这个文件来进行缺陷检测。
代码示例:使用训练好的模型进行检测
import os
import cv2
import time
from ultralytics import YOLO# 加载训练好的模型
model = YOLO('D:/Temp/dataset/runs/detect/train/weights/best.pt')# 定义输入和输出路径
input_folder = r'D:\BaiduNetdiskDownload\TEST'
output_folder = r'D:\BaiduNetdiskDownload\TEST_RESULTS'
os.makedirs(output_folder, exist_ok=True)# 遍历文件夹中的图片
for filename in os.listdir(input_folder):if filename.endswith('.jpg'):input_path = os.path.join(input_folder, filename)# 开始计时start_time = time.time()results = model(input_path)end_time = time.time()elapsed_time_ms = (end_time - start_time) * 1000# 加载原始图片img = cv2.imread(input_path)if img is None:print(f"无法加载图片: {input_path}")continue# 检测结果处理has_defect = Falsefor result in results:if result.boxes:has_defect = Truefor box in result.boxes:cls = result.names[int(box.cls)]conf = box.conf.item()x1, y1, x2, y2 = map(int, box.xyxy[0])color = (0, 0, 255)cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)label = f"{cls} ({x1},{y1},{x2 - x1},{y2 - y1})"cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)# 仅保存有缺陷的图片if has_defect:save_path = os.path.join(output_folder, filename)cv2.imwrite(save_path, img)print(f"图片: {filename} | 耗时: {elapsed_time_ms:.2f} 毫秒 | 已保存检测结果。")else:print(f"图片: {filename} 未检测到缺陷,跳过保存。")print(f"检测完成!检测结果已保存到: {output_folder}")
命令行检测
可以通过以下命令使用训练好的模型进行检测:
yolo detect predict model=D:/Temp/dataset/runs/detect/train/weights/best.pt source=D:/BaiduNetdiskDownload/TEST
7. yolo_label 标注工具
yolo_label可在 Linux 和 Windows 下运行,依赖 OpenCV 库,比 LabelImg 更好用。它在一些开源项目中被广泛使用,其标注速度相对较快,能够满足大部分图像检测任务数据集制作的需求。例如,在一个中等规模的数据集标注项目中,yolo_label 可以在较短的时间内完成对数千张图像的标注工作。
下载地址:developer0hye/Yolo_Label
使用步骤:
- 选择标注图片的目录
- 创建缺陷列表的文件(.txt)
Chip
Burr
In_Chip
In_Burr
- 选择缺陷列表的文件
- 标注
小结
- 标注数据:使用 Labelme 标注目标并保存为 JSON 文件。
- 转换格式:将 JSON 文件转换为 YOLO 格式的 TXT 文件。
- 划分数据集:将数据集划分为训练集和验证集。
- 训练模型:使用 YOLO 模型进行训练。
- 检测结果:使用训练好的模型检测新图像并保存结果。
这篇文章详细描述了从标注数据到训练模型和检测的完整流程,包含必要的代码示例和路径说明,帮助你顺利地完成目标检测任务。