YOLO 从标注到模型训练与检测

news/2024/11/26 9:49:33/

本篇文章将带你从数据标注开始,经过数据集转换和划分,最后训练 YOLO 模型并进行检测。包括必要的代码示例,以及路径和文件的详细说明,以帮助你完成整个流程。

1. 数据标注

首先,我们需要对目标检测的数据进行标注。这里我们使用 Labelme 进行标注,生成每张图片对应的 .json 标注文件。你可以在每个目标上绘制矩形框并标注类别。

使用 Labelme 进行标注
  1. 安装 Labelme:
    pip install labelme
    
  2. 启动 Labelme:
    labelme
    
  3. 打开包含待标注图片的目录,逐张标注目标区域,保存为 .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 工具转换
  1. 安装 labelme2yolo:
    pip install labelme2yolo
    
  2. 使用以下命令将 JSON 文件转换为 YOLO 格式:
    labelme2yolo --json_dir D:\Temp\dataset\TEST_LABELS --out_dir D:\Temp\dataset\YOLO_FORMAT
    
    这将把所有的 JSON 文件转换为 YOLO 格式的 .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_centery_centerwidthheight。这些值都是相对于图片宽高的归一化坐标。

  • class_id:类别编号,对应于 dataset.yaml 中的类别映射。
  • x_centery_center:标注框中心的坐标。
  • widthheight:标注框的宽和高。

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

使用步骤:

  1. 选择标注图片的目录
  2. 创建缺陷列表的文件(.txt)

Chip
Burr
In_Chip
In_Burr

  1. 选择缺陷列表的文件
  2. 标注
    在这里插入图片描述

小结

  1. 标注数据:使用 Labelme 标注目标并保存为 JSON 文件。
  2. 转换格式:将 JSON 文件转换为 YOLO 格式的 TXT 文件。
  3. 划分数据集:将数据集划分为训练集和验证集。
  4. 训练模型:使用 YOLO 模型进行训练。
  5. 检测结果:使用训练好的模型检测新图像并保存结果。

这篇文章详细描述了从标注数据到训练模型和检测的完整流程,包含必要的代码示例和路径说明,帮助你顺利地完成目标检测任务。


http://www.ppmy.cn/news/1550024.html

相关文章

lvgl学习复选框部件和进度条部件(基于正点原子)

复选框部件(lv_checkbox) 复选框部件常用于选择某个内容的开启和关闭,可以理解为自带标签的开关。 复选框部件组成部分: 主体(LV_PART_MAIN) 勾选框(LV PART INDICATOR) 知识点1:创建复选框部件 lv_obj_t *check…

lambda的作用

lambda 的定义 lambda 是 Python 中用于创建匿名函数的关键字。匿名函数是一种没有名字的函数,通常用来定义简单的、一次性的函数。 lambda 的语法 lambda 参数列表: 表达式 参数列表: 函数的输入,可以有多个,用逗号分隔。表达式: 函数的…

前端高能组件库 Shadcn-UI

你是不是用 element-ui 或者 ant-design ,然后,开发时常常遇到需要匹配设计稿时调样式的痛苦。 Shadcn-UI 结合tailwindcss ,即可与让你享受组件同时随意的设置样式。 支持 VUE 官方地址:shadcn/ui 项目地址:https:…

极客时间《Redis核心技术与实战》开篇词 知识点总结

Redis 主要的数据持久化方式 RDB(Redis Database Backup file) RDB 是 Redis 提供的一种数据快照持久化方式,它会在指定的时间间隔内生成数据集的时间点快照,并将这些快照保存到磁盘上的一个 RDB 文件中。RDB 文件是一个压缩的二…

Docker 配置 HTTP 和 HTTPS 网络代理

前言 在内网环境中,为了实现全局代理上网,Linux 系统通常通过修改 .bashrc 或 /etc/profile 等文件,设置 HTTP 和 HTTPS 代理。这种方式可以为大多数应用提供代理支持,但 Docker 并不会自动读取系统的环境变量,因此需…

谷粒商城-消息队列Rabbitmq

RabbitMq参考文档 在谷粒商城项目中使用消息队列主要有以下几个重要原因: 异步处理提高性能 场景示例:在订单系统中,当用户提交订单后,系统需要完成多个操作,如更新库存、生成订单记录、发送订单通知等。如果这些操作…

动态规划 详解

动态规划(Dynamic Programming, DP)详解 动态规划是一种通过分解问题为子问题并利用子问题的解来解决原问题的算法设计方法。它通常用于解决具有 重叠子问题 和 最优子结构 性质的问题。 1. 动态规划的核心思想 1.1 重叠子问题 问题可以分解为多个子问…

【C语言】野指针问题详解及防范方法

博客主页: [小ᶻ☡꙳ᵃⁱᵍᶜ꙳] 本文专栏: C语言 文章目录 💯前言💯什么是野指针?💯未初始化的指针代码示例问题分析解决方法 💯指针越界访问代码示例问题分析解决方法 💯指向已释放内存的…