AI本地化部署:C++运行图像超分

AI本地化部署:C++运行超分增强图像

前言

文章中包含部分Python部分,你只要会安装Python环境即可,如果不想要了解Python,可以直接进入到编写C++代码部分,我已经提供了一个示例工程

环境准备

  1. 安装onnxruntime和opencv:
pip install onnxruntime opencv-python
  1. 安装百度飞浆环境,我这里是一CPU推理为示例(详细可以参考官网安装:https://www.paddlepaddle.org.cn/install/quick ):
python -m pip install paddlepaddle==2.6.2 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/
  1. 按照官方教程安装准备PPGAN, 连接如下:
https://aistudio.baidu.com/projectdetail/3205183

模型转换

  • 新建一个Python脚本,运行下面代码,导出模型为onnx:
from ppgan.models.generators import RRDBNet
import paddle
import paddle.static as static

def export_onnx():
	# 加载百度飞浆的模型
    model = RRDBNet(3, 3, 64, 23)
    # 下载后的模型路径,注意按照实际更改
    state_dict = paddle.load("weight_params/DF2K_JPEG.pdparams")
    model.load_dict(state_dict)

    # 为模型指定输入的形状和数据类型,支持持 Tensor 或 InputSpec ,InputSpec 支持动态的 shape。
    x_spec = paddle.static.InputSpec([None, 3, None, None], 'float32', 'x') 
    # 指定导出的onnx的保存路径
    model_path = "./onnx_net/rrdbnet"
    # 导出onnx模型
    paddle.onnx.export(model, model_path, input_spec=[x_spec], opset_version=11)
 if __name__ == "__main__":
    export_onnx()
  • 执行完成后,你应该会指定的路径下面看到一个ONNX后缀的文件,我们可以通过下面代码验证一下:
import onnxruntime
def test():
    ort_sess = onnxruntime.InferenceSession("onnx_net/rrdbnet.onnx")
    in_img = Image.open("./test.png").convert("RGB")
    x = norm(in_img)
    x = np.expand_dims(x,axis=0).astype('float32')
    print(x.shape)
    print(ort_sess.get_inputs()[0])
    ort_outs  = ort_sess.run(None,{
         "x":x
    })
    y = ort_outs[0]
    out_img = Image.fromarray(denorm(y[0]))
    out_img.save("./test_out.png")

使用ONNX模型

C++项目准备
  1. 新建一个Visual Studio空白控制台工程:
    在这里插入图片描述

  2. 右键项目管理器的工程->管理NuGet程序包->搜索OnnxRuntime并下载
    在这里插入图片描述

  3. 然后搜索OpenCV并下载
    在这里插入图片描述

  4. 现在环境已经准备就绪了。

编写C++代码

示例工程可以直接从这里下载:

https://download.youkuaiyun.com/download/weixin_45809578/90364889
  1. 新建一个C++文件,粘贴如下代码:
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>

void PrintTensorContent(Ort::Value& tensor)
{
    // 获取张量类型和形状信息
    auto type_and_shape = tensor.GetTensorTypeAndShapeInfo();
    auto shape = type_and_shape.GetShape();
    size_t total_size = type_and_shape.GetElementCount();

    // 打印张量形状
    std::cout << "Tensor Shape: [";
    for (size_t i = 0; i < shape.size(); ++i) {
        std::cout << shape[i];
        if (i < shape.size() - 1) {
            std::cout << ", ";
        }
    }
    std::cout << "]\n";

    // 获取张量数据指针并打印内容
    const float* data = tensor.GetTensorMutableData<float>();
    std::cout << "Tensor Content:\n";
    for (size_t i = 0; i < total_size; ++i) {
        std::cout << data[i] << " ";
        if ((i + 1) % shape.back() == 0) {  // 每行换行
            std::cout << "\n";
        }
        if (i > 64) {
            std::cout << "..." << std::endl;
            break;
        }
    }
    std::cout << std::endl;
}
std::vector<cv::Mat> TensorToMatBatch(Ort::Value& tensor)
{
    // 获取张量类型和形状信息
    auto type_and_shape = tensor.GetTensorTypeAndShapeInfo();
    auto shape = type_and_shape.GetShape();

    // 验证张量维度是否为 4D (NCHW)
    if (shape.size() != 4) {
        throw std::runtime_error("Unsupported tensor shape for conversion to cv::Mat. Expected 4D (N, C, H, W).");
    }

    size_t batch_size = shape[0];  // N
    size_t channels = shape[1];    // C
    size_t height = shape[2];      // H
    size_t width = shape[3];       // W

    // 获取张量数据指针
    const float* tensor_data = tensor.GetTensorMutableData<float>();

    std::vector<cv::Mat> result_images;

    // 遍历每个图像批次
    for (size_t n = 0; n < batch_size; n++) {
        const float* batch_data = tensor_data + n * channels * height * width;

        // 将批次数据转换为 cv::Mat
        cv::Mat chw_mat(channels, height * width, CV_32F, const_cast<float*>(batch_data));

        // 转换为 HWC 格式的 cv::Mat
        std::vector<cv::Mat> channel_mats;
        for (size_t c = 0; c < channels; ++c) {
            cv::Mat channel(height, width, CV_32F, chw_mat.ptr<float>(c));
            channel_mats.push_back(channel);
        }

        cv::Mat hwc_mat;
        cv::merge(channel_mats, hwc_mat);
        result_images.push_back(hwc_mat);
    }

    return result_images;
}

std::vector<cv::Mat> test(const std::vector<std::string>& image_paths,const ORTCHAR_T* model_path)
{
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "ONNXRuntimeExample");

    // Create session options
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(4);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

    Ort::AllocatorWithDefaultOptions allocator;
    //加载超分模型
    //Ort::Session session(env, L"rrdbnet.onnx", session_options);
    Ort::Session session(env, model_path, session_options);
    std::string input_name = "";
    std::string output_name = "";
    //
    //查询到输入输出的名称
    for (size_t i = 0; i < session.GetInputCount(); i++)
    {
        auto name = session.GetInputNameAllocated(i, allocator);
        input_name = name.get();
        printf("input name: %s\n", input_name.c_str());
    }
    for (size_t i = 0; i < session.GetOutputCount(); i++)
    {
        auto name = session.GetOutputNameAllocated(i, allocator);
        output_name = name.get();
        printf("output name: %s\n", output_name.c_str());
    }

    std::vector<cv::Mat> images;
    std::vector<float> batch_data;
    int batch_size = image_paths.size();

    for (const auto& path : image_paths) {
        cv::Mat img = cv::imread(path, cv::IMREAD_COLOR);
        if (img.empty()) {
            std::cerr << "Failed to read image: " << path << std::endl;
            continue;
        }

        img.convertTo(img, CV_32FC3);
        img = img / 255.0f;
        images.push_back(img);

        // 将图像数据转换为 CHW 格式并追加到 batch_data
        for (int c = 0; c < 3; c++) {
            for (int h = 0; h < img.rows; h++) {
                for (int w = 0; w < img.cols; w++) {
                    batch_data.push_back(img.at<cv::Vec3f>(h, w)[c]);
                }
            }
        }
    }
    
    // 定义张量形状
    std::vector<int64_t> shape = { batch_size, 3, images[0].rows, images[0].cols };
    Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, batch_data.data(), batch_data.size(), shape.data(), shape.size()
        );

    // 推理
    const char* input_names[] = { input_name.c_str() };
    const char* output_names[] = { output_name.c_str() };
    std::chrono::system_clock::time_point start = std::chrono::system_clock::now();
    auto output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_names, &input_tensor, 1, output_names, 1);
    std::chrono::system_clock::time_point end = std::chrono::system_clock::now();
    std::cout << "inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

    // 处理输出
    std::vector<cv::Mat> output_mat_list = TensorToMatBatch(output_tensors[0]);
    return output_mat_list;
}

int main()
{
    std::vector<std::string> image_paths = {
        "D:\\test\\1.jpg",
        "D:\\test\\2.jpg",
        "D:\\test\\3.jpg",
        "D:\\test\\4.jpg"
    };
    // D:\\onnx_net\\rrdbnet.onnx 路径根据实际导出的位置更改
    std::vector<cv::Mat> images = test(image_paths, L"D:\\onnx_net\\rrdbnet.onnx");
    for (int i = 0; i < images.size(); i++)
    {
        auto& output_mat = images[i];
        output_mat *= 255.0f;
        output_mat.convertTo(output_mat, CV_8UC3);
        cv::imshow("output", output_mat);
        cv::waitKey(0);
        std::string output_path = "D:/output_" + std::to_string(i) + ".png";
        cv::imwrite(output_path, output_mat);
    }
    return 0;
}
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值