onnx查看工具 - 技术栈

一、主流ONNX查看工具

整理了几个主流的工具，你可以根据自己的具体任务来挑选：

工具名称	一句话概括	核心功能与特点	适合谁用
Netron	最流行的轻量级可视化工具，支持多种模型格式。	• 直观展示网络结构、节点连接、输入输出维度 • 点击查看单个节点的参数信息（如权重、偏置） • 支持在线版和桌面版（Windows/macOS/Linux），打开即用	几乎所有开发者，用于快速查看和分享模型结构，理解模型构成。
NVIDIA Nsight Deep Learning Designer	NVIDIA出品的强大集成式IDE，专为ONNX设计。	• 可视化与编辑：像IDE一样创建和编辑ONNX图（增删节点、修改参数） • 性能分析：集成TensorRT和ONNX Runtime，支持GPU性能剖析 • 模型优化：提供模型转换（如FP16）、子图提取等高级功能	面向高级开发者和性能调优工程师，需要在NVIDIA GPU上进行深度编辑和优化。
AMD AI Analyzer	AMD官方工具，专注于分析模型在Ryzen AI上的运行情况。	• 软硬件协同视图：可视化模型在CPU和NPU之间的算子分配和分区情况 • 性能剖析：分析NPU上的推理性能，查看层执行时间线	使用AMD Ryzen AI硬件进行部署和优化的开发者。
ONNX Runtime 性能分析	ONNX Runtime自带的代码内性能分析功能。	• 生成详细性能JSON ：记录每个算子（Operator）的执行延迟 • 兼容通用工具：生成的JSON文件可用Perfetto 、chrome://tracing等工具打开分析	需要量化性能瓶颈，进行精确延迟分析的开发者。
onnx-vis	一个基于客户端-服务器架构的Python库。	• 安装简单：`pip install onnx-vis` • 通过URL分享模型，方便协作	需要将模型可视化结果分享给他人，或者希望集成到Python工作流中。

如何选择？

只想看一眼模型长什么样？ 选 Netron 最简单。可以直接把模型文件拖到它的在线版网站上，或者下载桌面版。
想在NVIDIA GPU上优化模型性能？ 那么 NVIDIA Nsight Deep Learning Designer 是最佳选择，它提供了从编辑到性能剖析的一站式体验。
如果你是在AMD Ryzen AI平台上部署模型 ，可以了解一下 AMD AI Analyzer，它能让你清楚地看到模型在NPU上的运行情况。
需要定位某个具体算子的性能瓶颈？ 可以试试 ONNX Runtime 自带的性能分析，通过生成的追踪文件来精确定位问题。

二、Netron使用

Netron 的使用非常直观，核心就是"打开即看"。它支持多种打开方式，你可以根据当前的工作场景选择最顺手的一种。

安装与启动

Netron 是跨平台的开源工具，根据你的操作系统和习惯，有几种主流的选择：

使用方式	适用平台	安装/启动命令
桌面版	Windows	下载 `.exe` 文件或运行 `winget install netron`
	macOS	下载 `.dmg` 文件或运行 `brew cask install netron`
	Linux	下载 `.AppImage` 文件或运行 `snap install netron`
Python 集成	任意平台	`pip install netron`，然后在Python脚本中调用。
在线版	浏览器	直接访问 Netron 在线网站，无需安装任何东西。

核心使用步骤

无论哪种方式，最终都会进入同一个图形界面。你可以通过以下三种方法之一打开模型：

最直接：拖拽打开

启动 Netron 后，直接将你的模型文件（如 .onnx、.pt、.tflite 等）从文件夹拖拽到 Netron 窗口中即可。
最经典：菜单打开

在 Netron 界面中，点击菜单栏的 File -> Open，然后在文件系统中选择你的模型文件。
最集成：Python API 打开

如果你希望在 Python 工作流中快速预览模型，这是最方便的方法。安装 netron 库后，创建一个Python脚本或直接在命令行中操作：
- 命令行方式 ：netron 你的模型.onnx
- 脚本方式：
  
  python
  复制代码
```
import netron
# 指定你的模型文件路径
model_path = "你的模型.onnx"
# 启动Netron本地服务器，它会自动在浏览器中打开可视化界面
netron.start(model_path)
```
  执行后，Netron 会在本地启动一个服务，并自动弹出浏览器窗口显示模型。

界面操作与信息查看

打开模型后，你会看到一个清晰的网络结构图。

浏览结构 ：你可以用鼠标滚轮缩放 来查看模型全貌或局部细节，也可以拖拽空白处 来平移画布。对于复杂模型，界面左侧或上方通常会有概览图，帮助你快速定位。
查看详情 ：点击图中的任何一个节点（代表一个算子或一层），右侧面板会立刻显示它的详细信息，例如：
- 输入/输出（INPUTS/OUTPUTS）：张量的名称和维度。
- 属性（ATTRIBUTES） ：该算子的具体参数（如卷积层的 kernel_shape、strides）。
- 权重（Weights/Bias）：如果节点包含可训练参数（如卷积层），这里可以看到具体的数值，甚至可以保存下来。
探索子图 ：如果你的模型包含了子图结构（例如经过某些编译优化后的模型），双击或点击特定的子图节点，就能"进入"该子图内部查看其详细结构。想返回时，点击左上角的返回箭头即可。
导出与分享 ：想保存当前视图？可以通过 File -> Export 将整个模型结构图导出为图片，方便分享或归档。

三、在Netron 中分析模型的输入和输出

在 Netron 中分析模型的输入和输出是理解模型用法的关键步骤。这里详细说明如何查看和解析这些信息：

查看输入（Input）信息

当你用 Netron 打开 ONNX 模型后，找到输入节点的方法：

方法1：从模型入口查看

在模型图的最左侧 ，你会看到输入节点（通常标有 input、images 等名称）
点击该节点，右侧面板会显示详细信息：

text

复制代码

INPUTS
  name: "input"           # 输入张量的名称
  type: float32[1,3,224,224]  # 数据类型和形状

方法2：从模型属性面板查看

点击画布空白处，右侧会显示模型的全局信息
找到 "inputs" 部分，会列出所有输入：

text

复制代码

inputs
  └─ input
        type: float32[1,3,224,224]
        shape: [1,3,224,224]  # [batch, channels, height, width]

查看输出（Output）信息

类似地，查看输出的方法：

从模型出口查看

在模型图的最右侧 ，找到输出节点（通常标有 output、prob 等）
点击输出节点，右侧面板显示：

text

复制代码

OUTPUTS
  name: "prob"            # 输出张量的名称
  type: float32[1,1000]   # 数据类型和形状，通常是分类概率

从全局信息查看

点击画布空白处，在右侧找到 "outputs" 部分：

text

复制代码

outputs
  └─ prob
        type: float32[1,1000]
        shape: [1,1000]   # 1个样本，1000个类别的概率

深入分析输入输出

1. 理解数据形状（Shape）

常见的形状格式及其含义：

形状示例	含义	典型用途
`[1,3,224,224]`	[批大小, 通道数, 高度, 宽度]	图像分类（NCHW格式）
`[1,224,224,3]`	[批大小, 高度, 宽度, 通道数]	图像分类（NHWC格式）
`[1,3,128,128]`	[批大小, 通道数, 高度, 宽度]	目标检测输入
`[1,84,84,3]`	[批大小, 高度, 宽度, 通道数]	某些自定义模型
`[1,512]`	[批大小, 特征长度]	文本或特征向量
`[1,1000]`	[批大小, 类别数]	分类输出

2. 查看数据类型

Netron 会显示每个输入输出的数据类型：

float32：最常用的32位浮点数
float16：半精度浮点数
int32/int64：整数类型（用于索引或计数）
uint8：8位无符号整数（如图像像素）

3. 跟踪数据流动

想知道输入如何变换为输出？可以：

点击输入节点，记下其形状
依次点击后续节点，观察形状如何变化
最终点击输出节点，对比最终形状

如何判断你的模型需要什么范围？

方法1：查看Netron中的模型输入信息

在Netron中点击输入节点，查看是否有以下提示：

如果看到 mean 和 std 参数，通常意味着需要 [0, 1] 或归一化输入
如果是图像模型，常见的有：
- [0, 1]：分类模型（如MobileNet、ResNet）
- [0, 255]：检测模型（如YOLO系列）
- [-1, 1]：某些特殊模型

方法2：查看模型文档或训练代码

训练时的预处理通常决定了推理时的要求：

python

复制代码

# PyTorch 常见预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将 [0,255] 的PIL图像转为 [0,1] 的tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406],  # 进一步归一化
                        std=[0.229, 0.224, 0.225])
])

方法3：检查模型的输入数据类型

在Netron中：

如果输入类型是 float32[1,3,224,224]：通常是 [0,1] 或归一化后的值
如果输入类型是 uint8[1,224,224,3]：通常是 [0,255] 的原始像素值

完整预处理示例

假设你的模型需要 [0,1] 范围的输入：

cpp

复制代码

// 1. 读取图像 (uint8, [0,255])
cv::Mat image = cv::imread(image_path);

// 2. 缩放到模型输入尺寸
cv::Mat resized;
cv::resize(image, resized, cv::Size(input_w, input_h));

// 3. 转换为blob并归一化到 [0,1]
cv::Mat blob = cv::dnn::blobFromImage(
    resized,           // uint8 [0,255]
    1.0/255.0,         // 缩放到 [0,1]
    cv::Size(),        // 已经resize过了，这里可以留空
    cv::Scalar(),      // mean=0
    true,              // swapRB (OpenCV是BGR，很多模型需要RGB)
    false
);

// 4. 如果模型还需要特定的mean/std进一步归一化
// 例如 ImageNet 的 mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]
// 需要手动对每个通道进行操作
std::vector<cv::Mat> channels(3);
cv::split(blob, channels);  // blob是1x3xHxW格式

float mean[] = {0.485f, 0.456f, 0.406f};
float std[] = {0.229f, 0.224f, 0.225f};

for (int c = 0; c < 3; c++) {
    channels[c] = (channels[c] - mean[c]) / std[c];
}

cv::merge(channels, blob);

python

复制代码

import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np

# 定义预处理流程
transform = transforms.Compose([
    transforms.Resize((input_h, input_w)),  # 缩放到目标尺寸
    transforms.ToTensor(),                  # PIL Image [0,255] -> Tensor [0,1]
    transforms.Normalize(                   # 归一化
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# 读取图像
image = Image.open(image_path).convert('RGB')  # 确保RGB格式

# 应用预处理
blob = transform(image)  # shape: (3, H, W)

# 添加batch维度
blob = blob.unsqueeze(0)  # shape: (1, 3, H, W)

# 转换为numpy数组（如果需要用于ONNX Runtime）
blob_numpy = blob.numpy()

实用场景示例

场景1：图像分类模型

text

复制代码

输入：float32[1,3,224,224]  # 224x224的RGB图像
解释：1张图，3通道(RGB)，高224，宽224

输出：float32[1,1000]       # 1000个类别的概率
解释：1张图的预测结果，1000个类别对应的概率值

预处理要求：需要将图像resize到224x224，归一化到[0,1]或特定均值/标准差

场景2：目标检测模型

text

复制代码

输入：float32[1,3,416,416]  # 416x416的RGB图像
输出1：float32[1,2535,4]    # 边界框坐标
输出2：float32[1,2535,80]   # 类别概率
输出3：float32[1,2535]      # 置信度分数

后处理要求：需要进行NMS（非极大值抑制）过滤冗余框

场景3：BERT文本模型

text

复制代码

输入1：int64[1,128]         # input_ids：token IDs
输入2：int64[1,128]         # attention_mask：注意力掩码
输入3：int64[1,128]         # token_type_ids：段落标识
输出：float32[1,768]        # 句子级别的特征向量

预处理要求：需要分词器将文本转换为token IDs

高级技巧

1. 检查动态维度

某些模型支持动态批大小或动态尺寸：

text

复制代码

输入：float32[?,3,?,?]  # ?表示可变维度
含义：批大小和图像尺寸可变

2. 保存输入输出信息

右键点击节点 → "Copy" → 可粘贴到文档中
或使用截图工具保存关键信息

3. 查看具体数值要求

某些模型有特定的输入要求：

均值/标准差归一化参数
像素值范围（[0,1] 或 [0,255]）
通道顺序（RGB 或 BGR）

四、查看onnx模型，分析出输入，输出。编写代码进行推理（C++代码+opencv）。

将之前通过Netron查看到的模型信息，转化为实际的C++推理代码。整个过程会分为三个清晰的步骤：环境准备 -> 编写代码 -> 运行与验证。

第一步：准备你的开发环境

在编写代码之前，需要先安装好两个关键的库：

OpenCV：用于读取和预处理图像（如调整大小、归一化）。
ONNX Runtime：用于加载ONNX模型并执行高速推理。

你需要下载对应操作系统的C++库，并在项目中配置好头文件和库文件的路径。如果你使用CMake构建项目，可以参考类似下面的配置：

cmake

复制代码

# CMakeLists.txt 示例片段
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})

# 假设 ONNX Runtime 已解压到指定路径
include_directories(/path/to/onnxruntime/include)
link_directories(/path/to/onnxruntime/lib)

add_executable(your_project main.cpp)
target_link_libraries(your_project ${OpenCV_LIBS} onnxruntime)

请务必将/path/to/onnxruntime替换为你电脑上的实际路径。

第二步：编写C++推理代码

以下是一个通用的C++代码模板，它会自动读取你提供的ONNX模型，获取其输入输出信息，并用OpenCV处理一张图片进行推理。

cpp

复制代码

#include <iostream>
#include <vector>
#include <opencv2/opencv.hpp>
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
#include <onnxruntime/core/providers/cpu/cpu_provider_factory.h> // 如果只用CPU

int main(int argc, char* argv[]) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <onnx_model_path> <image_path>" << std::endl;
        return -1;
    }

    std::string model_path = argv[1];
    std::string image_path = argv[2];

    // 1. 初始化ONNX Runtime环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXInference");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    // 如果你有NVIDIA GPU并想使用CUDA，可以在这里添加CUDA执行提供程序
    // OrtSessionOptionsAppendExecutionProvider_CUDA(session_options);

    // 2. 加载ONNX模型并创建会话
    Ort::Session session(env, model_path.c_str(), session_options);
    Ort::AllocatorWithDefaultOptions allocator;

    // 3. 获取模型输入信息（名称、形状、类型）
    std::cout << "--- Model Inputs ---" << std::endl;
    size_t num_input_nodes = session.GetInputCount();
    std::vector<const char*> input_names;
    std::vector<int64_t> input_node_dims; // 用于存储动态或固定的维度
    for (size_t i = 0; i < num_input_nodes; i++) {
        // 获取输入名称
        char* input_name = session.GetInputName(i, allocator);
        input_names.push_back(input_name);
        std::cout << "Input [" << i << "] Name: " << input_name << std::endl;

        // 获取输入类型信息
        Ort::TypeInfo type_info = session.GetInputTypeInfo(i);
        auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
        ONNXTensorElementDataType type = tensor_info.GetElementType();
        std::cout << "   Type: " << type << std::endl;

        // 获取输入形状
        input_node_dims = tensor_info.GetShape();
        std::cout << "   Shape: ";
        for (long dim : input_node_dims) std::cout << dim << " ";
        std::cout << std::endl;

        // 注意：如果维度是-1，表示它是动态的（如batch size）
    }

    // 4. 获取模型输出信息
    std::cout << "--- Model Outputs ---" << std::endl;
    size_t num_output_nodes = session.GetOutputCount();
    std::vector<const char*> output_names;
    for (size_t i = 0; i < num_output_nodes; i++) {
        char* output_name = session.GetOutputName(i, allocator);
        output_names.push_back(output_name);
        std::cout << "Output [" << i << "] Name: " << output_name << std::endl;

        Ort::TypeInfo type_info = session.GetOutputTypeInfo(i);
        auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
        std::vector<int64_t> output_dims = tensor_info.GetShape();
        std::cout << "   Shape: ";
        for (long dim : output_dims) std::cout << dim << " ";
        std::cout << std::endl;
    }

    // 5. 使用OpenCV读取和预处理图像
    cv::Mat image = cv::imread(image_path);
    if (image.empty()) {
        std::cerr << "Error: Could not read image: " << image_path << std::endl;
        return -1;
    }

    // 假设模型需要固定尺寸的输入（例如224x224），从输入信息中获取目标尺寸
    // 注意：这里简化处理，假设输入形状是 [batch, channels, height, width]
    int64_t input_h = input_node_dims[2]; // 假设为NCHW格式
    int64_t input_w = input_node_dims[3];
    cv::Mat resized_image;
    cv::resize(image, resized_image, cv::Size(input_w, input_h));

    // 创建blob: 将HWC格式转换为CHW格式，并进行归一化（1.0/255），像素值从 [0, 255] 被归一化 到 [0, 1] 的浮点数范围。
    cv::Mat blob = cv::dnn::blobFromImage(resized_image, 1.0/255.0, cv::Size(input_w, input_h), cv::Scalar(), true, false);

    // 6. 准备ONNX Runtime的输入张量
    size_t input_tensor_size = blob.total(); // blob已经是连续的
    std::vector<float> input_tensor_values(blob.begin<float>(), blob.end<float>());

    // 创建输入张量
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_tensor_values.data(), input_tensor_size,
        input_node_dims.data(), input_node_dims.size());

    // 7. 运行推理
    try {
        auto output_tensors = session.Run(Ort::RunOptions{nullptr},
                                          input_names.data(), &input_tensor, 1,
                                          output_names.data(), output_names.size());

        // 8. 处理输出结果
        std::cout << "--- Inference Results ---" << std::endl;
        for (size_t i = 0; i < output_tensors.size(); ++i) {
            // 获取输出数据指针
            float* output_data = output_tensors[i].GetTensorMutableData<float>();
            // 获取输出形状信息
            auto output_info = output_tensors[i].GetTensorTypeAndShapeInfo();
            auto output_shape = output_info.GetShape();
            size_t output_count = output_info.GetElementCount();

            std::cout << "Output " << i << " (first 10 elements): ";
            for (size_t j = 0; j < std::min((size_t)10, output_count); ++j) {
                std::cout << output_data[j] << " ";
            }
            std::cout << "..." << std::endl;
        }
    } catch (const Ort::Exception& e) {
        std::cerr << "ONNX Runtime error: " << e.what() << std::endl;
        return -1;
    }

    // 清理动态分配的名称
    for (auto name : input_names) allocator.Free((void*)name);
    for (auto name : output_names) allocator.Free((void*)name);

    return 0;
}

第三步：编译与运行

将上述代码保存为main.cpp，配置好CMakeLists.txt后，按以下步骤操作：

编译：在你的环境中执行编译命令（如cmake和make，或在Visual Studio中生成解决方案）。
准备模型和图片：确保你的ONNX模型文件和待推理的图片路径正确。
运行：在终端执行生成的可执行文件。

bash
复制代码
```
./your_project path/to/your_model.onnx path/to/your_image.jpg
```
程序会首先打印出从模型中读取的输入输出信息 ，然后显示推理结果的前10个数值。如果是一个分类模型，你可能会看到代表各个类别概率的数值数组。

代码要点解析

动态获取信息 ：代码核心在于session.GetInputCount()、GetInputName()和GetInputTypeInfo()等API。这让你无需硬编码，即可处理任何ONNX模型。
图像预处理 ：cv::dnn::blobFromImage是OpenCV中一个非常强大的函数，它将图像转换为模型要求的格式（NCHW，即数量、通道、高、宽），并同时完成缩放和归一化。
形状处理 ：模型输入的维度可能是动态的（如用-1表示可变批处理大小）。在实际使用时，你需要根据当前输入（如单张图片）将其替换为具体的数值（例如将-1替换为1）。

五、查看onnx模型，分析出输入，输出。编写代码进行推理（python代码）。

根据Netron分析的信息，这里提供几个不同场景的代码模板：

模板1：通用分类模型（ImageNet风格）

python

复制代码

import numpy as np
import cv2
import onnxruntime as ort
import json

class ONNXClassifier:
    def __init__(self, model_path, class_names_path=None):
        """
        初始化分类器
        
        Args:
            model_path: ONNX模型路径
            class_names_path: 类别名称JSON文件路径（可选）
        """
        # 加载ONNX模型
        self.session = ort.InferenceSession(model_path)
        
        # 获取模型输入输出信息
        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        self.output_name = self.session.get_outputs()[0].name
        self.output_shape = self.session.get_outputs()[0].shape
        
        print("=== 模型信息 ===")
        print(f"输入名称: {self.input_name}")
        print(f"输入形状: {self.input_shape}")
        print(f"输入类型: {self.session.get_inputs()[0].type}")
        print(f"输出名称: {self.output_name}")
        print(f"输出形状: {self.output_shape}")
        print(f"输出类型: {self.session.get_outputs()[0].type}")
        
        # 加载类别名称
        self.class_names = None
        if class_names_path:
            with open(class_names_path, 'r') as f:
                self.class_names = json.load(f)
    
    def preprocess(self, image_path):
        """
        预处理图像（适用于ImageNet风格模型）
        假设输入格式: [batch, channels, height, width], RGB, [0,1]范围
        """
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # BGR转RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 获取目标尺寸（假设NCHW格式）
        if len(self.input_shape) == 4:
            target_h, target_w = self.input_shape[2], self.input_shape[3]
        else:
            target_h, target_w = 224, 224  # 默认值
        
        # 调整尺寸
        resized = cv2.resize(image, (target_w, target_h))
        
        # 转换为blob: HWC -> CHW, 归一化到[0,1]
        blob = cv2.dnn.blobFromImage(
            resized,
            scale=1.0/255.0,  # [0,255] -> [0,1]
            size=(target_w, target_h),
            mean=(0, 0, 0),
            swapRB=False,      # 已经手动转换为RGB
            crop=False
        )
        
        # ImageNet标准化: (x - mean) / std
        mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 3, 1, 1)
        std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 3, 1, 1)
        blob = (blob - mean) / std
        
        return blob
    
    def predict(self, image_path, top_k=5):
        """
        预测图像类别
        
        Args:
            image_path: 图像路径
            top_k: 返回前K个预测结果
        
        Returns:
            预测结果列表
        """
        # 预处理
        input_tensor = self.preprocess(image_path)
        
        # 推理
        outputs = self.session.run(
            [self.output_name],
            {self.input_name: input_tensor}
        )
        
        # 获取预测结果
        predictions = outputs[0][0]  # 去掉batch维度
        
        # 获取top-k索引
        top_indices = np.argsort(predictions)[-top_k:][::-1]
        
        # 构建结果
        results = []
        for idx in top_indices:
            score = predictions[idx]
            class_name = self.class_names[idx] if self.class_names else f"class_{idx}"
            results.append({
                'class_id': int(idx),
                'class_name': class_name,
                'score': float(score)
            })
        
        return results

# 使用示例
if __name__ == "__main__":
    # 初始化分类器
    classifier = ONNXClassifier(
        model_path="resnet50.onnx",
        class_names_path="imagenet_classes.json"  # 可选
    )
    
    # 预测
    results = classifier.predict("test_image.jpg", top_k=5)
    
    # 打印结果
    print("\n=== 预测结果 ===")
    for i, result in enumerate(results):
        print(f"{i+1}. {result['class_name']}: {result['score']:.4f}")

模板2：目标检测模型（YOLO风格）

python

复制代码

import numpy as np
import cv2
import onnxruntime as ort

class ONNXDetector:
    def __init__(self, model_path, conf_threshold=0.5, nms_threshold=0.45):
        """
        初始化检测器
        
        Args:
            model_path: ONNX模型路径
            conf_threshold: 置信度阈值
            nms_threshold: NMS阈值
        """
        self.session = ort.InferenceSession(model_path)
        self.conf_threshold = conf_threshold
        self.nms_threshold = nms_threshold
        
        # 获取输入信息
        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        
        print("=== 模型信息 ===")
        print(f"输入名称: {self.input_name}")
        print(f"输入形状: {self.input_shape}")
        
    def preprocess(self, image):
        """
        预处理图像（适用于YOLO风格模型）
        假设输入: [batch, channels, height, width], RGB, [0,1]范围
        """
        original_height, original_width = image.shape[:2]
        
        # 获取目标尺寸
        target_h, target_w = self.input_shape[2], self.input_shape[3]
        
        # 调整尺寸
        resized = cv2.resize(image, (target_w, target_h))
        
        # 归一化到[0,1]
        resized = resized.astype(np.float32) / 255.0
        
        # HWC -> CHW
        blob = resized.transpose(2, 0, 1)
        
        # 添加batch维度
        blob = np.expand_dims(blob, axis=0)
        
        return blob, original_height, original_width
    
    def postprocess(self, outputs, original_height, original_width):
        """
        后处理检测结果
        """
        # 这里需要根据具体的输出格式来解析
        # YOLOv8通常输出: [1, 84, 8400] 或 [1, 8400, 84]
        
        output = outputs[0]
        
        # 假设输出格式: [1, 84, 8400]
        if output.shape[1] == 84:
            output = output[0]  # [84, 8400]
            output = output.T    # [8400, 84]
        
        # 解析边界框和分数
        boxes = []
        scores = []
        class_ids = []
        
        for detection in output:
            # 提取类别分数
            class_scores = detection[4:]
            class_id = np.argmax(class_scores)
            confidence = class_scores[class_id]
            
            if confidence > self.conf_threshold:
                # 提取边界框坐标 (cx, cy, w, h)
                cx, cy, w, h = detection[:4]
                
                # 转换为xyxy格式
                x1 = (cx - w/2) * original_width
                y1 = (cy - h/2) * original_height
                x2 = (cx + w/2) * original_width
                y2 = (cy + h/2) * original_height
                
                boxes.append([x1, y1, x2, y2])
                scores.append(float(confidence))
                class_ids.append(int(class_id))
        
        # 应用NMS
        indices = cv2.dnn.NMSBoxes(
            boxes, scores, self.conf_threshold, self.nms_threshold
        )
        
        # 提取最终结果
        results = []
        if len(indices) > 0:
            for i in indices.flatten():
                results.append({
                    'bbox': boxes[i],
                    'score': scores[i],
                    'class_id': class_ids[i]
                })
        
        return results
    
    def detect(self, image_path):
        """
        执行目标检测
        
        Args:
            image_path: 图像路径
        
        Returns:
            检测结果列表
        """
        # 读取图像
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法读取图像: {image_path}")
        
        # BGR转RGB
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 预处理
        input_tensor, orig_h, orig_w = self.preprocess(image_rgb)
        
        # 推理
        outputs = self.session.run(None, {self.input_name: input_tensor})
        
        # 后处理
        detections = self.postprocess(outputs, orig_h, orig_w)
        
        return detections, image

# 使用示例
if __name__ == "__main__":
    detector = ONNXDetector(
        model_path="yolov8n.onnx",
        conf_threshold=0.5,
        nms_threshold=0.45
    )
    
    detections, image = detector.detect("test_image.jpg")
    
    print(f"\n检测到 {len(detections)} 个目标:")
    for det in detections:
        print(f"类别 {det['class_id']}: {det['score']:.3f}, "
              f"框: [{det['bbox'][0]:.0f}, {det['bbox'][1]:.0f}, "
              f"{det['bbox'][2]:.0f}, {det['bbox'][3]:.0f}]")

模板3：通用推理模板（适用于任何模型）

python

复制代码

import numpy as np
import onnxruntime as ort
import cv2
from typing import Dict, Any

class ONNXInference:
    def __init__(self, model_path):
        """
        通用的ONNX推理类
        """
        # 加载模型
        self.session = ort.InferenceSession(model_path)
        
        # 获取输入输出信息
        self.inputs = self.session.get_inputs()
        self.outputs = self.session.get_outputs()
        
        print("=== 模型输入信息 ===")
        for i, input_info in enumerate(self.inputs):
            print(f"Input {i}:")
            print(f"  名称: {input_info.name}")
            print(f"  形状: {input_info.shape}")
            print(f"  类型: {input_info.type}")
        
        print("\n=== 模型输出信息 ===")
        for i, output_info in enumerate(self.outputs):
            print(f"Output {i}:")
            print(f"  名称: {output_info.name}")
            print(f"  形状: {output_info.shape}")
            print(f"  类型: {output_info.type}")
    
    def get_input_details(self) -> Dict[str, Any]:
        """
        获取输入详细信息
        """
        details = {}
        for input_info in self.inputs:
            details[input_info.name] = {
                'shape': input_info.shape,
                'type': input_info.type
            }
        return details
    
    def get_output_details(self) -> Dict[str, Any]:
        """
        获取输出详细信息
        """
        details = {}
        for output_info in self.outputs:
            details[output_info.name] = {
                'shape': output_info.shape,
                'type': output_info.type
            }
        return details
    
    def infer(self, input_data: Dict[str, np.ndarray]):
        """
        执行推理
        
        Args:
            input_data: 字典，键为输入名称，值为numpy数组
        
        Returns:
            输出结果列表
        """
        # 验证输入
        for name, data in input_data.items():
            if name not in [inp.name for inp in self.inputs]:
                print(f"警告: 输入 '{name}' 不在模型输入中")
        
        # 运行推理
        outputs = self.session.run(
            [out.name for out in self.outputs],
            input_data
        )
        
        return outputs

# 使用示例
if __name__ == "__main__":
    # 初始化
    inference = ONNXInference("model.onnx")
    
    # 获取输入输出信息
    input_details = inference.get_input_details()
    output_details = inference.get_output_details()
    
    # 准备输入数据（这里需要根据实际模型调整）
    # 例如，如果模型需要图像输入
    if input_details:
        input_name = list(input_details.keys())[0]
        input_shape = input_details[input_name]['shape']
        
        # 创建随机输入（仅用于测试）
        random_input = np.random.randn(*input_shape).astype(np.float32)
        
        # 推理
        results = inference.infer({input_name: random_input})
        
        print(f"\n推理结果形状: {[r.shape for r in results]}")

快速测试脚本

python

复制代码

import onnxruntime as ort
import numpy as np
import cv2
import json

def quick_test(model_path, image_path=None):
    """
    快速测试ONNX模型
    """
    # 加载模型
    session = ort.InferenceSession(model_path)
    
    # 获取输入信息
    input_info = session.get_inputs()[0]
    input_name = input_info.name
    input_shape = input_info.shape
    
    # 创建测试输入
    if image_path:
        # 从图像创建输入
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError("无法读取图像")
        
        # 简单预处理
        if len(input_shape) == 4:
            target_h, target_w = input_shape[2], input_shape[3]
            image = cv2.resize(image, (target_w, target_h))
            image = image.astype(np.float32) / 255.0
            input_data = image.transpose(2, 0, 1)[np.newaxis, ...]
        else:
            input_data = image.astype(np.float32)[np.newaxis, ...]
    else:
        # 创建随机输入
        input_data = np.random.randn(*input_shape).astype(np.float32)
    
    # 推理
    outputs = session.run(None, {input_name: input_data})
    
    # 显示结果
    print("推理完成!")
    for i, output in enumerate(outputs):
        print(f"输出{i}形状: {output.shape}")
        print(f"输出{i}范围: [{output.min():.4f}, {output.max():.4f}]")
        if output.size < 100:
            print(f"输出{i}值: {output.flatten()[:10]}...")
    
    return outputs

# 运行测试
if __name__ == "__main__":
    results = quick_test("model.onnx", "test.jpg")