目标检测的事儿,咱们不绕弯子了。这篇文章就是教你一件事:

3分钟教你用YOLOv11检测口罩、手套、防护服,工业、医疗、安防一网打尽。

这篇文章不会深究模型原理,而是手把手带你跑通全流程训练与推理,使用的数据是HuggingFace上经典的CPPE-5数据集,涵盖多种个人防护装备(PPE)目标检测任务。

✳️什么是 CPPE-5?

这里采用rishitdagli/cppe-5公开数据集,含以下5类标签:

类别英文 中文含义
Coverall 防护服 / 连体衣
Face_Shield 面罩 / 防护面屏
Gloves 手套
Goggles 护目镜
Mask 口罩

共1000张图片,真实复杂场景,适合目标检测实战测试。

🚀 快速开始(带你跑通)

1. 下载数据集

1
2
3
wget https://huggingface.co/datasets/rishitdagli/cppe-5/resolve/main/data/test-00000-of-00001.parquet

wget https://huggingface.co/datasets/rishitdagli/cppe-5/resolve/main/data/train-00000-of-00001.parquet

2. 转换成yolo格式

2.1 数据预处理脚本

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

import pathlib

import torch
from datasets import load_dataset
from PIL import Image, ImageDraw
from sympy import im

# 这里路径根据实际路径来
cppe5 = load_dataset(
'parquet',
data_files={
"train": pathlib.Path(__file__).parent.parent.joinpath("train-00000-of-00001.parquet").__str__(),
"test": pathlib.Path(__file__).parent.parent.joinpath("test-00000-of-00001.parquet").__str__()
}
)
categories = cppe5["train"].features["objects"].feature["category"].names
id2label = {index: x for index, x in enumerate(categories, start=0)}
label2id = {v: k for k, v in id2label.items()}
# ['Coverall', 'Face_Shield', 'Gloves', 'Goggles', 'Mask']
print(categories)


def cppe5_to_yolo_format(examples, data_type='train'):
# 注意这里,images/train放的是训练图片,images/test放的是测试图片
# 同理,labels/train放的是对应图片标签
pathlib.Path(__file__).parent.joinpath(
f"images/{data_type}").mkdir(exist_ok=True, parents=True)
pathlib.Path(__file__).parent.joinpath(
f"labels/{data_type}").mkdir(exist_ok=True, parents=True)
for index, sample in enumerate(examples):
image = sample["image"].convert("RGB")
yolo_labels = []

# draw = ImageDraw.ImageDraw(image)
# for bbox, category in zip(sample["objects"]["bbox"], sample["objects"]["category"]):
# x, y, w, h = tuple(bbox)
# if max(bbox) > 1.0:
# x1, y1 = int(x), int(y)
# x2, y2 = int(x + w), int(y + h)
# else:
# raise NotImplementedError
# draw.rectangle((x1, y1, x2, y2), outline="red")
# draw.text((x1, y1), str(id2label[category]), fill="red")
# image.show()

for bbox, category in zip(sample["objects"]["bbox"], sample["objects"]["category"]):

x, y, w, h = tuple(bbox)
if max(bbox) > 1.0:
x1, y1 = int(x), int(y)
x2, y2 = int(x + w), int(y + h)
else:
raise NotImplementedError
img_w, img_h = image.size
x_center = (x1 + x2) / 2.0 / img_w
y_center = (y1 + y2) / 2.0 / img_h
width = (x2 - x1) / img_w
height = (y2 - y1) / img_h
yolo_labels.append(
f'{category} {x_center} {y_center} {width} {height}'
)
image.save(pathlib.Path(__file__).parent.joinpath(
f"images/{data_type}/{index}.jpg"))
with open(pathlib.Path(__file__).parent.joinpath(f"labels/{data_type}/{index}.txt"), "w") as f:
f.write("\n".join(yolo_labels))


cppe5_train = cppe5_to_yolo_format(cppe5["train"], data_type='train')
cppe5_test = cppe5_to_yolo_format(cppe5["test"], data_type='test')

2.2 配置文件data.yaml

1
2
3
4
5
train: /object_detection/yolo_impl/images/train
val: /object_detection/yolo_impl/images/test
nc: 5
names: ["Coverall", "Face_Shield", "Gloves", "Goggles", "Mask"]

3. 启动YOLOv11模型训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

import os.path

from ultralytics import YOLO

os.environ['WANDB_DISABLED'] = 'true'
os.environ['YOLO_OFFLINE'] = 'true'


model_path = '/object_detection/yolo_impl/yolo11n.pt'
assert os.path.exists(model_path)

model = YOLO(model_path)

model.train(
data=os.path.join(os.path.dirname(__file__), 'data.yaml'),
epochs=100, # 训练轮数
imgsz=640, # 输入图像尺寸
batch=64, # 批量大小,根据显存可调
workers=4, # 数据加载线程数

# 优化器与学习率
lr0=0.001, # 初始学习率
lrf=0.01, # 最低学习率(cosine decay 最终值)
optimizer='SGD', # 优化器
weight_decay=0.0005, # 权重衰减
warmup_epochs=3, # 热身轮数
patience=20, # 提前停止的容忍度

# 增强与正则
augment=True, # 启用数据增强
mosaic=1.0,
mixup=0.1,
hsv_h=0.015,
hsv_s=0.7,
hsv_v=0.4,
scale=0.5,
translate=0.1,
degrees=0.0, # 不启用旋转

# 损失函数超参
box=7.5, # 边界框损失权重
cls=0.5, # 类别损失权重

iou=0.2, # IoU 阈值(用于损失计算)

device=0 # 使用GPU:0

)

📊 训练效果如下(示例):

4. 推理预测

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

import pathlib

from ultralytics import YOLO

model = YOLO(pathlib.Path(__file__).parent.joinpath(
"runs/detect/train/weights/best.pt"
)
)

results = model.predict(

source='/object_detection/yolo_impl/images/test',
save=True

)
# 处理结果
for result in results:
print(f"检测到的对象数量: {len(result.boxes)}")
for box, conf, cls in zip(result.boxes.xyxy, result.boxes.conf, result.boxes.cls):
print(f"类别: {result.names[int(cls)]}, 置信度: {conf:.2f}, 坐标: {box}")
print(f"类别: {result.names[int(cls)]}, 置信度: {conf:.2f}, 坐标: {box}")

🔍 示例结果:

预测结果1 预测结果2

✅ 总结

恭喜你,现在你已经掌握了如何用YOLOv11快速完成PPE检测模型训练与推理,下一步可以尝试:

  • 加入你的实际工业/医疗图像
  • 微调超参数
  • 部署模型到边缘设备

📬 有问题欢迎评论交流,觉得有用记得一键三连 🙌!