如何使用 ITK 将 PNG 转换为 PyTorch 的张量

Question

我正在尝试运行 C++ PyTorch 框架和运行解决以下问题。

我成功编写了模型脚本，现在可以运行。现在我必须将 png 图像输入模型。

我在网上发现有人有类似的问题，他的想法是使用ITK模块读取PNG文件并将其转换为数组，然后将其制成Tensor。

PNG -> RGBPixel[] -> tensor

下面是我现在正在尝试的。

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

typename ImageType::RegionType region = itk_img->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->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

当我尝试打印日志数据时，它包含三个数字，如 84 85 83，所以我认为 PNG 文件已成功读取。

但是，我无法使代码的最后一部分起作用。我需要的是 3:1920:1080 张量，但我认为函数无法正确理解三个 RGBPixel 值。

除此之外，我不明白为什么维度设置为 3。

如有任何帮助，我将不胜感激。

Answer 1

不需要维度 3，Dimension = 2 就足够了。如果您需要的布局是 RGBx1920x1080，则 PixelType* rowdata = itk_img->GetBufferPointer(); 无需进一步处理即可为您提供该布局。由于 torch::from_blob 没有取得缓冲区的所有权，其他人试图使用 .clone()。您也不必这样做，假设您将 itk_img 保持在范围内，或者处理它的引用计数和 deleter.

崩溃可能是因为说缓冲区有短像素 (torch::kShort)，当它有 uchar.

时

如何使用 ITK 将 PNG 转换为 PyTorch 的张量

How to use ITK to convert PNG into tensor for PyTorch

c++

itk

pytorch