1. 边缘设备上的YOLOv5模型转换与部署

任务目标

将PyTorch训练的YOLOv5s模型转换为ONNX格式，并在树莓派4B（4GB RAM，ARM64架构）上使用ONNX Runtime完成端到端目标检测推理。最终通过摄像头或静态图片输入，输出带边界框的检测结果，单帧推理时间需满足基本实时性要求（<1s）。

环境与依赖

PC端（模型转换）：Python 3.8+，PyTorch ≥1.8，YOLOv5官方仓库（v6.0+），onnx ≥1.12

树莓派端（推理）：Python 3.9（系统默认），ONNX Runtime（ARM64版本），OpenCV，NumPy

树莓派安装命令：

pip install onnxruntime opencv-python-headless numpy
# 若需摄像头支持：sudo apt install libcblas-dev libatlas-base-dev

模型转换（PC端执行）

使用YOLOv5官方仓库的export.py脚本，将预训练权重转为ONNX。关键参数：--include onnx指定导出格式，--dynamic保持动态输入尺寸以适配不同分辨率输入。

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
python export.py --weights yolov5s.pt --include onnx --dynamic --opset 12

转换完成后生成yolov5s.onnx（约29MB）。将文件拷贝至树莓派。

树莓派推理实现

以下代码为完整可运行的检测脚本，包含预处理、ONNX Runtime推理、后处理（NMS）与结果绘制。模型输入为640×640 RGB图像，输出为1×25200×85张量（80类COCO + 4个框坐标 + 1个置信度）。

import cv2
import numpy as np
import onnxruntime as ort

class YOLOv5ONNX:
    def __init__(self, onnx_path, conf_thres=0.4, iou_thres=0.5):
        self.session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.input_shape = (640, 640)
        self.classes = self._load_coco_names()
        
    def _load_coco_names(self):
        # COCO 80类名称（此处仅列出前3个示例，完整列表需自行补全）
        return ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
                'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
                'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
                'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
                'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
                'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
                'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
                'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
                'hair drier', 'toothbrush']
    
    def preprocess(self, img):
        # 保持宽高比缩放并填充至640×640
        h, w = img.shape[:2]
        scale = min(self.input_shape[0]/h, self.input_shape[1]/w)
        new_w, new_h = int(w*scale), int(h*scale)
        resized = cv2.resize(img, (new_w, new_h))
        canvas = np.full((self.input_shape[0], self.input_shape[1], 3), 114, dtype=np.uint8)
        x_offset = (self.input_shape[1]-new_w)//2
        y_offset = (self.input_shape[0]-new_h)//2
        canvas[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = resized
        # BGR→RGB，HWC→CHW，归一化
        blob = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        blob = np.transpose(blob, (2, 0, 1))[None]
        return blob, scale, x_offset, y_offset
    
    def postprocess(self, outputs, scale, x_offset, y_offset, orig_shape):
        # outputs shape: (1, 25200, 85)
        preds = outputs[0][0]  # (25200, 85)
        # 过滤低置信度
        conf = preds[:, 4]
        mask = conf > self.conf_thres
        preds = preds[mask]
        if len(preds) == 0:
            return []
        # 获取类别与类别置信度
        class_conf = preds[:, 5:] * preds[:, 4:5]
        class_id = np.argmax(class_conf, axis=1)
        final_conf = np.max(class_conf, axis=1)
        # 再次过滤
        mask = final_conf > self.conf_thres
        preds = preds[mask]
        class_id = class_id[mask]
        final_conf = final_conf[mask]
        if len(preds) == 0:
            return []
        # 解码边界框（xywh → xyxy）
        boxes = preds[:, :4].copy()
        boxes[:, 0] = (boxes[:, 0] - boxes[:, 2]/2 - x_offset) / scale
        boxes[:, 1] = (boxes[:, 1] - boxes[:, 3]/2 - y_offset) / scale
        boxes[:, 2] = (boxes[:, 0] + boxes[:, 2]/scale)
        boxes[:, 3] = (boxes[:, 1] + boxes[:, 3]/scale)
        # 裁剪到图像范围
        boxes[:, 0] = np.clip(boxes[:, 0], 0, orig_shape[1])
        boxes[:, 1] = np.clip(boxes[:, 1], 0, orig_shape[0])
        boxes[:, 2] = np.clip(boxes[:, 2], 0, orig_shape[1])
        boxes[:, 3] = np.clip(boxes[:, 3], 0, orig_shape[0])
        # NMS（按类别分别执行）
        indices = cv2.dnn.NMSBoxes(boxes.tolist(), final_conf.tolist(), self.conf_thres, self.iou_thres)
        results = []
        if len(indices) > 0:
            for i in indices.flatten():
                results.append({
                    'box': boxes[i].astype(int),
                    'conf': float(final_conf[i]),
                    'class_id': int(class_id[i]),
                    'class_name': self.classes[int(class_id[i])]
                })
        return results
    
    def detect(self, img):
        blob, scale, x_off, y_off = self.preprocess(img)
        outputs = self.session.run(None, {self.session.get_inputs()[0].name: blob})
        return self.postprocess(outputs, scale, x_off, y_off, img.shape[:2])
    
    def draw(self, img, detections):
        for det in detections:
            x1, y1, x2, y2 = det['box']
            label = f"{det['class_name']} {det['conf']:.2f}"
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img, label, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
        return img

# 使用示例
if __name__ == '__main__':
    model = YOLOv5ONNX('yolov5s.onnx')
    # 静态图片检测
    img = cv2.imread('test.jpg')
    dets = model.detect(img)
    result = model.draw(img, dets)
    cv2.imwrite('result.jpg', result)
    print(f"检测到 {len(dets)} 个目标")
    for d in dets:
        print(f"{d['class_name']}: {d['conf']:.2f}, box={d['box']}")

运行与验证

运行方式：将yolov5s.onnx与测试图片置于同一目录，执行python detect.py。若需摄像头实时检测，将主循环替换为cv2.VideoCapture(0)逐帧调用model.detect()即可。

验证方法：

确认输出图片result.jpg中目标框位置与类别正确（与PC端PyTorch推理结果对比，允许坐标微小偏差）；
在代码中记录单帧推理耗时（time.time()包围session.run），树莓派4B上640×640输入应小于1秒；
连续运行10分钟无内存溢出或崩溃。

预期输出示例：对于包含行人的测试图片，输出应包含person: 0.87, box=[x1 y1 x2 y2]格式的检测结果，且result.jpg中对应位置绘制绿色边界框。

常见问题

ONNX Runtime导入失败：确认安装的是onnxruntime而非onnxruntime-gpu，且Python版本为3.9（树莓派默认）；
推理结果全为空：检查置信度阈值conf_thres是否过高（建议0.25～0.4），或输入图片是否正常读取；
内存不足：降低输入尺寸至416×416（需同步修改input_shape并重新导出ONNX），或使用yolov5n轻量模型。

📚 课程目录