基于Pytorch和yolov8n手搓安全帽目标检测的全过程
一.背景
还是之前的主题，使用开源软件为公司搭建安全管理平台，从视觉模型识别安全帽开始。主要参考学习了开源项目 https://github.com/jomarkow/Safety-Helmet-Detection，我是从运行、训练、标注倒过来学习的。由于工作原因，抽空学习了vscode的使用、python语法，了解了pytorch、yolo、ultralytics、cuda、cuda toolkit、cudnn、AI的分类、AI的相关专业名词等等。到这里，基本可以利用工程化的方式解决目标检测环境搭建、AI标注、训练、运行全过程了。
二.捋一捋思路
1.人类目前的AI本质是一堆数据喂出个没有智慧的算命师
大量的数量，AI最大程度（如95%）的找到了满足了数据输入与结果的因果关系。你说不智能吧？你人不一定能有它预测的准。你说智能吧？其实受限于有限的输入数据，毕竟不能穷举所有输入数据，人类掌握的数据既有限，也不一定就是客观的。个人愚见，不喜勿喷。
2.图形目标检测就是划圈圈、多看看、试一试的过程
划圈圈就是数据标注，在图片上面框一下。多看看就是让算法自己看图片，好比我们教小孩子指着自己反复说“爸爸”，而后小孩慢慢的学会了叫“爸爸”的过程。试一试，就是把训练的结果（AI模型或者算法模型）拿来运行，就像让小孩对着另外一个男人让他去称呼，他可能也叫“爸爸”。那么，就需要纠正告诉他只有自己才是叫“爸爸”，其他的男人应该叫“叔叔”。图形目标检测就是这么个过程，没有什么神秘的。当然，我们是站在前辈肩膀上的，那些算法在那个年代写出来，确实是值得敬佩的。
3.环境与工具说明
windows11家庭版、 显卡NVIDIA GeForce GT 730(我没有用GPU训练，主要还是显卡太老了，版本兼容性问题把我弄哭了，CPU慢点~~就慢点吧)、vscode1.83.0 、conda 24.9.2、Python 3.12.7、pytorch2.5.0、yolo8（后面换最新的yolo11试试）
三.开始手搓
1.创建空的工程结构
1）在非中文、空格目录中创建object-detection-hello文件夹
2）vscode打开文件夹
3）在vscode中创建目录及文件等
下图中的文件夹，并不是必须的，但是推荐这样
4）编写训练需要的参数文件train_config.yaml
train: ../../datas/images
val: ../../datas/images
#class count
nc: 1
# names: ['helmet']
names: ['helmet']
labels: ../../datas/labels5)下载yolov8n.pt到models目录中
在我之前上传的也有，详见本文章关联的资源。也可以去安全帽开源项目GitHub - jomarkow/Safety-Helmet-Detection: YoloV8 model, trained for recognizing if construction workers are wearing their protection helmets in mandatory areas中去下载，在根目录就有。
2.图片中安全帽标注
1）图片准备
去把安全帽的开源下载下来，里面有图片。我只选择了0-999，共1千张图片，毕竟我是cpu，训练慢，1千张估计也能有个效果了。
2）图片标注
参考之前的文章在windows系统中使用labelimg对图片进行标注之工具安装及简单使用-CSDN博客
3）标注数据处理
我标注后的文件是xml，需要转为txt文件。内容分别是
<annotation>
<folder>images</folder>
<filename>hard_hat_workers2.png</filename>
<path>D:\zsp\works\temp\20241119-zsp-helmet\Safety-Helmet-Detection-main\data\images\hard_hat_workers2.png</path>
<source>
<database>Unknown</database>
</source>
<size>
<width>416</width>
<height>415</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>helmet</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>295</xmin>
<ymin>219</ymin>
<xmax>326</xmax>
<ymax>249</ymax>
</bndbox>
</object>
<object>
<name>helmet</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>321</xmin>
<ymin>212</ymin>
<xmax>365</xmax>
<ymax>244</ymax>
</bndbox>
</object>
</annotation>0 0.745192 0.565060 0.072115 0.060241
0 0.826923 0.549398 0.081731 0.072289有时候，默认就是文本文件的格式了。如果不是，创建converter.py直接转换：
from xml.dom import minidom
import os
classes={"helmet":0}
def convert_coordinates(size, box):
dw = 1.0/size[0]
dh = 1.0/size[1]
x = (box[0]+box[1])/2.0
y = (box[2]+box[3])/2.0
w = box[1]-box[0]
h = box[3]-box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def converter(classes):
old_labels_path = "datas/images/raw_data/"
new_labels_path = "datas/images/raw_data/"
current_path = os.getcwd()
# 打印当前工作目录
print("当前路径是:", current_path)
for file_name in os.listdir(old_labels_path):
if ".xml" in file_name:
old_file = minidom.parse(f"{old_labels_path}/{file_name}")
name_out = (file_name[:-4]+'.txt')
with open(f"{new_labels_path}/{name_out}", "w") as new_file:
itemlist = old_file.getElementsByTagName('object')
size = old_file.getElementsByTagName('size')[0]
width = int((size.getElementsByTagName('width')[0]).firstChild.data)
height = int((size.getElementsByTagName('height')[0]).firstChild.data)
for item in itemlist:
# get class label
class_name = (item.getElementsByTagName('name')[0]).firstChild.data
if class_name in classes:
label_str = str(classes[class_name])
else:
label_str = "-1"
print (f"{class_name} not in function classes")
# get bbox coordinates
xmin = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('xmin')[0]).firstChild.data
ymin = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('ymin')[0]).firstChild.data
xmax = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('xmax')[0]).firstChild.data
ymax = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('ymax')[0]).firstChild.data
b = (float(xmin), float(xmax), float(ymin), float(ymax))
bb = convert_coordinates((width,height), b)
#print(bb)
#new_file.write(f"{label_str} {' '.join([(f'{a}.6f') for a in bb])}\n")
new_file.write(f"{label_str} {' '.join([(f'{a:.6f}') for a in bb])}\n")
print (f"wrote {name_out}")
def main():
converter(classes)
if __name__ == '__main__':
main()4）偷懒直接用开源项目标注的labels
当然，也可以偷懒，复制开源项目的labels中0-999的txt文件到我们的labels目录。但是它文件是多目标检测，我们只保留下我们的安全帽标注，也就是txt文件中0开始的行。所以，我在AI的帮助下写了这个utils/deleteOtherclass.py程序，文件内容如下：
import os
PROY_FOLDER = os.getcwd().replace("\\","/")
INPUT_FOLDER = f"{PROY_FOLDER}/datas/labels/"
files = os.listdir(INPUT_FOLDER)
def process_file(file_path):
# 存储处理后的行
processed_lines = []
try:
with open(file_path, 'r') as file:
for line in file:
if len(line) > 0 and line[0] == '0': # 检查行的第一个字符是否为 0
processed_lines.append(line)
except FileNotFoundError:
print(f"文件 {file_path} 未找到")
return
try:
with open(file_path, 'w') as file: # 以写入模式打开文件，会清空原文件
file.writelines(processed_lines)
except Exception as e:
print(f"写入文件时出现错误: {e}")
for file_name in files:
file_path = INPUT_FOLDER + file_name
print(file_path)
process_file(file_path)
执行这个程序后，txt文件中就只剩下0开始的了，也就是我们安全帽的标注了。
比如，hard_hat_workers0.txt文件前后内容分别如下：
0 0.914663 0.349760 0.112981 0.141827
0 0.051683 0.396635 0.084135 0.091346
0 0.634615 0.379808 0.052885 0.091346
0 0.748798 0.391827 0.055288 0.086538
0 0.305288 0.397837 0.052885 0.069712
0 0.216346 0.397837 0.048077 0.069712
1 0.174279 0.379808 0.050481 0.067308
1 0.801683 0.383413 0.055288 0.088942
1 0.443510 0.411058 0.045673 0.072115
1 0.555288 0.400240 0.043269 0.074519
1 0.500000 0.383413 0.038462 0.064904
0 0.252404 0.360577 0.033654 0.048077
1 0.399038 0.393029 0.043269 0.064904
0 0.914663 0.349760 0.112981 0.141827
0 0.051683 0.396635 0.084135 0.091346
0 0.634615 0.379808 0.052885 0.091346
0 0.748798 0.391827 0.055288 0.086538
0 0.305288 0.397837 0.052885 0.069712
0 0.216346 0.397837 0.048077 0.069712
0 0.252404 0.360577 0.033654 0.0480773.编写训练的程序
Ultralytics的配置文件，一般存放在C:\Users\Dell\AppData\Roaming\Ultralytics\settings.json中，路径中的Dell你要换成你的用户名哦。当然，这里只是知道就好了。看看它内部的内容如下：
{
"settings_version": "0.0.6",
"datasets_dir": "D:\\zsp\\works\\temp\\20241218-zsp-pinwei\\object-detection-hello",
"weights_dir": "weights",
"runs_dir": "runs",
"uuid": "09253350c3bd45fd265c2e8346acaaa599711c1c3ef91e7e78ceff31d4132a83",
"sync": true,
"api_key": "",
"openai_api_key": "",
"clearml": true,
"comet": true,
"dvc": true,
"hub": true,
"mlflow": true,
"neptune": true,
"raytune": true,
"tensorboard": true,
"wandb": false,
"vscode_msg": true
}从配置的内容看，我们可能需要修改的是datasets_dir，为了更优雅，我写了代码来修改。
1）用程序去修改配置文件的代码scripts/ultralytics_init.py
from ultralytics import settings
import os
def update_ultralytics_settings(key, value):
try:
#settings.update(key, value) # 假设存在 update 方法
settings[key]=value
print(f"Updated {key} to {value} in ultralytics settings.")
except AttributeError:
print(f"Failed to update {key}, the update method may not exist in the settings module.")
def init():
current_path = os.getcwd()
print(current_path)
# 调用函数，使用形参，参数值用引号括起来
update_ultralytics_settings("datasets_dir",current_path)
print(settings)
2）建立训练的主程序scripts/train.py
import ultralytics_init as uinit
uinit.init()
from ultralytics import YOLO
import os
# Return a specific setting
# value = settings["runs_dir"]
model = YOLO("models/yolov8n.pt")
model.train(data="config/train_config.yaml", epochs=10)
result = model.val()
path = model.export(format="onnx")
代码我我觉得不解释了，一看就明白。
3）配置文件config/train_config.yaml的设置
#训练的图片集合
train: ../../datas/images
#过程验证的图片集合
val: ../../datas/images
#目标类型的数量
nc: 1
#label的英文名称
names: ['helmet']4.执行训练
右上角，点三角形运行。
训练了10代，训练过程约1小时。日志如下：
PS D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello> & C:/Users/Dell/.conda/envs/myenv/python.exe d:/zsp/works/temp/20241218-zsp-pinwei/object-detection-hello/scripts/train.py
D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello
Updated datasets_dir to D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello in ultralytics settings.
JSONDict("C:\Users\Dell\AppData\Roaming\Ultralytics\settings.json"):
{
"settings_version": "0.0.6",
"datasets_dir": "D:\\zsp\\works\\temp\\20241218-zsp-pinwei\\object-detection-hello",
"weights_dir": "weights",
"runs_dir": "runs",
"uuid": "09253350c3bd45fd265c2e8346acaaa599711c1c3ef91e7e78ceff31d4132a83",
"sync": true,
"api_key": "",
"openai_api_key": "",
"clearml": true,
"comet": true,
"dvc": true,
"hub": true,
"mlflow": true,
"neptune": true,
"raytune": true,
"tensorboard": true,
"wandb": false,
"vscode_msg": true
}
New https://pypi.org/project/ultralytics/8.3.55 available 😃 Update with 'pip install -U ultralytics'
Ultralytics 8.3.49 🚀 Python-3.12.7 torch-2.5.0 CPU (12th Gen Intel Core(TM) i7-12700)
engine\trainer: task=detect, mode=train, model=models/yolov8n.pt, data=config/train_config.yaml, epochs=10, time=None, patience=100, batch=16, imgsz=640, save=True, save_period=-1, cache=False, device=None, workers=8, project=None, name=train, exist_ok=False, pretrained=True, optimizer=auto, verbose=True, seed=0, deterministic=True, single_cls=False, rect=False, cos_lr=False, close_mosaic=10, resume=False, amp=True, fraction=1.0, profile=False, freeze=None, multi_scale=False, overlap_mask=True, mask_ratio=4, dropout=0.0, val=True, split=val, save_json=False, save_hybrid=False, conf=None, iou=0.7, max_det=300, half=False, dnn=False, plots=True, source=None, vid_stride=1, stream_buffer=False, visualize=False, augment=False, agnostic_nms=False, classes=None, retina_masks=False, embed=None, show=False, save_frames=False, save_txt=False, save_conf=False, save_crop=False, show_labels=True, show_conf=True, show_boxes=True, line_width=None, format=torchscript, keras=False, optimize=False, int8=False, dynamic=False, simplify=True, opset=None, workspace=None, nms=False, lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=7.5, cls=0.5, dfl=1.5, pose=12.0, kobj=1.0, nbs=64, hsv_h=0.015, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, bgr=0.0, mosaic=1.0, mixup=0.0, copy_paste=0.0, copy_paste_mode=flip, auto_augment=randaugment, erasing=0.4, crop_fraction=1.0, cfg=None, tracker=botsort.yaml, save_dir=runs\detect\train
Overriding model.yaml nc=80 with nc=1
from n params module arguments
0 -1 1 464 ultralytics.nn.modules.conv.Conv [3, 16, 3, 2]
1 -1 1 4672 ultralytics.nn.modules.conv.Conv [16, 32, 3, 2]
2 -1 1 7360 ultralytics.nn.modules.block.C2f [32, 32, 1, True]
3 -1 1 18560 ultralytics.nn.modules.conv.Conv [32, 64, 3, 2]
4 -1 2 49664 ultralytics.nn.modules.block.C2f [64, 64, 2, True]
5 -1 1 73984 ultralytics.nn.modules.conv.Conv [64, 128, 3, 2]
6 -1 2 197632 ultralytics.nn.modules.block.C2f [128, 128, 2, True]
7 -1 1 295424 ultralytics.nn.modules.conv.Conv [128, 256, 3, 2]
8 -1 1 460288 ultralytics.nn.modules.block.C2f [256, 256, 1, True]
9 -1 1 164608 ultralytics.nn.modules.block.SPPF [256, 256, 5]
10 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
11 [-1, 6] 1 0 ultralytics.nn.modules.conv.Concat [1]
12 -1 1 148224 ultralytics.nn.modules.block.C2f [384, 128, 1]
13 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
14 [-1, 4] 1 0 ultralytics.nn.modules.conv.Concat [1]
15 -1 1 37248 ultralytics.nn.modules.block.C2f [192, 64, 1]
16 -1 1 36992 ultralytics.nn.modules.conv.Conv [64, 64, 3, 2]
17 [-1, 12] 1 0 ultralytics.nn.modules.conv.Concat [1]
18 -1 1 123648 ultralytics.nn.modules.block.C2f [192, 128, 1]
19 -1 1 147712 ultralytics.nn.modules.conv.Conv [128, 128, 3, 2]
20 [-1, 9] 1 0 ultralytics.nn.modules.conv.Concat [1]
21 -1 1 493056 ultralytics.nn.modules.block.C2f [384, 256, 1]
22 [15, 18, 21] 1 751507 ultralytics.nn.modules.head.Detect [1, [64, 128, 256]]
Model summary: 225 layers, 3,011,043 parameters, 3,011,027 gradients, 8.2 GFLOPs
train: Scanning D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello\datas\labels.cache.
val: Scanning D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello\datas\labels.cache...
Plotting labels to runs\detect\train\labels.jpg...
optimizer: 'optimizer=auto' found, ignoring 'lr0=0.01' and 'momentum=0.937' and determining best 'optimizer', 'lr0' and 'momentum' automatically...
optimizer: AdamW(lr=0.002, momentum=0.9) with parameter groups 57 weight(decay=0.0), 64 weight(decay=0.0005), 63 bias(decay=0.0)
Image sizes 640 train, 640 val
Using 0 dataloader workers
Logging results to runs\detect\train
Starting training for 10 epochs...
Closing dataloader mosaic
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
1/10 0G 1.592 2.148 1.278 40 640: 100%|██████████|
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [02:58<00:00, 5.57s/it]
all 1000 3792 0.977 0.033 0.423 0.229
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
2/10 0G 1.483 1.464 1.215 23 640: 100%|██████████| 63/63 [07:53<00:00, 7.51s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:20<00:00, 2.51s/it]
all 1000 3792 0.697 0.647 0.687 0.398
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
3/10 0G 1.489 1.318 1.242 15 640: 100%|██████████| 63/63 [02:52<00:00, 2.74s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:13<00:00, 2.30s/it]
all 1000 3792 0.783 0.662 0.744 0.401
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
4/10 0G 1.47 1.183 1.22 19 640: 100%|██████████| 63/63 [02:50<00:00, 2.71s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:12<00:00, 2.25s/it]
all 1000 3792 0.837 0.749 0.832 0.496
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
5/10 0G 1.42 1.041 1.196 31 640: 100%|██████████| 63/63 [02:51<00:00, 2.72s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:11<00:00, 2.22s/it]
all 1000 3792 0.867 0.776 0.87 0.537
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
6/10 0G 1.4 0.9758 1.196 30 640: 100%|██████████| 63/63 [02:51<00:00, 2.72s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:11<00:00, 2.24s/it]
all 1000 3792 0.898 0.818 0.902 0.565
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
7/10 0G 1.352 0.8787 1.156 37 640: 100%|██████████| 63/63 [02:52<00:00, 2.74s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:20<00:00, 2.51s/it]
all 1000 3792 0.921 0.843 0.922 0.576
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
8/10 0G 1.307 0.825 1.13 17 640: 100%|██████████| 63/63 [06:18<00:00, 6.01s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:11<00:00, 2.22s/it]
all 1000 3792 0.906 0.845 0.924 0.58
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
9/10 0G 1.294 0.7867 1.133 29 640: 100%|██████████| 63/63 [02:51<00:00, 2.72s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:10<00:00, 2.21s/it]
all 1000 3792 0.922 0.87 0.938 0.611
Epoch GPU_mem box_loss cls_loss dfl_loss Instances Size
10/10 0G 1.257 0.7387 1.119 57 640: 100%|██████████| 63/63 [02:51<00:00, 2.72s/it]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [01:18<00:00, 2.47s/it]
all 1000 3792 0.933 0.884 0.95 0.62
10 epochs completed in 0.934 hours.
Optimizer stripped from runs\detect\train\weights\last.pt, 6.2MB
Optimizer stripped from runs\detect\train\weights\best.pt, 6.2MB
Validating runs\detect\train\weights\best.pt...
Ultralytics 8.3.49 🚀 Python-3.12.7 torch-2.5.0 CPU (12th Gen Intel Core(TM) i7-12700)
Model summary (fused): 168 layers, 3,005,843 parameters, 0 gradients, 8.1 GFLOPs
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 32/32 [00:59<00:00, 1.86s/it]
all 1000 3792 0.933 0.884 0.95 0.62
Speed: 1.4ms preprocess, 51.6ms inference, 0.0ms loss, 0.4ms postprocess per image
Results saved to runs\detect\train
Ultralytics 8.3.49 🚀 Python-3.12.7 torch-2.5.0 CPU (12th Gen Intel Core(TM) i7-12700)
Model summary (fused): 168 layers, 3,005,843 parameters, 0 gradients, 8.1 GFLOPs
val: Scanning D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello\datas\labels.cache... 1000 images, 76 backgrounds, 0 corrupt: 100%|██████████| 1000/1000 [00:00<?, ?it/s]
Class Images Instances Box(P R mAP50 mAP50-95): 100%|██████████| 63/63 [00:54<00:00, 1.17it/s]
all 1000 3792 0.933 0.884 0.95 0.62
Speed: 1.1ms preprocess, 46.3ms inference, 0.0ms loss, 0.4ms postprocess per image
Results saved to runs\detect\train2
Ultralytics 8.3.49 🚀 Python-3.12.7 torch-2.5.0 CPU (12th Gen Intel Core(TM) i7-12700)
PyTorch: starting from 'runs\detect\train\weights\best.pt' with input shape (1, 3, 640, 640) BCHW and output shape(s) (1, 5, 8400) (6.0 MB)
ONNX: starting export with onnx 1.17.0 opset 19...
ONNX: slimming with onnxslim 0.1.43...
ONNX: export success ✅ 1.0s, saved as 'runs\detect\train\weights\best.onnx' (11.7 MB)
Export complete (1.1s)
Results saved to D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello\runs\detect\train\weights
Predict: yolo predict task=detect model=runs\detect\train\weights\best.onnx imgsz=640
Validate: yolo val task=detect model=runs\detect\train\weights\best.onnx imgsz=640 data=config/train_config.yaml
Visualize: https://netron.app5.运行训练后的模型看看效果
1）把训练后的模型best.pt准备好
训练结果模型在哪里？看看日志啊Results saved to D:\zsp\works\temp\20241218-zsp-pinwei\object-detection-hello\runs\detect\train\weights。我去把它复制到了test文件夹中。
2）把测试图片1.jpg复制到test目录下
注意：图片中的马赛克是为了保护同事隐私添加的，并非程序效果。
3）编写验证代码script/test.py
代码也不解释了，一看就明白的
import os
from ultralytics import YOLO
import cv2
PROY_FOLDER = os.getcwd().replace("\\","/")
INPUT_FOLDER = f"{PROY_FOLDER}/test/"
OUTPUT_FOLDER = f"{PROY_FOLDER}/test_out/"
MODEL_PATH = f"{PROY_FOLDER}/test/best.pt"
if not os.path.exists(OUTPUT_FOLDER):
os.mkdir(OUTPUT_FOLDER)
model = YOLO(MODEL_PATH)
files = os.listdir(INPUT_FOLDER)
def draw_box(params, frame, threshold = 0.2):
x1, y1, x2, y2, score, class_id = params
if score > threshold:
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
textValue=results.names[int(class_id)].upper()
if "HELMET" in textValue :
textValue="yes"
cv2.putText(frame, textValue, (int(x1), int(y1 - 10)),
cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
elif "HEAD" in textValue :
textValue="no!!!"
cv2.putText(frame, textValue, (int(x1), int(y1 - 10)),
cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
else:
cv2.putText(frame, textValue, (int(x1), int(y1 - 10)),
cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
return frame
for file_name in files:
file_path = INPUT_FOLDER + file_name
if ".jpg" in file_name:
image_path_out = OUTPUT_FOLDER + file_name[:-4] + "_out.jpg"
image = cv2.imread(file_path,cv2.IMREAD_COLOR)
results = model(image)[0]
for result in results.boxes.data.tolist():
image = draw_box(result, image)
cv2.imwrite(image_path_out, image)
cv2.destroyAllWindows()4）运行结果查看
注意：图片中的马赛克是为了保护同事隐私添加的，并非程序效果。
四.总结
到这里，我们就完成了从标注、写代码训练、验证训练结果的全过程。为我们后面搭建一个安全帽检测的服务奠定了基础，当然对于训练结果的调优干预还是我们的短板，毕竟是初学，我想未来都不是问题。我有编程基础，过程中还是出现了不少的问题，但只要努力尝试去看输出日志，都能解决。不行的话，把输出日志拿去问AI都能找到解决问题的思路。