2. 使用ONNX Runtime在Jetson Nano上运行YOLOv5
任务目标
将YOLOv5s模型导出为ONNX格式，并在Jetson Nano上使用ONNX Runtime完成目标检测推理，在输入图像上绘制边界框与类别标签，输出标注后的图像文件。
环境与依赖
- Jetson Nano（JetPack 4.6+，aarch64架构）
- Python 3.6+
- onnxruntime-gpu（或onnxruntime，若仅用CPU）
- PyTorch 1.8+（仅导出阶段需要）
- OpenCV、NumPy
- YOLOv5官方仓库（导出用，
git clone https://github.com/ultralytics/yolov5）
关键实现
导出阶段在x86 PC或Jetson上均可执行，推理阶段在Jetson Nano上运行。以下代码为完整推理脚本，包含ONNX模型加载、输入预处理、NMS后处理与画框逻辑。
import cv2
import numpy as np
import onnxruntime as ort
# ---------- 配置 ----------
ONNX_PATH = "yolov5s.onnx"
IMG_PATH = "bus.jpg"
OUT_PATH = "output.jpg"
CONF_THRES = 0.5
IOU_THRES = 0.45
INPUT_SIZE = 640
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 preprocess(img):
h, w = img.shape[:2]
r = INPUT_SIZE / max(h, w)
new_w, new_h = int(round(w * r)), int(round(h * r))
resized = cv2.resize(img, (new_w, new_h))
canvas = np.full((INPUT_SIZE, INPUT_SIZE, 3), 114, dtype=np.uint8)
canvas[:new_h, :new_w] = resized
blob = canvas[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0
return blob, r, new_w, new_h
# ---------- NMS ----------
def nms(boxes, scores, iou_thres):
x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
areas = (x2 - x1) * (y2 - y1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
iou = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(iou <= iou_thres)[0]
order = order[inds + 1]
return keep
# ---------- 后处理 ----------
def postprocess(output, r, new_w, new_h):
# output shape: (1, 25200, 85)
preds = output[0][0] # (25200, 85)
boxes_xywh = preds[:, :4]
obj_conf = preds[:, 4]
cls_conf = preds[:, 5:]
scores = obj_conf * cls_conf.max(axis=1)
cls_ids = cls_conf.argmax(axis=1)
mask = scores > CONF_THRES
if not mask.any():
return []
boxes_xywh = boxes_xywh[mask]
scores = scores[mask]
cls_ids = cls_ids[mask]
# xywh -> xyxy (letterbox坐标)
x_center, y_center, w, h = boxes_xywh.T
x1 = x_center - w / 2
y1 = y_center - h / 2
x2 = x_center + w / 2
y2 = y_center + h / 2
boxes = np.stack([x1, y1, x2, y2], axis=1)
# 去除letterbox填充
boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]] / r, 0, new_w / r)
boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]] / r, 0, new_h / r)
# 按类别分别NMS
final_boxes, final_scores, final_cls = [], [], []
for c in np.unique(cls_ids):
idx = np.where(cls_ids == c)[0]
keep = nms(boxes[idx], scores[idx], IOU_THRES)
for k in keep:
final_boxes.append(boxes[idx[k]])
final_scores.append(scores[idx[k]])
final_cls.append(c)
return final_boxes, final_scores, final_cls
# ---------- 主流程 ----------
def main():
# 加载ONNX模型
sess = ort.InferenceSession(ONNX_PATH, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
input_name = sess.get_inputs()[0].name
# 读取并预处理图像
img = cv2.imread(IMG_PATH)
blob, r, new_w, new_h = preprocess(img)
# 推理
outputs = sess.run(None, {input_name: blob[None, ...]})
# 后处理
results = postprocess(outputs, r, new_w, new_h)
if not results:
print("未检测到目标")
return
boxes, scores, cls_ids = results
# 绘制
for box, score, cls_id in zip(boxes, scores, cls_ids):
x1, y1, x2, y2 = box.astype(int)
label = f"{CLASS_NAMES[cls_id]} {score:.2f}"
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(img, label, (x1, max(0, y1 - 5)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
cv2.imwrite(OUT_PATH, img)
print(f"检测完成，结果已保存至 {OUT_PATH}，共 {len(boxes)} 个目标")
if __name__ == "__main__":
main()
模型导出命令
在YOLOv5仓库目录下执行（需先安装依赖 pip install -r requirements.txt）：
python export.py --weights yolov5s.pt --include onnx --img-size 640
导出后得到yolov5s.onnx文件，将其与测试图像bus.jpg放在同一目录下。
运行与验证
- 安装推理依赖：
pip install onnxruntime-gpu opencv-python numpy - 运行脚本：
python detect_onnx.py - 通过条件：程序正常退出，输出
output.jpg，图像中目标（如行人、公交车）被绿色矩形框标注，框上方显示类别名称与置信度。
常见问题
- 若CUDA不可用，将providers改为
['CPUExecutionProvider']，推理速度会下降但功能正常。 - 导出时若报错，确认PyTorch版本与YOLOv5仓库版本匹配（建议使用官方最新release分支）。
- 检测框偏移时，检查预处理中letterbox填充是否与后处理裁剪逻辑一致。