YOLO 从标注到模型训练与检测
本篇文章将带你从数据标注开始，经过数据集转换和划分，最后训练YOLO模型并进行检测。包括必要的代码示例，以及路径和文件的详细说明，以帮助你完成整个流程。
·
本篇文章将带你从数据标注开始，经过数据集转换和划分，最后训练 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}")
continue
class_id = class_name_to_id[label]
# 获取标注框的坐标点（左上角和右下角）
x1, y1 = shape['points'][0]
x2, y2 = shape['points'][1]
# 计算 YOLO 格式的中心点坐标和宽高
x_center = (x1 + x2) / 2 / image_width
y_center = (y1 + y2) / 2 / image_height
box_width = abs(x2 - x1) / image_width
box_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 文件
每个 .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/val
nc: 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 = False
for result in results:
if result.boxes:
has_defect = True
for 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 可以在较短的时间内完成对数千张图像的标注工作。
使用步骤：
- 选择标注图片的目录
- 创建缺陷列表的文件(.txt)
Chip
Burr
In_Chip
In_Burr
- 选择缺陷列表的文件
- 标注
小结
- 标注数据：使用 Labelme 标注目标并保存为 JSON 文件。
- 转换格式：将 JSON 文件转换为 YOLO 格式的 TXT 文件。
- 划分数据集：将数据集划分为训练集和验证集。
- 训练模型：使用 YOLO 模型进行训练。
- 检测结果：使用训练好的模型检测新图像并保存结果。
这篇文章详细描述了从标注数据到训练模型和检测的完整流程，包含必要的代码示例和路径说明，帮助你顺利地完成目标检测任务。
GitCode 天启AI是一款由 GitCode 团队打造的智能助手，基于先进的LLM（大语言模型）与多智能体 Agent 技术构建，致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话，还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力，真正做到“一句话，让 Al帮你完成复杂任务”。
更多推荐
11
0- 0
已为社区贡献2条内容
所有评论(0)