2. YOLOv5 TensorRT 引擎构建与推理加速
任务目标
将训练好的YOLOv5模型导出为ONNX格式，并使用ONNX Runtime在纯CPU环境下对单张输入图片执行目标检测推理，最终输出检测到的目标类别、置信度与边界框坐标。
环境与依赖
- Python 3.8+
- onnxruntime（CPU版本）：
pip install onnxruntime - opencv-python：
pip install opencv-python - numpy：
pip install numpy - PyTorch（仅用于导出ONNX，推理阶段不需要）：
pip install torch torchvision
实现步骤
本案例使用YOLOv5官方仓库的预训练权重（yolov5s.pt）完成导出与推理。核心流程分为两步：（1）将PyTorch权重导出为ONNX格式；（2）使用ONNX Runtime加载ONNX模型进行CPU推理。以下代码已包含完整的预处理（letterbox缩放）、NMS后处理与结果绘制逻辑，可直接运行验证。
完整代码
import cv2
import numpy as np
import onnxruntime as ort
import torch
# ========== 1. 导出ONNX（仅需执行一次） ==========
def export_onnx():
import torch
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
model.eval()
dummy_input = torch.zeros(1, 3, 640, 640)
torch.onnx.export(model.model, dummy_input, "yolov5s.onnx",
opset_version=12, input_names=['images'],
output_names=['output'])
print("ONNX导出完成: yolov5s.onnx")
# ========== 2. ONNX Runtime CPU推理 ==========
class YOLOv5ONNX:
def __init__(self, onnx_path, conf_thres=0.25, iou_thres=0.45):
self.session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
self.conf_thres = conf_thres
self.iou_thres = iou_thres
self.input_shape = self.session.get_inputs()[0].shape[2:] # [640, 640]
self.class_names = ['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 letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114)):
shape = img.shape[:2]
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
dw, dh = dw // 2, dh // 2
img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = dh, dh + (new_shape[0] - new_unpad[1]) % 2
left, right = dw, dw + (new_shape[1] - new_unpad[0]) % 2
img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
return img, r, dw, dh
def preprocess(self, img_path):
img0 = cv2.imread(img_path)
img, ratio, dw, dh = self.letterbox(img0)
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR->RGB, HWC->CHW
img = np.ascontiguousarray(img).astype(np.float32) / 255.0
img = np.expand_dims(img, axis=0)
return img, img0, ratio, dw, dh
def postprocess(self, output, ratio, dw, dh, orig_shape):
# output shape: [1, 25200, 85]
preds = output[0] # (25200, 85)
boxes = preds[:, :4]
scores = preds[:, 4:]
# 获取每个box的最高类别分数
class_scores = scores.max(axis=1)
class_ids = scores.argmax(axis=1)
conf_mask = class_scores > self.conf_thres
boxes, class_scores, class_ids = boxes[conf_mask], class_scores[conf_mask], class_ids[conf_mask]
if len(boxes) == 0:
return []
# 将xywh格式转为xyxy
xyxy = np.zeros_like(boxes)
xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
# 坐标还原到原图尺寸
xyxy[:, [0, 2]] = (xyxy[:, [0, 2]] - dw) / ratio
xyxy[:, [1, 3]] = (xyxy[:, [1, 3]] - dh) / ratio
xyxy = np.clip(xyxy, 0, [orig_shape[1], orig_shape[0], orig_shape[1], orig_shape[0]])
# NMS
indices = cv2.dnn.NMSBoxes(xyxy.tolist(), class_scores.tolist(), self.conf_thres, self.iou_thres)
results = []
if len(indices) > 0:
for i in indices.flatten():
x1, y1, x2, y2 = xyxy[i].astype(int)
results.append({
'bbox': [x1, y1, x2, y2],
'score': float(class_scores[i]),
'class_id': int(class_ids[i]),
'class_name': self.class_names[int(class_ids[i])]
})
return results
def infer(self, img_path):
img, img0, ratio, dw, dh = self.preprocess(img_path)
output = self.session.run(None, {self.session.get_inputs()[0].name: img})
results = self.postprocess(output[0], ratio, dw, dh, img0.shape[:2])
return results, img0
def draw_results(self, img, results):
for r in results:
x1, y1, x2, y2 = r['bbox']
label = f"{r['class_name']} {r['score']:.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), 2)
return img
# ========== 3. 运行入口 ==========
if __name__ == '__main__':
# 首次运行先导出ONNX（若已有yolov5s.onnx可注释掉）
# export_onnx()
detector = YOLOv5ONNX('yolov5s.onnx')
results, img0 = detector.infer('bus.jpg') # 使用YOLOv5官方示例图片
img0 = detector.draw_results(img0, results)
print(f"检测到 {len(results)} 个目标:")
for r in results:
print(f" {r['class_name']}: {r['score']:.2f}, bbox={r['bbox']}")
cv2.imwrite('output.jpg', img0)
print("结果已保存至 output.jpg")运行与验证
- 准备测试图片：下载YOLOv5官方示例图片
bus.jpg（包含行人与公交车），放置于脚本同目录。 - 导出ONNX模型：首次运行时取消
export_onnx()的注释，执行一次生成yolov5s.onnx（约28MB）。后续推理可注释掉该行。 - 执行推理：运行脚本，控制台将打印检测到的目标列表。
- 判定通过标准：
output.jpg中正确绘制出绿色边界框，且至少检测到图片中的公交车（bus）和行人（person），置信度均高于0.25。控制台输出的类别名称与边界框坐标与图片内容一致。
关键说明
- ONNX Runtime的CPU推理无需GPU，本案例完全在CPU上完成前处理、推理与后处理。
- 导出ONNX时使用
opset_version=12，兼容onnxruntime 1.x全系列版本。 - 后处理中的letterbox还原逻辑确保检测框坐标映射回原始图像尺寸，不会出现偏移。
- 若需测试其他图片，替换
infer()中的图片路径即可，无需修改其他代码。