Как использовать ITK для преобразования PNG в тензор для PyTorch - PullRequest
Новая
       27

Как использовать ITK для преобразования PNG в тензор для PyTorch

0 голосов
/ 09 апреля 2020

Я пытаюсь запустить среду C ++ PyTorch и столкнулся со следующей проблемой:

Я успешно создал сценарий для модели и теперь готов к запуску. Теперь мне нужно подать png изображение на модель.

Я нашел кого-то с схожей проблемой в inte rnet, его идея заключалась в том, чтобы использовать модуль 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

Когда я пытаюсь распечатать журнал data, он содержит три числа, например 84 85 83, поэтому я предполагаю, что файл PNG успешно прочитан.

Однако я не могу заставить работать последнюю часть кода. Мне нужен тензор 3:1920:1080, но я не думаю, что значение трех RGBPixel правильно понимается функцией.

И кроме этого, я не понимаю, почему измерение установлено в 3.

Буду признателен за любую помощь.

1 Ответ

1 голос
/ 09 апреля 2020

Вам не нужно измерение 3, достаточно Dimension = 2. Если вам нужен макет RGBx1920x1080, то PixelType* rowdata = itk_img->GetBufferPointer(); даст вам этот макет без дальнейшей обработки. Поскольку torch::from_blob не становится владельцем буфера, другой человек пытался использовать .clone(). Вам не нужно этого делать, предполагая, что вы удерживаете itk_img в области видимости, или возитесь с его счетчиком ссылок и deleter.

cra sh, вероятно, происходит из-за того, что буфер имеет короткие пиксели (torch::kShort), когда он имеет uchar.

...