Loading opencv image to pytorch tensor

[jainshobhit]

❓ Questions and Help

I am trying to load an image in OpenCV Mat variable and then converting it into tensor for passing it into my TorchScript model. I followed #12506 for loading the image however, I am not sure whether it is the correct way or not.
Here is my code

#include <torch/script.h> // One-stop header.
#include <iostream>
#include <memory>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>

using namespace cv;
using namespace std;

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }
  // Deserialize the ScriptModule from a file using torch::jit::load().
  std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);

  assert(module != nullptr);
  std::cout << "ok\n";

  Mat image_bgr, image;
  image_bgr = imread("/home/landscape_org.jpg");
  cvtColor(image_bgr, image, COLOR_BGR2RGB);
  for (int j=0;j<10;j++)
  {
    cout<<image.at<Vec3b>(0,j)<<endl;
  }
  at::Tensor tensor_image = torch::from_blob(image.data, {1, 3, image.rows, image.cols}, at::kByte);
  tensor_image = tensor_image.to(at::kFloat);
  cout<<tensor_image.slice(2,0,1)<<endl;

  // Create a vector of inputs.
  std::vector<torch::jit::IValue> input;
  input.emplace_back(tensor_image);

  // Execute the model and turn its output into a tensor.
  auto output = module->forward(input).toTuple()->elements()[6].toTensor().clone().clamp(0,255);
  Mat output_mat(cv::Size(1920,1080), CV_8UC3, output.data<float>());

  Mat output8, output_bgr;
  cvtColor(output8, output_bgr, COLOR_RGB2BGR);
  imwrite("landscape_output.jpg", output_bgr);
}

The output for the first 10 pixel values is

[53, 149, 249] 
[52, 148, 248] 
[53, 149, 249] 
[55, 151, 251]
[58, 154, 254]
[58, 154, 254]
[61, 155, 255]
[61, 155, 255]
[58, 152, 252]
[58, 152, 252]

And the output on calling the slice function on the resulting tensor (cout<<tensor_image.slice(2,0,1)<<endl;) is (only mentioning the first few columns of the R color channel):

(1,1,.,.) =
Columns 1 to 15   53  149  249   52  148  248   53  149  249   55  151  251   58  154  254
Columns 16 to 30   58  154  254   61  155  255   61  155  255   58  152  252   58  152  252

This clearly shows that the values are not being copied correctly in the tensor.

I am not able to figure out what is the correct way to perform this step since I was not able to find adequate documentation.

@goldsborough

[SsnL]

in opencv, the channel dimension is the last dimension. so you would do {1, image.rows, image.cols, 3} in torch::from_blob and transpose afterwards.

[jainshobhit]

Thanks for that information. I had just figured it out.
I am still trying to figure out how to convert the output tensor to OpenCV Mat variable.
Mat output_mat({1920, 1080}, CV_8UC3, output.data<int>()); does not work as mentioned here in Mat constructor documentation.
Haven't been able to find any other documentation yet on this.

@SsnL Any comments?

[bhack]

@jainshobhit If you want an 8UC the output is float so you need to convert in at::kByte and access with data<uint8_t>() not with int.

[jainshobhit]

Thanks @bhack

[TomHeaven]

@SsnL Do you know how to transpose a tensor? I cannot find documentation related to it.

[TomHeaven]

I found it at https://pytorch.org/cppdocs/api/classat_1_1_tensor.html#exhale-class-classat-1-1-tensor.
Here is a demo:

 tensor_image = tensor_image.permute({0, 3, 1, 2});

[jainshobhit]

A much cleaner way to transpose the axes is also provided in the ATen library. You can find an example for how to use it at https://github.com/jainshobhit/pytorch-cpp-examples/blob/master/libtorch_inference.cpp#L39.
Specifically,
tensor_image = at::transpose(tensor_image, 1, 2);

[martinruenz]

@jainshobhit I am curious, why is at::transpose cleaner? Is it faster? I see the advantage of tensor_image.permute that you only need 1 rather than 2 function calls.

[jainshobhit]

tensor_image.permute also uses abstractions from the ATen library which makes it slower than at::transpose. Moreover, at::transpose gives you more control over how you wish to swap the axes such that the overall permutation is faster.

[xhuvom]

What about the vice-versa? I mean Pytorch tensor array to cv::Mat image?