2026/2/18 16:52:09
网站建设
项目流程
网站出现风险如何处理,无锡关键词优化价格,咖啡色网站模板,建立医防融合工作的机制不包括导出ONNX后如何用Python加载#xff1f;代码示例来了
1. 为什么需要加载ONNX模型#xff1f;
你已经在WebUI里点了几下#xff0c;成功导出了model_800x800.onnx——但接下来呢#xff1f; 导出只是第一步#xff0c;真正让模型“活起来”的#xff0c;是把它加载进Pyt…导出ONNX后如何用Python加载代码示例来了1. 为什么需要加载ONNX模型你已经在WebUI里点了几下成功导出了model_800x800.onnx——但接下来呢导出只是第一步真正让模型“活起来”的是把它加载进Python环境跑通推理流程。很多用户卡在这一步下载了ONNX文件却不知道怎么读取照着文档改了代码报错说输入维度不对图片预处理结果和模型期待的格式对不上输出全是空列表甚至不确定该用onnxruntime还是torch.onnx该装CPU版还是GPU版。别急。这篇不是理论课不讲ONNX规范、不画计算图、不推导IR转换逻辑。我们只做一件事用最简明的步骤把cv_resnet18_ocr-detection这个OCR文字检测模型在本地Python环境中稳稳跑起来。所有代码可直接复制、粘贴、运行适配你从WebUI导出的ONNX文件。2. 前置准备安装必要依赖2.1 安装核心库一行搞定pip install onnxruntime opencv-python numpy说明onnxruntime是加载和运行ONNX模型的工业级引擎轻量、跨平台、支持CPU/GPUopencv-python负责图像读取、缩放、通道变换等预处理numpy是数据流转的底层支撑无需额外指定版本当前主流版本均兼容。避坑提示如果你有NVIDIA GPU且想启用CUDA加速请安装带GPU支持的版本pip install onnxruntime-gpu安装后onnxruntime.InferenceSession会自动识别CUDA设备无需修改代码。2.2 验证安装是否成功新建一个test_env.py运行以下代码import onnxruntime as ort print(ONNX Runtime版本:, ort.__version__) print(可用提供器:, ort.get_available_providers())正常输出应类似ONNX Runtime版本: 1.19.2 可用提供器: [CUDAExecutionProvider, CPUExecutionProvider]出现CUDAExecutionProvider表示GPU已就绪即使只有CPUExecutionProvider模型也能正常运行只是速度稍慢。3. 加载ONNX模型三行核心代码3.1 最简加载方式推荐新手import onnxruntime as ort # ① 指定ONNX模型路径替换为你实际下载的文件 model_path model_800x800.onnx # ② 创建推理会话自动选择最优执行提供器 session ort.InferenceSession(model_path, providers[CUDAExecutionProvider, CPUExecutionProvider]) # ③ 查看模型输入信息调试必备 input_name session.get_inputs()[0].name input_shape session.get_inputs()[0].shape print(f模型输入名: {input_name}) print(f模型期望输入形状: {input_shape})输出示例模型输入名: input 模型期望输入形状: [1, 3, 800, 800]这说明模型要求输入是1张、3通道RGB、高800、宽800的图片。关键认知WebUI中你设置的“输入高度800×宽度800”直接决定了ONNX模型的固定输入尺寸。不能传640×480的图也不能传[1,3,1024,1024]——必须严格匹配。后面预处理环节就是为这个目标服务的。4. 图片预处理让输入“长得像”模型想要的样子4.1 预处理四步法缺一不可步骤操作为什么必须做① 读取cv2.imread()OpenCV默认BGR而训练时用的是RGB需转换② 缩放cv2.resize()强制拉到模型要求的800×800或你导出时设的尺寸③ 通道变换.transpose(2,0,1)把HWC→CHW符合ONNX标准输入格式④ 归一化/ 255.0训练时像素值被归一化到[0,1]推理必须保持一致4.2 完整预处理函数可直接复用import cv2 import numpy as np def preprocess_image(image_path, target_height800, target_width800): 将任意图片预处理为ONNX模型可接受的输入格式 Args: image_path (str): 图片文件路径 target_height (int): 模型输入高度如800 target_width (int): 模型输入宽度如800 Returns: np.ndarray: shape[1,3,H,W]dtypefloat32值域[0,1] # ① 读取图片BGR格式 img cv2.imread(image_path) if img is None: raise ValueError(f无法读取图片: {image_path}) # ② BGR → RGB关键 img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # ③ 缩放到目标尺寸直接拉伸不保持长宽比 # 注意这是cv_resnet18_ocr-detection模型的要求与PaddleOCR DBNet一致 img_resized cv2.resize(img, (target_width, target_height)) # ④ HWC → CHW 扩展batch维度 归一化 img_tensor img_resized.transpose(2, 0, 1) # [H,W,C] → [C,H,W] img_tensor np.expand_dims(img_tensor, axis0) # [C,H,W] → [1,C,H,W] img_tensor img_tensor.astype(np.float32) / 255.0 # 归一化到[0,1] return img_tensor # 使用示例 input_blob preprocess_image(test.jpg, target_height800, target_width800) print(预处理后输入形状:, input_blob.shape) # 应输出: (1, 3, 800, 800)重要提醒不要使用cv2.INTER_AREA或cv2.INTER_LANCZOS4等高质量插值——该模型训练时用的就是双线性插值cv2.INTER_LINEAR保持一致才能保证效果稳定不要做直方图均衡、锐化等增强操作——模型没见过可能引入噪声如果你导出的是640x640模型请把target_height和target_width都改为640。5. 执行推理获取原始检测输出5.1 一行代码完成推理# 假设 session 已创建input_blob 已准备好 outputs session.run(None, {input: input_blob})session.run()返回一个Python列表顺序对应模型输出节点定义。对于cv_resnet18_ocr-detection典型输出为outputs[0]: 文本区域概率图shrink mapoutputs[1]: 阈值图threshold mapoutputs[2]: 二值化预测图binary map这和PaddleOCR DBNet的输出结构完全一致也是WebUI能解析JSON结果的基础。5.2 查看原始输出形状调试黄金法则for i, out in enumerate(outputs): print(f输出{i}: shape{out.shape}, dtype{out.dtype})典型输出输出0: shape(1, 1, 800, 800), dtypefloat32 输出1: shape(1, 1, 800, 800), dtypefloat32 输出2: shape(1, 1, 800, 800), dtypefloat32三个图都是单通道、800×800值域在[0,1]之间。6. 后处理从概率图到文本框坐标DBNet核心逻辑6.1 为什么不能跳过后处理ONNX模型输出的是热力图不是最终坐标。就像医生拍完CT得到的是像素强度分布图而你需要的是“这里有个肿瘤坐标是(x1,y1,x2,y2,x3,y3,x4,y4)”。后处理就是那个“影像科医生”。cv_resnet18_ocr-detection基于DBNetDifferentiable Binarization其后处理包含三步二值化用自适应阈值将概率图转为0/1二值图轮廓提取找连通区域过滤小噪点多边形拟合把不规则轮廓拟合成4点文本框。6.2 轻量级后处理实现无依赖纯NumPyimport numpy as np import cv2 def db_postprocess(binary_map, threshold_map, shrink_map, binary_thresh0.3, box_thresh0.6, unclip_ratio1.5): DBNet后处理从三张图生成文本框坐标 Args: binary_map: 二值化预测图 (1,1,H,W) threshold_map: 阈值图 (1,1,H,W) shrink_map: 收缩图 (1,1,H,W) binary_thresh: 二值化阈值0~1 box_thresh: 文本框置信度阈值0~1 unclip_ratio: 扩展比例控制框大小 Returns: List[np.ndarray]: 每个元素是4×2的文本框顶点坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]] # 去除batch和channel维度 binary np.squeeze(binary_map) # (H,W) threshold np.squeeze(threshold_map) shrink np.squeeze(shrink_map) # ① 二值化binary (shrink threshold * scale (1-scale) * threshold) # DBNet经典公式简化为binary shrink (threshold * binary_thresh (1-binary_thresh)*0.5) approx_binary shrink (threshold * binary_thresh (1 - binary_thresh) * 0.5) # ② 提取连通区域 contours, _ cv2.findContours( (approx_binary * 255).astype(np.uint8), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE ) boxes [] for contour in contours: # 过滤太小的区域避免噪点 if cv2.contourArea(contour) 10: continue # 拟合最小外接矩形旋转框 rect cv2.minAreaRect(contour) box cv2.boxPoints(rect) # ③ unclip按比例向外扩展 box unclip(box, unclip_ratio) # ④ 过滤计算该框在shrink图上的平均置信度 score box_score_fast(shrink, box) if score box_thresh: continue boxes.append(box) return boxes def unclip(box, unclip_ratio): 扩展文本框 poly Polygon(box) distance poly.area * unclip_ratio / poly.length offset pyclipper.PyclipperOffset() offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON) expanded offset.Execute(distance) if len(expanded) 0: return box return np.array(expanded[0]).astype(np.int32) def box_score_fast(bitmap, box): 计算box区域在bitmap上的平均值 h, w bitmap.shape[:2] box np.clip(box, 0, [w-1, h-1]).astype(np.int32) mask np.zeros((h, w), dtypenp.uint8) cv2.fillPoly(mask, [box], 1) return cv2.mean(bitmap, maskmask)[0] # 注意上面unclip需要pyclipper如未安装请先运行 # pip install pyclipper更实用的简化版推荐首次运行如果你只想快速看到效果用OpenCV自带的cv2.boundingRect获取轴对齐矩形非旋转框def simple_postprocess(binary_map, box_thresh0.5): 极简后处理只返回轴对齐矩形框 binary np.squeeze(binary_map) box_thresh binary (binary * 255).astype(np.uint8) contours, _ cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) boxes [] for cnt in contours: if cv2.contourArea(cnt) 20: # 过滤小噪点 continue x, y, w, h cv2.boundingRect(cnt) # 转为4点格式[[x,y], [xw,y], [xw,yh], [x,yh]] boxes.append(np.array([[x,y], [xw,y], [xw,yh], [x,yh]], dtypenp.int32)) return boxes # 使用 boxes simple_postprocess(outputs[2]) # 用binary map做简单后处理 print(f检测到 {len(boxes)} 个文本框)7. 完整端到端示例从图片到可视化结果7.1 一键运行的完整脚本# ocr_inference.py import cv2 import numpy as np import onnxruntime as ort def preprocess_image(image_path, target_height800, target_width800): img cv2.imread(image_path) if img is None: raise ValueError(f无法读取图片: {image_path}) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img_resized cv2.resize(img, (target_width, target_height)) img_tensor img_resized.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32) / 255.0 return img_tensor def simple_postprocess(binary_map, box_thresh0.5): binary np.squeeze(binary_map) box_thresh binary (binary * 255).astype(np.uint8) contours, _ cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) boxes [] for cnt in contours: if cv2.contourArea(cnt) 20: continue x, y, w, h cv2.boundingRect(cnt) boxes.append(np.array([[x,y], [xw,y], [xw,yh], [x,yh]], dtypenp.int32)) return boxes def draw_boxes(image_path, boxes, output_pathresult.jpg): img cv2.imread(image_path) for i, box in enumerate(boxes): # 绘制四边形绿色线宽2 cv2.polylines(img, [box], isClosedTrue, color(0,255,0), thickness2) # 标注序号白色文字 cv2.putText(img, str(i1), tuple(box[0]), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,255), 1) cv2.imwrite(output_path, img) print(f结果已保存至: {output_path}) # 主流程 if __name__ __main__: # 1. 加载模型 session ort.InferenceSession(model_800x800.onnx) # 2. 预处理 input_blob preprocess_image(test.jpg, 800, 800) # 3. 推理 outputs session.run(None, {input: input_blob}) # 4. 后处理用binary map boxes simple_postprocess(outputs[2], box_thresh0.3) # 5. 可视化 draw_boxes(test.jpg, boxes) # 6. 打印坐标供后续OCR识别使用 print(检测框坐标x,y:) for i, box in enumerate(boxes): print(f{i1}. {box.tolist()})7.2 运行效果假设你有一张商品截图test.jpg运行后生成result.jpg清晰标出所有文字区域控制台输出类似检测到 5 个文本框 检测框坐标x,y: 1. [[21, 732], [782, 735], [780, 786], [20, 783]] 2. [[105, 620], [320, 625], [318, 670], [103, 665]] ...这些坐标正是WebUI中“检测框坐标 (JSON)”字段的来源。你可以把它们传给另一个OCR识别模型如CRNN、SVTR完成“检测→识别”全流程。8. 常见问题速查表问题现象可能原因解决方案ValueError: Input tensor has incorrect dimensions输入图片尺寸≠模型期望尺寸检查preprocess_image中target_height/width是否与导出时设置一致RuntimeError: CUDA errorGPU显存不足或驱动不匹配改用CPUsession ort.InferenceSession(..., providers[CPUExecutionProvider])检测结果为空boxes[]box_thresh设得太高降低simple_postprocess中的box_thresh尝试0.1~0.3框体严重变形、错位图片通道顺序错误BGR未转RGB确保cv2.cvtColor(img, cv2.COLOR_BGR2RGB)已执行推理速度极慢CPU下5秒ONNX Runtime未启用优化安装最新版pip install --upgrade onnxruntime或启用--enable_onnx_checker重导出输出坐标超出图片范围后处理未做坐标截断在draw_boxes前加box np.clip(box, 0, [img_w-1, img_h-1])9. 进阶建议让部署更工程化9.1 批量处理多张图片from pathlib import Path def batch_inference(image_dir, model_path, output_diroutputs): session ort.InferenceSession(model_path) image_paths list(Path(image_dir).glob(*.jpg)) list(Path(image_dir).glob(*.png)) for img_path in image_paths: input_blob preprocess_image(str(img_path)) outputs session.run(None, {input: input_blob}) boxes simple_postprocess(outputs[2]) draw_boxes(str(img_path), boxes, str(Path(output_dir) / fresult_{img_path.stem}.jpg)) # 使用 batch_inference(input_images/, model_800x800.onnx)9.2 封装为可调用函数供Flask/FastAPI集成def ocr_detect(image_bytes: bytes) - dict: 接收图片字节流返回检测结果字典 nparr np.frombuffer(image_bytes, np.uint8) img cv2.imdecode(nparr, cv2.IMREAD_COLOR) # ...中间预处理、推理、后处理 return { boxes: [box.tolist() for box in boxes], count: len(boxes), inference_time_ms: int((end-start)*1000) } # FastAPI示例 # app.post(/detect) # async def detect(file: UploadFile File(...)): # result ocr_detect(await file.read()) # return result9.3 模型量化减小体积提升CPU速度# 安装工具 pip install onnxsim onnxruntime-tools # 优化并量化INT8 python -m onnxruntime_tools.optimizer_cli \ --input model_800x800.onnx \ --output model_800x800_opt.onnx \ --float16 \ --opt_level 99量化后模型体积减少约50%CPU推理速度提升1.5~2倍精度损失1%对OCR检测任务可忽略。10. 总结你已经掌握了ONNX落地的核心链路回顾一下从WebUI导出ONNX到本地Python运行你完成了环境确认装对了onnxruntime知道怎么选CPU/GPU模型加载三行代码创建InferenceSession并验证输入形状数据对齐用OpenCV精准完成BGR→RGB、缩放、CHW变换、归一化推理执行session.run()拿到三张热力图结果解码用simple_postprocess把概率图变成坐标框效果验证画框打印坐标和WebUI输出完全一致。这不仅是“跑通一个模型”更是建立了一套可复用的ONNX推理范式换任何OCR检测模型DBNet、PSENet、TextSnake只要输出结构相似后处理逻辑几乎不用改换成图像分类、目标检测模型只需调整预处理和后处理加载和推理代码一模一样。下一步你可以 把检测框送给识别模型构建端到端OCR流水线 将脚本封装成API服务供业务系统调用 在Jetson Nano或树莓派上部署实现边缘OCR。技术没有玄学只有清晰的步骤和可验证的结果。你已经走完了最难的第一公里。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。