保姆级教程:用Python和ONNX Runtime部署YOLOv8 OBB旋转框模型(附完整代码)
PythonONNX Runtime实战YOLOv8旋转框模型部署全流程解析旋转目标检测OBB在遥感图像分析、文档识别、工业质检等领域应用广泛。相比水平框检测旋转框能更精确地定位倾斜物体。本文将手把手教你用ONNX Runtime部署YOLOv8 OBB模型从环境准备到可视化输出覆盖完整技术链路。1. 环境准备与模型转换部署前需要确保基础环境就绪。建议使用Python 3.8环境避免版本兼容性问题pip install onnxruntime-gpu1.16.0 # GPU版本 pip install opencv-python4.8.0 pip install numpy1.23.5YOLOv8官方支持直接导出OBB模型为ONNX格式from ultralytics import YOLO model YOLO(yolov8n-obb.pt) # 加载预训练OBB模型 model.export(formatonnx, dynamicFalse, opset12) # 导出静态形状ONNX关键参数说明dynamicFalse固定输入输出维度提升推理效率opset12确保旋转框相关算子兼容性注意导出时建议固定输入尺寸如640x640避免动态尺寸带来的预处理复杂度2. 核心代码模块解析2.1 图像预处理标准化YOLOv8需要特定的letterbox预处理保持宽高比def letterbox(im, new_shape(640, 640), color(114, 114, 114)): # 原始图像尺寸 shape im.shape[:2] # 计算缩放比例保持长宽比 ratio min(new_shape[0]/shape[0], new_shape[1]/shape[1]) # 新尺寸计算 new_unpad int(round(shape[1]*ratio)), int(round(shape[0]*ratio)) dw, dh (new_shape[1]-new_unpad[0])/2, (new_shape[0]-new_unpad[1])/2 # 添加灰边填充 im cv2.resize(im, new_unpad, interpolationcv2.INTER_LINEAR) im cv2.copyMakeBorder(im, int(dh), int(dh), int(dw), int(dw), cv2.BORDER_CONSTANT, valuecolor) return im预处理后的图像需要转换为CHW格式并归一化image letterbox(src_img, (640, 640)) image image[:, :, ::-1].transpose(2, 0, 1) # BGR→RGB, HWC→CHW image np.ascontiguousarray(image, dtypenp.float32) / 255.02.2 ONNX Runtime推理引擎初始化针对不同硬件配置优化执行提供器providers [ (CUDAExecutionProvider, { device_id: 0, arena_extend_strategy: kNextPowerOfTwo, gpu_mem_limit: 4 * 1024 * 1024 * 1024, # 4GB cudnn_conv_algo_search: EXHAUSTIVE, do_copy_in_default_stream: True, }), CPUExecutionProvider ] session onnxruntime.InferenceSession( yolov8n-obb.onnx, providersproviders )输入输出张量名称可通过以下方式获取input_name session.get_inputs()[0].name output_names [output.name for output in session.get_outputs()]2.3 旋转框后处理详解YOLOv8 OBB输出包含5个关键参数中心点坐标 (cx, cy)宽高 (w, h)旋转角度 (θ)典型输出张量结构维度含义0:4框坐标 (cx,cy,w,h)4:4nc类别置信度-1旋转角度角度处理需要特别注意单位转换def process_angle(angle_rad): # 将弧度限制在[-π/2, π/2]范围内 angle_rad angle_rad % math.pi if angle_rad math.pi/2: angle_rad - math.pi return angle_rad旋转框NMS实现需要特殊处理def rotated_nms(boxes, scores, threshold0.5): 基于ProbIoU的旋转框NMS实现 :param boxes: [N,5] 包含cx,cy,w,h,angle :param scores: [N,] 置信度分数 :param threshold: IoU阈值 :return: 保留的索引 order scores.argsort()[::-1] keep [] while order.size 0: i order[0] keep.append(i) # 计算当前框与剩余框的ProbIoU ious [probiou(boxes[i], boxes[j]) for j in order[1:]] # 保留IoU低于阈值的框 inds np.where(np.array(ious) threshold)[0] order order[inds 1] return keep3. 完整部署流程实现整合各模块的端到端流程class YOLOv8OBB: def __init__(self, onnx_path, class_names): self.session onnxruntime.InferenceSession(onnx_path) self.class_names class_names self.input_size (640, 640) def predict(self, image_path): # 1. 图像预处理 orig_img cv2.imread(image_path) img self.preprocess(orig_img) # 2. ONNX推理 outputs self.session.run( None, {self.session.get_inputs()[0].name: img} ) # 3. 后处理 boxes self.postprocess(outputs, orig_img.shape) # 4. 可视化 result self.visualize(orig_img, boxes) return result关键参数配置建议参数推荐值说明score_threshold0.25置信度过滤阈值nms_threshold0.45NMS重叠阈值max_det300每图最大检测数4. 性能优化技巧4.1 推理加速方案启用ONNX Runtime的图优化session_options onnxruntime.SessionOptions() session_options.graph_optimization_level ( onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL ) session_options.enable_profiling True # 用于性能分析批处理实现技巧def batch_inference(images): # 堆叠多个预处理后的图像 batch np.stack([preprocess(img) for img in images]) outputs session.run( None, {input_name: batch} ) # 批处理结果解析 return [postprocess(out) for out in outputs]4.2 内存优化策略使用固定内存减少拷贝开销input_tensor np.ascontiguousarray( preprocessed_image, dtypenp.float32 )释放临时变量del intermediate_results # 显式释放大内存对象4.3 多线程处理模式from concurrent.futures import ThreadPoolExecutor def parallel_process(image_paths): with ThreadPoolExecutor(max_workers4) as executor: results list(executor.map( detector.predict, image_paths )) return results5. 常见问题解决方案5.1 角度跳变问题旋转框角度在90度附近可能出现突变解决方案def smooth_angle(current, previous): 角度平滑处理 diff current - previous if diff math.pi/2: current - math.pi elif diff -math.pi/2: current math.pi return current5.2 框裁剪优化旋转框可能超出图像边界需要裁剪def clip_rotated_box(box, img_width, img_height): cx, cy, w, h, angle box # 计算旋转后的四个角点 corners get_rotated_corners(cx, cy, w, h, angle) # 限制在图像范围内 corners[:,0] np.clip(corners[:,0], 0, img_width-1) corners[:,1] np.clip(corners[:,1], 0, img_height-1) # 从裁剪后的角点重新计算旋转框参数 return fit_rotated_rect(corners)5.3 部署架构选择不同场景下的部署方案对比场景推荐方案优势云端部署ONNX Runtime Triton高吞吐量边缘设备ONNX Runtime DirectML跨平台兼容移动端ONNX Runtime Mobile内存优化实际测试数据RTX 3090模型输入尺寸FPS内存占用YOLOv8n-obb640x6401421.2GBYOLOv8s-obb640x640981.5GBYOLOv8m-obb640x640532.4GB