cv::imread() 默认读取为三通道BGR,需要进行B/R通道交换,这里采用 cv::cvtColor()实现。
图像尺寸需要调整到 224 × 224 224\times 224224×224,通过 cv::resize() 实现。
opencv读取的图像矩阵存储形式:H x W x C, 但是pytorch中 Tensor的存储为:N x C x H x W, 因此需要进行变换,就是np.transpose()操作,这里使用tensor.permut()实现,效果是一样的。
数据归一化,采用 tensor.div(255) 实现。
// test_model.cpp
include
include <torch/torch.h>
include <torch/script.h>
include <opencv2/core.hpp>
include <opencv2/imgproc/imgproc.hpp>
include <opencv2/highgui/highgui.hpp>
int main(int argc, char* argv[]) {
// 加载JIT模型
auto module = torch::jit::load(argv[1]);
// 加载图像
auto image = cv::imread(argv[2], cv::ImreadModes::IMREAD_COLOR);
cv::Mat image_transfomed;
cv::resize(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, {image_transfomed.rows, image_transfomed.cols, 3},torch::kByte);
tensor_image = tensor_image.permute({2, 0, 1});
// tensor_image = tensor_image.toType(torch::kFloat);
tensor_image = tensor_image.div(255.);
// tensor_image = tensor_image.sub(0.5);
// tensor_image = tensor_image.div(0.5);
tensor_image = tensor_image.unsqueeze(0);
// 运行模型
torch::Tensor output = module.forward({tensor_image}).toTensor();
// 结果处理
int result = output.argmax().item();
std::cout << "The classifiction index is: " << result << std::endl;
return 0;
}