BiRefNet512*512onnx模型c++推理

BiRefNet主页：github图像前景分割

onnx模型下载：githu releases

前处理使用opencv,后处理使用libtorch

参考：onnxruntime c++ inference

           U_NET 网络onnx部署

           libtorch 输出tensor 到文件并用ImageJ 查看

环境：win11-wsl

依赖opencv,onnxruntime,libtorch

opencv编译参考：rk3588使用交叉编译opencv静态库，后面有wsl报openjp2错误解决方法

onnxruntime下载编译好的包解压：onnxruntime releases

libtorch下载编译好的包解压，参考：torch官方教程

目录结构：onnxruntime不一定要放在这里

main.cpp:

#include <vector>
#include <onnxruntime_cxx_api.h>
#include <opencv2/opencv.hpp>
#include <torch/torch.h>
#include <torch/script.h>
#include <fstream>

void tensor2txt(torch::Tensor &tsor, std::string path)
{
    std::ofstream outfile;
    outfile.open(path);
    assert(outfile && "failed to open the file!");    
    // tsor = tsor.reshape({-1});
    tsor.squeeze(0);
    for (int i=0; i<tsor.sizes()[0]; i++) {
      for (int j = 0; j < tsor.sizes()[1]; j++) {
        auto s = tsor[i][j].item<float>();
        outfile << s << ","; 
      }
        outfile << "\n"; 
    }
    outfile.close();
}

int main(int argc, char* argv[])
{
    printf("main start!!!");
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "test");
    Ort::SessionOptions session_options;
    session_options.SetInterOpNumThreads(1);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
    const char* model_path = "./models/BiRefNet-general-resolution_512x512-fp16-epoch_216.onnx";
    
    Ort::Session session(env, model_path, session_options);
    Ort::AllocatorWithDefaultOptions allocator;
    
    size_t num_input_nodes = session.GetInputCount();
    Ort::Value input_tensor{nullptr};
    std::array<int64_t, 4> input_shape{1, 3, 512, 512};
    
    Ort::Value output_tensor{nullptr};
    std::array<int64_t, 4> output_shape{1, 1, 512, 512};
    
    std::array<float, 512*512*3> input_image{};
    std::array<float, 512*512*1> output{};
    
    std::string imgPath = "./models/20250321111449.jpg";
    cv::Mat img = cv::imread(imgPath);
    cv::cvtColor(img, img, cv::COLOR_BGR2RGB);
    cv::resize(img, img, cv::Size(512, 512));
    cv::Mat img_float;
    img.convertTo(img_float, CV_32FC3, 1/255.0, 0);    
    std::cout<<"dim:"<< img_float.rows<<":"<<img_float.cols<<":"<<img_float.channels()<<std::endl;
    int index = 0;
    for (unsigned c = 0; c < 3; c++)
    {
        for (unsigned i = 0; i < img_float.rows; i++)
        {
            for (unsigned j = 0; j < img_float.cols; j++)
            {
                input_image[index++] = (static_cast<float>(img_float.at<cv::Vec3f>(i, j)[c]));
            }
        }
    }

    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeDefault);
    input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_image.data(), input_image.size(),input_shape.data(), input_shape.size());
    output_tensor = Ort::Value::CreateTensor<float>(memory_info, output.data(), output.size(), output_shape.data(), output_shape.size());
    const char* input_names[] = {"input_image"};
    const char* output_names[] = {"output_image"};

    session.Run(Ort::RunOptions{nullptr}, input_names, &input_tensor, 1, output_names, &output_tensor, 1);

    float* out = output_tensor.GetTensorMutableData<float>();
    torch::Tensor tensor = torch::from_blob((void*)out, {1, 1, 512, 512});
    torch::Tensor output2 = torch::sigmoid(tensor);
    output2.print();
    // 1*C*H*W -> H*W*C
    output2 = output2.squeeze(0).detach().permute({1, 2, 0});
    output2.print();
    tensor2txt(output2, "./output1.txt");
    output2 = output2.mul(255).clamp(0, 255).to(torch::kU8);
    output2.print();
    tensor2txt(output2, "./output2.txt");
    output2 = output2.cpu();
    cv::Mat resultImg(512, 512, CV_8UC1, output2.data_ptr());
    cv::imwrite("result.jpg", resultImg);

    // for(int i = 0; i < 10; i++)
    // {
    //     std::cout<< out[i]<<":";
    // }

    return 0;
}

CMakeLists.txt:

cmake_minimum_required(VERSION 3.18)

project(onnx_example LANGUAGES CXX)

find_package(Torch REQUIRED)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

# 替换成自己的路径
set(ONNXRUNTIME_DIR "xxxx/3rdparty/onnxruntime-linux-x64-1.21.0")
set(OpenCV_DIR "xxxx/opencv/install/x64/lib/cmake/opencv4")
include_directories(${ONNXRUNTIME_DIR}/include)
find_package(Iconv)
find_package(OpenCV REQUIRED )

include_directories(
    ${OpenCV_INCLUDE_DIRS}/include
    )


add_executable(onnx_example main.cpp)
target_link_libraries(onnx_example PRIVATE ${ONNXRUNTIME_DIR}/lib/libonnxruntime.so ${OpenCV_LIBS} ${TORCH_LIBRARIES})

install(TARGETS ${PROJECT_NAME} DESTINATION .)
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../models/20250321111449.jpg DESTINATION models)
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../models/BiRefNet-general-resolution_512x512-fp16-epoch_216.onnx DESTINATION models)

进入build编译：替换成自己的libtorch路径

cmake ../ -DCMAKE_INSTALL_PREFIX=../lib -DCMAKE_PREFIX_PATH=xxx/depends/libtorch

make

make install

可执行文件在lib路径下

添加libtorch和onnxruntime的环境变量

export ONNXRUNTIME_DIR="xxx/depends/onnxruntime-linux-x64-1.21.0"

export LIBTORCH_DIR="xxx/depends/libtorch"
export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH

export LD_LIBRARY_PATH=$LIBTORCH_DIR/lib:$LD_LIBRARY_PATH

运行：./onnx_example