实战项目
准备训练数据
在 kaggle 平台搜索相关的数据,打开 kaggle 平台,搜索 Safety Helmet 或者 worker safety。
YOLOv8 - Safety Helmet Detection
编辑
模型训练时使用的yaml配置文件(safehat.yaml)
# 模型训练时使用的yaml配置文件。该文件说明了数据的地址和待训练的类别 # 配置三个路径,分别对应训练数据、验证数据、测试数据的地址 # 训练数据用于模型的训练 train: /root/css-data/train/images/ # 验证数据用于模型训练过程中的评估和参数调试 val: /root/css-data/valid/images/ # 测试数据用于模型完成训练后的测试 test: /root/css-data/test/images/ # number of classes nc: 10 # class names #names: ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'] #Hardhat - 安全帽 #Mask - 口罩 #NO-Hardhat - 无安全帽 #NO-Mask - 无口罩 #NO-Safety Vest - 无安全背心 #Person - 人 #Safety Cone - 安全锥 #Safety Vest - 安全背心 #machinery - 机械(挖掘机) #vehicle - 车辆 names: [Hardhat, Mask, NO-Hardhat, NO-Mask, NO-Safety Vest, Person, Safety Cone, Safety Vest, machinery, vehicle]
定义训练文件(yolo_train.py)
from ultralytics import YOLO # 导入Yolo模块 # 加载yolov8的预训练模型,这个模型是yolov8使用了coco数据集训练的通用目标检测模型, # 我们将它作为基础模型,在该模型的基础上,训练安全帽模型 model = YOLO('yolov8n.pt') # 用于加载模型 # 训练用户自定义的数据集,数据的配置保存在safehat.yaml中,epochs等于100表示100轮迭代 model.train(data='safehat.yaml', epochs=100) # 使用验证集验证效果 model.val()
模型训练结束后,会在这个runs > detect > train 路径下保存训练结果
测试yolo_test.py
from ultralytics import YOLO # 在代码中导入yolo模块 # 导入训练好的模型best.pt model = YOLO('best.pt') # 自己构造一些数据 # 在识别自己构造的数据时,传入了classes = [0, 2], # 代表只输出0和2,也就是安全帽是否佩戴这两个类别 # line_width = 30表示指定识别框的字体大小为30 model.predict('/root/css-data/valid/images/YouTube_FreeStockFootage_People-wearing-face-mask_Empty-Street_Covid19-D_DbgrvhlGs-720p_mp4-66_jpg.rf.72af13fb1f20520e0d41547037ced88b.jpg', save=True, classes = [0, 2], line_width = 30)
识别结果会保存在当前目录下的runs/detect/predict中
模型部署
部署前的准备:模型格式转换
转换方式非常简单,Ultralytics 一行代码搞定:
from ultralytics import YOLO model = YOLO("best.pt") model.export(format="onnx") # 导出 ONNX model.export(format="engine") # 导出 TensorRT(需要 CUDA 环境)
场景一:服务器 / 云端部署
适合什么情况?
- 需要多人、多系统同时调用
- 算力集中在服务器侧
- 用户通过网页或 App 上传图片进行检测
核心思路
用户上传图片 → 服务器推理 → 返回检测结果
模型始终运行在服务器上,不暴露给用户。
第一步:用 FastAPI 封装推理接口
# app.py from fastapi import FastAPI, UploadFile, File from fastapi.responses import JSONResponse from ultralytics import YOLO import numpy as np import cv2 app = FastAPI() model = YOLO("best.onnx") # 服务启动时加载一次模型 @app.post("/detect") async def detect(file: UploadFile = File(...)): # 读取图片 img_bytes = await file.read() nparr = np.frombuffer(img_bytes, np.uint8) img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # 推理 results = model.predict(img, conf=0.25) # 整理结果 detections = [] for r in results: for box in r.boxes: detections.append({ "class_name": model.names[int(box.cls)], "confidence": round(float(box.conf), 3), "bbox": { "x1": round(float(box.xyxy[0][0])), "y1": round(float(box.xyxy[0][1])), "x2": round(float(box.xyxy[0][2])), "y2": round(float(box.xyxy[0][3])), } }) return JSONResponse({"count": len(detections), "detections": detections})
第二步:用 Docker 打包
# Dockerfile FROM python:3.10-slim WORKDIR /app COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txt COPY app.py . COPY best.onnx . EXPOSE 8000 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
# 构建并运行 docker build -t yolo-api . docker run -d -p 8000:8000 --gpus all --name yolo-service yolo-api
第三步:外部调用测试
import requests with open("test.jpg", "rb") as f: response = requests.post( "http://your-server-ip:8000/detect", files={"file": f} ) print(response.json())
高并发场景:Triton 推理服务器
如果有大量用户同时调用(例如同时供几十个工厂使用),普通 FastAPI 扛不住,推荐使用 NVIDIA Triton 推理服务器:
工厂 A ──┐ 工厂 B ──┤ 工厂 C ──┼──→ Nginx 负载均衡 → Triton 推理服务器 → TensorRT 模型 工厂 D ──┤ 工厂 E ──┘
场景二:Jetson 边缘设备部署
适合什么情况?
- 需要在设备本地实时推理,不依赖网络
- 场景:工厂产线检测、智能摄像头、机器人视觉
- 设备:Jetson Nano / Xavier / Orin 系列
核心:转换为 TensorRT 格式
Jetson 自带 NVIDIA GPU,使用 TensorRT 可以充分利用硬件加速。
from ultralytics import YOLO model = YOLO("best.pt") model.export( format="engine", # TensorRT 格式 device=0, half=True, # FP16 加速,速度翻倍,显存减半 imgsz=640, ) # 生成 best.engine
在 Jetson 上推理
from ultralytics import YOLO # 加载 TensorRT 模型,调用方式和 .pt 完全一样 model = YOLO("best.engine") results = model.predict("test.jpg", conf=0.25) for r in results: print(r.boxes)
场景三:RK3588 嵌入式部署
适合什么情况?
- 低成本嵌入式设备,无 NVIDIA GPU
- 场景:工业检测设备、国产嵌入式平台
- 设备:RK3588 / RK3566 / RK1808 等瑞芯微芯片
核心:转换为 RKNN 格式
RK3588 内置 NPU(算力 6 TOPS),必须转换为 格式才能用 NPU 加速。.rknn
转换链路如下:
best.pt → best.onnx → best.rknn
(第一步) (第二步)
pt → onnx(PC 上)
from ultralytics import YOLO model = YOLO("best.pt") model.export( format="onnx", opset=12, # rknn-toolkit2 推荐 opset 12 simplify=True, dynamic=False, imgsz=640, )
onnx → rknn(PC 上)
# convert.py from rknn.api import RKNN rknn = RKNN() rknn.config( mean_values=[[0, 0, 0]], std_values=[[255, 255, 255]], target_platform="rk3588", quantized_dtype="asymmetric_quantized-8", # INT8 量化 ) rknn.load_onnx(model="best.onnx") # dataset.txt:每行一张图片路径,用于 INT8 量化校准(100~200 张) rknn.build(do_quantization=True, dataset="dataset.txt") rknn.export_rknn("best.rknn") print("✅ 转换完成") rknn.release()
第三步:将模型传到 RK3588
scp best.rknn user@rk3588-ip:/home/user/
第四步:在 RK3588 上推理
# 板子上只需安装轻量级运行时 pip install rknnlite2
# inference.py(在 RK3588 上运行) from rknnlite.api import RKNNLite import cv2 import numpy as np rknn = RKNNLite() rknn.load_rknn("best.rknn") rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_ALL) # 三核 NPU 全开 # 注意:RK3588 输入格式是 NHWC,和 PyTorch 的 NCHW 不同! img = cv2.imread("test.jpg") img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (640, 640)) img = np.expand_dims(img, 0) # [1, 640, 640, 3] outputs = rknn.inference(inputs=[img]) print("输出 shape:", outputs[0].shape) rknn.release()
需要多人/多系统远程调用? └── 是 → 服务器部署(ONNX 或 TensorRT + FastAPI) 并发量大?→ 加上 Triton 推理服务器 需要设备本地实时推理? ├── 有 NVIDIA GPU(Jetson)→ TensorRT 格式 └── 无 NVIDIA GPU ├── 瑞芯微芯片(RK3588)→ RKNN 格式 └── 树莓派 / ARM 通用 → ONNX + onnxruntime
参考内容
【1】