如何使用 ITK 將 PNG 轉換為 PyTorch 的張量

Question

我正在嘗試運行 C++ PyTorch 框架並遇到以下問題。

我成功編寫了 model 腳本，現在可以運行了。 現在我必須將png圖像輸入 model。

我在互聯網上發現了一個有類似問題的人，他的想法是使用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張量，但我認為 function 不能正確理解這三個 RGBPixel 值。

除此之外，我不明白為什么將維度設置為 3。

我將不勝感激任何幫助。

Answer 1

您不需要維度 3， Dimension = 2就足夠了。 如果你需要的布局是RGBx1920x1080，那么PixelType* rowdata = itk_img->GetBufferPointer(); 無需進一步處理即可獲得該布局。 由於torch::from_blob不擁有緩沖區的所有權，因此其他人試圖使用.clone() 。 您也不必這樣做，假設您將itk_img保留在 scope 中，或者使用它的引用計數和deleter亂七八糟。

崩潰可能來自說緩沖區有短像素（ torch::kShort ），當它有uchar時。