YOLO 模型 ONNX 导出与跨平台部署流程
1. 模型导出(YOLO .pt → ONNX)
-
导出工具:使用 YOLO 框架(如 PyTorch YOLOv5/v7/v8)自带的导出脚本。
-
示例命令 :
bashexport model=best.pt format=onnx task=detect dynamic=True opset=17
-
输出文件 :生成
.onnx
格式的模型文件,支持跨平台部署(如 ONNX Runtime、TensorRT、OpenVINO 等)。
2. 预处理流程
python
import numpy as np
import cv2
def preprocessing(image_path):
input_height, input_width = 1280, 1280 # 模型输入尺寸
original_image = cv2.imread(image_path)
orig_height, orig_width = original_image.shape[:2]
global scale # 全局变量用于后处理
scale = min(input_height / orig_height, input_width / orig_width)
new_height = int(orig_height * scale)
new_width = int(orig_width * scale)
# 等比缩放并填充为正方形
resized = cv2.resize(original_image, (new_width, new_height))
padded_image = np.zeros((input_height, input_width, 3), dtype=np.uint8)
top = (input_height - new_height) // 2
bottom = input_height - new_height - top
left = (input_width - new_width) // 2
right = input_width - new_width - left
padded_image[top:top+new_height, left:left+new_width] = resized
# 归一化与通道转换
image_data = padded_image.transpose(2, 0, 1) # HWC → CHW
image_data = np.expand_dims(image_data, axis=0).astype(np.float32) / 255.0 # 归一化到 [0,1]
return image_data
- 关键步骤 :
- 等比缩放:保持宽高比缩放图像。
- 填充为正方形 :使用黑色填充(
np.zeros
)确保输入尺寸为(1280, 1280)
。 - 通道转换 :
HWC → CHW
(符合 ONNX 模型输入格式)。 - 归一化 :除以 255 将像素值归一化到
[0, 1]
。
3. ONNX 推理
python
import onnxruntime as ort
def predict(model_path, image_data):
session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
outputs = session.run([output_name], {input_name: image_data})
return outputs[0] # 输出形状: (1, 5, 33600)
- 输入 :预处理后的
image_data
(形状[1, 3, 1280, 1280]
)。 - 输出 :检测结果
detections
(形状[1, 5, 33600]
),每行包含[x_center, y_center, width, height, confidence]
。
4. 后处理与结果判断
python
def isdetected(boxes):
confidence_threshold = 0.1
if np.max(boxes[:, 4]) < confidence_threshold:
print("未检测到任何目标")
else:
max_conf_index = np.argmax(boxes[:, 4])
best_box = boxes[max_conf_index]
print(f"检测到目标,最高置信度: {np.max(boxes[:, 4]):.4f}")
print(f"检测框数据: {best_box}")
- 逻辑 :
- 过滤置信度低于阈值的检测框。
- 打印置信度最高的检测框信息。
5. 检测框绘制
python
def draw_boxes(image_path, boxes, confidence_threshold=0.5):
global scale # 使用预处理阶段的缩放比例
original_image = cv2.imread(image_path)
image = original_image.copy()
filtered_boxes = boxes[boxes[:, 4] >= confidence_threshold]
for box in filtered_boxes:
x_center, y_center, width, height, confidence = box
# 坐标反向转换(从模型输出到原始图像坐标)
x_center = x_center / scale
y_center = (y_center - 280) / scale # ← 注意:此处可能需要修正
width = width / scale
height = height / scale
x_min = int(x_center - width / 2)
y_min = int(y_center - height / 2)
x_max = int(x_center + width / 2)
y_max = int(y_center + height / 2)
# 绘制矩形框与标签
cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
cv2.putText(image, f"{confidence:.2f}", (x_min, y_min - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
cv2.imshow("Detection", image)
cv2.waitKey(0)
cv2.destroyAllWindows()
- 关键点 :
- 坐标反向转换 :根据
scale
和填充信息(top
,left
)将检测框坐标还原到原始图像。
- 坐标反向转换 :根据
6. 完整调用示例
python
# 1. 预处理图像
image_data = preprocessing("example.jpg")
# 2. ONNX 推理
model_path = "best.onnx"
detections = predict(model_path, image_data)
# 3. 判断检测结果
boxes = detections[0].T # 转置为 (33600, 5)
isdetected(boxes)
# 4. 绘制检测框
draw_boxes("example.jpg", boxes)