读取图片进行分类
#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)