cv_resnet18 ONNX模型如何调用？Python推理代码实例

cv_resnet18 ONNX模型如何调用？Python推理代码实例

1. 模型背景与定位

1.1 什么是cv_resnet18_ocr-detection？

cv_resnet18_ocr-detection 是一个专为中文场景优化的轻量级OCR文字检测模型，由科哥基于ResNet-18主干网络构建。它不负责文字识别（OCR中的Recognition部分），只专注解决“文字在哪”的问题——也就是在任意图片中精准定位所有文字区域的位置。

你可以把它理解成一个“文字雷达”：输入一张图，它会快速扫描出所有可能包含文字的矩形框，并返回每个框的坐标和置信度。后续再交给识别模型（如CRNN、PaddleOCR识别模块）去读取框内具体是什么字。

这个模型特别适合部署在边缘设备或对响应速度有要求的服务中，因为ResNet-18结构简洁、参数量小、推理快，同时在常规文档、截图、商品图等常见OCR场景下保持了不错的检测精度。

1.2 为什么选择ONNX格式？

ONNX（Open Neural Network Exchange）是一种开放的模型格式标准，就像PDF之于文档——它让模型脱离原始训练框架（比如PyTorch），变成一种通用“可执行文件”。这意味着：

• 你可以在没有PyTorch环境的服务器上运行它
• 能用C++、Java、Go甚至JavaScript调用（通过对应runtime）
• 支持TensorRT、ONNX Runtime等高性能推理引擎加速
• WebUI中的“ONNX导出”功能，正是为了帮你把训练好的模型打包成这种跨平台、易部署的形态

简单说：ONNX不是新模型，而是同一个cv_resnet18_ocr-detection模型的“便携安装包”。

2. ONNX模型获取与验证

2.1 从WebUI导出ONNX模型

根据手册第六章，你已在WebUI中完成ONNX导出操作。导出后你会得到类似这样的文件：

model_800x800.onnx

文件名中的800x800表示该模型固定接受800×800像素的输入图像。这是关键信息——调用时必须严格匹配尺寸，否则会报错或结果异常。

验证小技巧：用文本编辑器打开.onnx文件，开头几行是可读的ASCII字符，能看到类似ONNX、ir_version等字样，说明不是损坏文件；若打开全是乱码或提示“无法读取”，大概率导出失败，需重试。

2.2 检查模型输入输出结构

在写推理代码前，必须确认模型“长什么样”。我们用ONNX官方工具快速探查：

import onnx # 加载模型 model = onnx.load("model_800x800.onnx") # 打印输入信息 print("=== 输入信息 ===") for inp in model.graph.input: print(f"名称: {inp.name}") print(f"形状: {[dim.dim_value for dim in inp.type.tensor_type.shape.dim]}") print(f"数据类型: {inp.type.tensor_type.elem_type}") # 打印输出信息 print("\n=== 输出信息 ===") for out in model.graph.output: print(f"名称: {out.name}") print(f"形状: {[dim.dim_value for dim in out.type.tensor_type.shape.dim]}")

典型输出如下（以800×800模型为例）：

=== 输入信息 === 名称: input 形状: [1, 3, 800, 800] 数据类型: 1 === 输出信息 === 名称: output 形状: [1, 1, 800, 800] 数据类型: 1

说明：

• 输入张量名为input，形状为[batch=1, channel=3, height=800, width=800]，即单张RGB图，归一化前（值域0–255）
• 输出张量名为output，形状为[1, 1, 800, 800]，是一个单通道热力图（heatmap），值越大代表该位置越可能是文字中心

注意：这里输出不是坐标列表，而是像素级热力图——你需要自己后处理（如阈值分割+连通域分析）才能提取出最终的文本框。

3. 完整Python推理代码详解

3.1 环境准备与依赖安装

确保已安装以下库（推荐使用conda或pip）：

pip install onnxruntime opencv-python numpy # 如需GPU加速（NVIDIA显卡），额外安装： pip install onnxruntime-gpu

小贴士：onnxruntime和onnxruntime-gpu不能共存，装了GPU版就别装CPU版。若不确定，先装CPU版测试通路。

3.2 核心推理代码（含注释）

以下代码完整实现：读图 → 预处理 → 推理 → 后处理 → 可视化，每一步都附带原理说明：

import cv2 import numpy as np import onnxruntime as ort def preprocess_image(image_path, input_size=(800, 800)): """ 图像预处理：读取、缩放、HWC→CHW、归一化、增加batch维度 注意：本模型未做标准化（如减均值除方差），仅做/255.0归一化 """ # 1. 读取BGR格式图像 img = cv2.imread(image_path) if img is None: raise ValueError(f"无法读取图片: {image_path}") # 2. 缩放到模型输入尺寸（保持宽高比？否！本模型要求严格resize） # 因为ONNX导出时固定了输入shape，所以必须拉伸（非填充），否则报错 img_resized = cv2.resize(img, input_size) # (800, 800, 3) # 3. BGR → RGB（OpenCV默认BGR，但模型训练时用的是RGB） img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB) # 4. HWC → CHW + 归一化 + 增加batch维度 input_blob = img_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0 # (3, 800, 800) input_blob = np.expand_dims(input_blob, axis=0) # (1, 3, 800, 800) return input_blob, img # 返回原始图用于画框 def postprocess_heatmap(heatmap, original_img, threshold=0.2, min_area=100): """ 热力图后处理：二值化 → 连通域 → 外接矩形 → 过滤小区域 heatmap: (1, 1, 800, 800) → squeeze成(800, 800) original_img: 原始图，用于按比例还原坐标 """ # 1. 提取单通道并压缩维度 heat = np.squeeze(heatmap) # (800, 800) # 2. 二值化（大于threshold设为255，否则0） _, binary = cv2.threshold(heat, threshold, 1.0, cv2.THRESH_BINARY) # 3. 查找连通域（OpenCV要求uint8） binary_uint8 = (binary * 255).astype(np.uint8) num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(binary_uint8, connectivity=8) boxes = [] for i in range(1, num_labels): # 跳过背景（label 0） x, y, w, h, area = stats[i] if area < min_area: # 过滤太小的噪声块 continue # 还原到原始图尺寸（假设原始图宽高比不同，此处按等比缩放计算） orig_h, orig_w = original_img.shape[:2] scale_x = orig_w / 800.0 scale_y = orig_h / 800.0 box = [ int(x * scale_x), int(y * scale_y), int((x + w) * scale_x), int((y + h) * scale_y) ] boxes.append(box) return boxes def draw_boxes(image, boxes, color=(0, 255, 0), thickness=2): """在图上画矩形框""" img_copy = image.copy() for box in boxes: cv2.rectangle(img_copy, (box[0], box[1]), (box[2], box[3]), color, thickness) return img_copy # —————— 主流程开始 —————— if __name__ == "__main__": # 1. 加载ONNX模型（自动选择CPU/GPU） session = ort.InferenceSession("model_800x800.onnx", providers=['CUDAExecutionProvider', 'CPUExecutionProvider']) # 2. 预处理图像 input_data, original_img = preprocess_image("test.jpg", input_size=(800, 800)) # 3. 执行推理 outputs = session.run(None, {"input": input_data}) heatmap = outputs[0] # shape: (1, 1, 800, 800) # 4. 后处理提取文本框 detected_boxes = postprocess_heatmap( heatmap, original_img, threshold=0.25, # 可根据效果调整，0.2~0.4较常用 min_area=150 # 过滤小于150像素面积的噪声 ) # 5. 可视化结果 result_img = draw_boxes(original_img, detected_boxes) cv2.imwrite("detection_result.jpg", result_img) print(f" 检测完成！共找到 {len(detected_boxes)} 个文本区域") print("坐标格式：[x1, y1, x2, y2]（左上角→右下角）") for i, box in enumerate(detected_boxes): print(f" {i+1}. {box}")

3.3 关键细节说明

• 为什么不用padding而用resize？
  因为导出ONNX时模型输入被硬编码为固定尺寸（800×800），cv2.resize是最直接满足要求的方式。若想支持任意尺寸，需在导出时使用动态轴（-1），但当前WebUI导出的是静态模型。

• 热力图怎么变成框？
  模型输出的是每个像素属于“文字中心”的概率。我们用OpenCV的连通域分析（connectedComponentsWithStats）把高概率区域聚合成块，再取每个块的外接矩形——这就是OCR检测中最常用的“heatmap→box”范式。

• 阈值0.25怎么选？
  它控制热力图二值化的敏感度。值越小，框越多（可能包含误检）；越大，框越少（可能漏检）。建议从0.2开始试，看结果再微调。

• GPU加速开关
  providers=['CUDAExecutionProvider', 'CPUExecutionProvider']表示优先用GPU，失败则降级到CPU。确保你已安装onnxruntime-gpu且CUDA驱动正常。

4. 实际效果调试与优化

4.1 快速验证：三步判断模型是否正常工作

1. 检查输出热力图是否“有反应”
   在postprocess_heatmap函数开头加一句：

   print("热力图统计：", heat.min(), heat.max(), heat.mean())

   正常应输出类似0.0 0.92 0.03—— 说明有明显峰值，不是全黑或全灰。

2. 可视化热力图本身
   临时保存热力图为图片：

   cv2.imwrite("heatmap.jpg", (heat * 255).astype(np.uint8))

   用看图软件打开，应能看到白色斑点（文字区域）分布在图中。

3. 用纯色图测试
   创建一张全白图：

   white = np.ones((800, 800, 3), dtype=np.uint8) * 255 cv2.imwrite("white.jpg", white)

   如果对这张图也检测出大量框，说明阈值设得太低或模型过拟合；如果完全没框，可能是预处理顺序错了（比如忘了BGR→RGB）。

4.2 常见问题与修复方案

5. 进阶用法：集成到业务系统

5.1 批量处理多张图片

只需将主流程封装为函数，并遍历文件列表：

import glob import os def batch_detect(image_dir, output_dir, model_path="model_800x800.onnx"): session = ort.InferenceSession(model_path) image_paths = glob.glob(os.path.join(image_dir, "*.jpg")) + \ glob.glob(os.path.join(image_dir, "*.png")) for img_path in image_paths: try: input_data, orig = preprocess_image(img_path) heatmap = session.run(None, {"input": input_data})[0] boxes = postprocess_heatmap(heatmap, orig, threshold=0.25) # 保存结果图 result_img = draw_boxes(orig, boxes) out_name = os.path.join(output_dir, f"result_{os.path.basename(img_path)}") cv2.imwrite(out_name, result_img) print(f"✓ {img_path} → {out_name}") except Exception as e: print(f"✗ {img_path} 处理失败: {e}") # 使用示例 batch_detect("/data/input/", "/data/output/")

5.2 构建轻量API服务（Flask示例）

from flask import Flask, request, jsonify, send_file import io app = Flask(__name__) session = ort.InferenceSession("model_800x800.onnx") @app.route('/detect', methods=['POST']) def detect(): if 'image' not in request.files: return jsonify({"error": "缺少图片"}), 400 file = request.files['image'] img_bytes = np.frombuffer(file.read(), np.uint8) img = cv2.imdecode(img_bytes, cv2.IMREAD_COLOR) # 预处理 & 推理（同上） input_blob = cv2.resize(img, (800, 800)) input_blob = cv2.cvtColor(input_blob, cv2.COLOR_BGR2RGB) input_blob = input_blob.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32) / 255.0 heatmap = session.run(None, {"input": input_blob})[0] boxes = postprocess_heatmap(heatmap, img, threshold=0.25) return jsonify({ "boxes": boxes, "count": len(boxes) }) if __name__ == '__main__': app.run(host='0.0.0.0', port=5000)

调用方式：

curl -X POST -F "image=@test.jpg" http://localhost:5000/detect

6. 总结

6.1 你已掌握的核心能力

• 理解cv_resnet18_ocr-detection模型的定位：它是专注文字定位的轻量检测器，不是端到端OCR
• 知道ONNX模型的本质：同一模型的跨平台“可执行格式”，导出后无需PyTorch即可运行
• 能独立编写完整推理链：从图像读取、预处理、ONNX加载、热力图推理，到坐标后处理与可视化
• 具备调试能力：通过热力图统计、中间结果保存、阈值调节快速定位问题
• 可扩展落地：批量处理脚本、HTTP API服务，真正融入生产环境

6.2 下一步建议

• 搭配识别模型：将本模型输出的boxes裁剪后，送入CRNN或PaddleOCR识别模型，组成完整OCR流水线
• 尝试量化模型：用ONNX Runtime的量化工具生成INT8模型，进一步提速（尤其在边缘设备）
• 探索动态尺寸：若需支持任意分辨率，可重新导出带-1动态轴的ONNX，再用cv2.resize+letterbox预处理

记住：所有OCR工程落地，第一步永远是“准确定位文字在哪”。而你手上的这个cv_resnet18_ocr-detection ONNX模型，就是那个可靠、快速、开箱即用的起点。

--- > **获取更多AI镜像** > > 想探索更多AI镜像和应用场景？访问 [CSDN星图镜像广场](https://ai.csdn.net/?utm_source=mirror_blog_end)，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。