Load a single image in a pretrained pytorch net

Question

Total newbie here, I'm using this pytorch SegNet implementation with a '.pth' file containing weights from a 50 epochs training. How can I load a single test image and see the net prediction? I know this may sound like a stupid question but I'm stuck. What I've got is:

from segnet import SegNet
import torch

model = SegNet(2)
model.load_state_dict(torch.load('./model_segnet_epoch50.pth'))

How do I "use" the net on a single test picture?

Answer 1

I provide with an example of ResNet152 pre-trained model.

def image_loader(loader, image_name):
    image = Image.open(image_name)
    image = loader(image).float()
    image = torch.tensor(image, requires_grad=True)
    image = image.unsqueeze(0)
    return image

data_transforms = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor()
])


model_ft = models.resnet152(pretrained=True)
model_ft.eval()

print( np.argmax(model_ft(image_loader(data_transforms, $FILENAME)).detach().numpy()))

$FILENAME is the path and name of your image to be loaded. I got necessary help from this post .

Answer 2

output = model(image) .

Note that the image should be a Variable object and that the output will be as well. If your image is, for example, a Numpy array, you can convert it like so:

var_image = Variable(torch.Tensor(image))