Libtorch初试教程《三》

读取图片进行分类

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>

using namespace std;
using namespace cv;

int main(int argc, const char* argv[]) {
    if (argc != 3) {
        std::cerr << "usage: example-app <path-to-exported-script-module>\n";
        return -1;
    }

    torch::jit::script::Module module;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        module = torch::jit::load(argv[1]);
        module.to(at::kCUDA);
    }
    catch (const c10::Error &e) {
        std::cout << "Error loading model: " << e.msg() << std::endl;
        std::cerr << "error loading the module model\n";

        return -1;
    }
    //输入图像
    auto src_image_ = imread(argv[2]);
    if (src_image_.empty())
        fprintf(stderr, "Can not load image\n");
    cv::Mat image_transfomed;
    cv::resize(src_image_, image_transfomed, cv::Size(224, 224));
    cv::cvtColor(image_transfomed, image_transfomed, cv::COLOR_BGR2RGB);

    // 图像转换为Tensor
    torch::Tensor tensor_image = torch::from_blob(image_transfomed.data, {1, image_transfomed.rows, image_transfomed.cols,3},torch::kByte);
    tensor_image = tensor_image.permute({0,3,1,2});
    tensor_image = tensor_image.toType(torch::kFloat);
    tensor_image = tensor_image.div(255);
    tensor_image = (tensor_image-0.5).div(0.5);

    tensor_image = tensor_image.to(torch::kCUDA);
    torch::Tensor result = module.forward({tensor_image}).toTensor();

    std::cout << result.slice(/*dim=*/1, /*start=*/0, /*end=*/10) << '\n';
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.15)
project(libtorch_opencv_demo1)

set(CMAKE_CXX_STANDARD 14)
set(OpenCV_DIR /home/nvidia/opencv-3.0.0/build)
list(APPEND CMAKE_PREFIX_PATH "abs_path/libtorch-shared-with-deps-1.3.1/libtorch")

find_package(Torch REQUIRED)
find_package(OpenCV REQUIRED)
include_directories( ${OpenCV_INCLUDE_DIRS} )

add_executable(libtorch_opencv_demo1 main.cpp)
target_link_libraries(libtorch_opencv_demo1 "${TORCH_LIBRARIES}" "${OpenCV_LIBS}")
set_property(TARGET libtorch_opencv_demo1 PROPERTY CXX_STANDARD 11)