Maxpool изображения в pytorch - PullRequest
Новая
       109

Maxpool изображения в pytorch

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

Я пытаюсь просто применить maxpool2d (из torch.nn) к одному изображению (а не как слой maxpool). Вот мой код прямо сейчас: 

name = 'astronaut'
imshow(images[name], name)
img = images[name]
# pool of square window of size=3, stride=1
m = nn.MaxPool2d(3,stride = 1)
img_transform = torch.Tensor(images[name])
plt.imshow(m(img_transform).view((512,510)))

Проблема в том, что этот код дает мне очень зеленое изображение в результате. Я уверен, что проблема связана с размерами вида, но я не смог найти, как применить maxpool только к одному изображению, поэтому я не смог это исправить. Размер изображения, которое я рассматриваю, составляет 512x512. Аргументы для представления сейчас не имеют смысла для меня, это всего лишь единственное число, которое дает результат ...

Если, например, я дал 512,512 в качестве аргумента для просмотра, я получаю следующую ошибку: 

RuntimeError: shape '[512, 512]' is invalid for input of size 261120

Если кто-нибудь подскажет, как применить maxpool, avgpool или minpool к изображению и отобразить результат, я был бы очень признателен!

Спасибо (:

1 Ответ

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

Предполагая, что ваше изображение равно numpy.array при загрузке (пожалуйста, смотрите комментарии для объяснения каждого шага): 

import numpy as np
import torch

# Assuming you have 3 color channels in your image
# Assuming your data is in Width, Height, Channels format
numpy_img = np.random.randint(low=0, high=255, size=(512, 512, 3))

# Transform to tensor
tensor_img = torch.from_numpy(numpy_img)
# PyTorch takes images in format Channels, Width, Height
# We have to switch their dimensions using `permute`
tensor_img = tensor_img.permute(2, 0, 1)
tensor_img.shape # Shape [3, 512, 512]

# Layers always need batch as first dimension (even for one image)
# unsqueeze will add it for you    
ready_tensor_img = tensor_img.unsqueeze(dim=0)
ready_tensor_img.shape # Shape [1, 3, 512, 512]

pooling = torch.nn.MaxPool2d(kernel_size=3, stride=1)

# You need to cast your image to float as
# pooling is not implemented for Tensors of type long
new_img = pooling(ready_tensor_img.float())

Если ваше изображение черно-белое, вам понадобится форма [1, 1, 512, 512] (один канал только), вы не можете оставить / сжать эти размеры, они всегда должны быть там для любого torch.nn.Module!

Чтобы снова преобразовать тензор в изображение, вы можете использовать аналогичные шаги: 

# Cast to long and squeeze batch dimension
no_batch = new_img.long().squeeze(dim=0)

# Unpermute
width_height_channels = no_batch.permute(1, 2, 0)
width_height_channels.shape  # Shape: [510, 510, 3]

# Cast to numpy and you have your image
final_image = width_height_channels.numpy()
...