CODEWITHSUNDEEP
pytorchdimensionconv-neural-networkmax-pooling
user11764168
user11764168

Reputation:

Maxpool of an image in pytorch

I'm trying to just apply maxpool2d (from torch.nn) on a single image (not as a maxpool layer). Here is my code right now:

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)))

The issue is, this code gives me a very green image as a result. I am sure the problem is with the dimensions of view, but I was unable to find how to apply maxpool to just one image so I couldn't fix it. The dimension of the image I'm considering is 512x512. The arguments for view make no sense for me right now, it's just the only number that gives a result...

If for example, I gave 512,512 as the argument for view, I get the following error:

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

If anyone can tell me how to apply maxpool, avgpool, or minpool to an image and display the result I would be super grateful!

Thanks (:

Upvotes: 1

Views: 3845

Answers (1)

Szymon Maszke
Szymon Maszke

Reputation: 24691

Assuming your image is a numpy.array upon loading (please see comments for explanation of each step):

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())

If your image is black and white you would need shape [1, 1, 512, 512] (single channel only), you can't leave/squeeze those dimensions, they always have to be there for any torch.nn.Module!

To transform tensor into image again you could use similar steps:

# 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()

Upvotes: 3

Related Questions