How to use ITK to convert PNG into tensor for PyTorch
I'm trying to run a C++ PyTorch framework and ran into the following problem.

I successfully scripted a model and is now ready to run. Now I have to feed in a png image to the model.

I found someone with simmilar issue on the internet, his idea was to use ITK module to read in a PNG file and convert it to a array, then make it to Tensor.

PNG -&gt; RGBPixel[] -&gt; tensor


So the following is what I am trying out right now.

using PixelTyupe = itk::RGBPixel&lt;unsinged char&gt;;
const unsigned int Dimension = 3;
typedef itk::Image&lt;PixelType, Dimension&gt;      ImageType;
typedef itk::ImageFileReader&lt;ImageType&gt;       ReaderType;
typedef itk::ImageRegionIterator&lt;ImageType&gt;   IteratorType;

typename ImageType::RegionType region = itk_img-&gt;GetLargestPossibleRegion();
const typename ImageType::SizeType size = region.GetSize();

int len = size[0] * size[1] * size[2]; // This ends up 1920 * 1080 * 1
PixelType rowdata[len];
int count = 0;
IteratorType iter(itk_img, itk_img-&gt;GetRequestedRegion());

// convert itk to array
for (iter.GoToBegin(); !iter.IsAtEnd(); ++iter) {
   rowdata[count] = iter.Get();
   count++;
} // count = 1920 * 1080

// convert array to tensor
tensor_img = torch::from_blob(rowdata, {3, (int)size[0], (int)size[1]}, torch::kShort). clone(); // Segmenation Fault


When I try to print log data, it holds three numbers like 84 85 83, so I suppose the PNG file is successfully read.

However, I can't get the last part of the code to work. What I need is 3:1920:1080 tensor, but I don't think the three RGBPixel value is understood properly by the function.

And apart from that, I don't see why the dimension is set to 3.

I would appreciate any kind of help.

You don't need dimension 3, Dimension = 2 is sufficient. If the layout you need is RGBx1920x1080, then PixelType* rowdata = itk_img-&gt;GetBufferPointer(); would get you that layout without further processing. As torch::from_blob does not take ownership of the buffer, the other person was trying to use .clone(). You don't have to do that either, assuming you keep the itk_img in scope, or muck around with its reference count and deleter.

The crash probably comes from saying the buffer has short pixels (torch::kShort), when it has uchar.
