如何在C++(VS)中用DirectML加载图像作为ONNX模型输入并获取输出

一、现有代码的核心问题 #

你的代码完成了模型加载和图像预处理的基础工作，但缺少输入维度的正确适配和模型推理、输出解析的关键步骤，同时图像预处理环节没有强制将图像缩放到ResNet18要求的224x224标准输入尺寸。

二、完整修正代码 #

#include <iostream>
#include <vector>
#include <string>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "stb_image_resize.h"
#include <onnxruntime_cxx_api.h>
#include <dml_provider_factory.h>

// 加载并预处理图像：resize到224x224、归一化、转为CHW格式
std::vector<float> load_image_and_preprocess(const std::string& filename) 
{
    const int target_width = 224;
    const int target_height = 224;

    int orig_w, orig_h, orig_channels;
    // 加载原始图像
    unsigned char* orig_data = stbi_load(filename.c_str(), &orig_w, &orig_h, &orig_channels, 3);
    if (!orig_data) {
        std::cerr << "Failed to load image: " << filename << std::endl;
        return {};
    }

    // 缩放到目标尺寸
    unsigned char* resized_data = new unsigned char[target_width * target_height * 3];
    stbir_resize_uint8(orig_data, orig_w, orig_h, 0,
                       resized_data, target_width, target_height, 0,
                       3);

    // 预处理：归一化+转为CHW格式
    std::vector<float> input_data(target_width * target_height * 3);
    float* r_ptr = &input_data[0];
    float* g_ptr = &input_data[target_width * target_height];
    float* b_ptr = &input_data[2 * target_width * target_height];

    for (int i = 0; i < target_width * target_height; ++i) {
        *r_ptr++ = (static_cast<float>(resized_data[3*i + 0]) / 255.0f - 0.485f) / 0.229f;
        *g_ptr++ = (static_cast<float>(resized_data[3*i + 1]) / 255.0f - 0.456f) / 0.224f;
        *b_ptr++ = (static_cast<float>(resized_data[3*i + 2]) / 255.0f - 0.406f) / 0.225f;
    }

    // 释放内存
    stbi_image_free(orig_data);
    delete[] resized_data;

    return input_data;
}

int main()
{
    // 模型路径
    const std::string model_path = "C:/Downloads/resnet18-v1-7.tar/resnet18-v1-7/resnet18-v1-7.onnx";
    std::basic_string<ORTCHAR_T> model_ort_path(model_path.begin(), model_path.end());

    // 初始化ONNX Runtime环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ResNet18-Inference");
    Ort::SessionOptions session_options;
    // 添加DirectML执行提供者（不需要可注释，改用CPU）
    OrtSessionOptionsAppendExecutionProvider_DML(session_options, 0);
    Ort::Session session(env, model_ort_path.c_str(), session_options);

    // 获取输入输出信息
    Ort::AllocatorWithDefaultOptions allocator;
    auto input_name = session.GetInputNameAllocated(0, allocator);
    auto output_name = session.GetOutputNameAllocated(0, allocator);
    std::cout << "Input name: " << input_name.get() << std::endl;
    std::cout << "Output name: " << output_name.get() << std::endl;

    // 定义输入Tensor的shape：[batch, channel, height, width]，batch设为1
    std::vector<int64_t> input_shape = {1, 3, 224, 224};
    const size_t input_tensor_size = 1 * 3 * 224 * 224;

    // 加载并预处理图像
    std::string image_path = "C:/Desktop/Image.png";
    std::vector<float> input_data = load_image_and_preprocess(image_path);
    if (input_data.empty()) {
        return -1;
    }

    // 创建输入Tensor
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_data.data(), input_tensor_size,
        input_shape.data(), input_shape.size()
    );

    // 准备输入输出名称数组
    const char* input_names[] = {input_name.get()};
    const char* output_names[] = {output_name.get()};

    // 执行模型推理
    Ort::Value output_tensor = session.Run(
        Ort::RunOptions{nullptr}, input_names, &input_tensor, 1,
        output_names, 1
    );

    // 解析输出结果
    float* output_data = output_tensor.GetTensorMutableData<float>();
    auto output_shape = output_tensor.GetTensorTypeAndShapeInfo().GetShape();
    size_t output_size = 1;
    for (auto dim : output_shape) {
        output_size *= dim;
    }

    // 找到置信度最高的类别（对应ImageNet 1000分类）
    float max_confidence = -1.0f;
    int max_class_idx = -1;
    for (size_t i = 0; i < output_size; ++i) {
        if (output_data[i] > max_confidence) {
            max_confidence = output_data[i];
            max_class_idx = static_cast<int>(i);
        }
    }

    std::cout << "Predicted class index: " << max_class_idx << std::endl;
    std::cout << "Confidence: " << max_confidence << std::endl;

    return 0;
}

三、关键步骤解析 #

1. 图像预处理优化

   • 新增图像resize：用stb_image_resize将任意尺寸图像缩放到224x224，匹配ResNet输入要求
   • CHW格式转换：将R、G、B通道分别存入连续内存块，符合ONNX模型通道优先的输入规范
   • 标准化：严格使用ResNet预训练模型的归一化参数(像素值/255 - mean)/std，保证推理精度

2. 输入Tensor构建

   • 明确输入shape：指定为{1, 3, 224, 224}，第一个维度是batch size，单张图像推理设为1
   • 内存匹配：确保输入数据的内存布局与Tensor要求一致，避免维度不匹配导致的推理失败

3. 模型推理与输出解析

   • 执行推理：调用session.Run，传入输入名称、输入Tensor、输出名称，获取输出Tensor
   • 提取结果：通过GetTensorMutableData<float>()拿到分类置信度数组
   • 解析最大值：遍历置信度数组，找到对应最高置信度的类别索引

四、注意事项 #

• 确保项目已链接ONNX Runtime库，同时添加stb_image和stb_image_resize的头文件
• 若不需要DirectML加速，可注释OrtSessionOptionsAppendExecutionProvider_DML，改用CPU执行
• 若模型输入尺寸不是224x224，需对应调整预处理和输入shape参数

内容的提问来源于stack exchange，提问作者Ron