How to test single image for a model trained with DataLoader-loaded dataset?

vision
1

I’m currently working on replicating the original training results with this SRResNet/SRGAN Implementation. This repository did not provide a script for testing purposes so I want to write my own.

This implementation uses a DataLoader to read individual images for training, in my testing script I attempt to write a test by pump in a single image file:

import argparse
import os
import numpy as np
import math
import itertools
import sys

import torchvision.transforms as transforms
from torchvision.utils import save_image, make_grid
from torch.utils.data import DataLoader
from torch.autograd import Variable
import PIL.Image as pil_image

from models import *
from datasets import *

import torch.nn as nn
import torch.nn.functional as F
import torch

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights-file', type=str, required=True)
    parser.add_argument('--image-file', type=str, required=True)
    args = parser.parse_args()

    cuda = torch.device('cuda:0')

    test = GeneratorResNet()
    calc_ssim = pytorch_ssim.SSIM(window_size = 11)

    state = test.state_dict()
    for n, p in torch.load(args.weights_file, map_location=lambda storage, loc: storage).items():
        if n in state.keys():
            state[n].copy_(p)
        else:
            raise KeyError(n)
    # test.load_state_dict(torch.load(args.weights_file))
    test.to(cuda)
    test.eval()

    image = pil_image.open(args.image_file).convert('RGB')
    image_width = (image.width // 4) * 4
    image_height = (image.height // 4) * 4
    image = image.resize((image_width, image_height), resample=pil_image.BICUBIC)
    image_bicubic = image.resize((image.width // 4, image.height // 4), resample=pil_image.BICUBIC)
    image_bicubic = image.resize((image.width * 4, image.height * 4), resample=pil_image.BICUBIC)
    image_bicubic.save(args.image_file.replace('.', '_bicubic_x4.'))

    image = np.array(image).astype(np.float32)
    image /= 255.
    image = torch.from_numpy(image).to(cuda)
    image = image.unsqueeze(-1)
    preds = test(image)

However I have an error regarding the image’s dimension does not match my tensor dimension in the input layer:

Traceback (most recent call last):
  File "test.py", line 56, in <module>
    preds = test(image)
  File "C:\Python38\lib\site-packages\torch\nn\modules\module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "D:\FYP\SRGAN-pytorch\models.py", line 77, in forward
    out1 = self.conv1(x)
  File "C:\Python38\lib\site-packages\torch\nn\modules\module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Python38\lib\site-packages\torch\nn\modules\container.py", line 100, in forward
    input = module(input)
  File "C:\Python38\lib\site-packages\torch\nn\modules\module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "C:\Python38\lib\site-packages\torch\nn\modules\conv.py", line 353, in forward
    return self._conv_forward(input, self.weight)
  File "C:\Python38\lib\site-packages\torch\nn\modules\conv.py", line 349, in _conv_forward
    return F.conv2d(input, weight, self.bias, self.stride,
RuntimeError: Given groups=1, weight of size [64, 3, 9, 9], expected input[720, 840, 3, 1] to have 3 channels, but got 840 channels instead

Since I’m still very new to the inner working of tensors and models I want to know if there’s anything wrong with my current approach, or what do I need to do further to pipe the image into the model to test?

2

Hi,

The error is about the channel logic. numpy is channel last, pytorch is channel first, the error says you are using numpy logic for pytorch tensors. Instead of converting your PIL images to numpy then reading those arrays as tensor using from_numpy, you can directly convert the PIL image to proper tensor using transforms.ToTensor. Other workaround would be permuting channels of the image after converting it to tensor using from_numpy by using torch.permute.

The idea is you should get [1, 3, 720, 840], instead of [720, 840, 3, 1].

Bests

1 Like
3

Problem solved using torch.permute. Much Thanks!

4

You can use this code for test single image for your model train:

import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader,Dataset
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
def pre_image(image_path,model):
   img = Image.open(image_path)
   mean = [0.485, 0.456, 0.406] 
   std = [0.229, 0.224, 0.225]
   transform_norm = transforms.Compose([transforms.ToTensor(), 
   transforms.Resize((224,224)),transforms.Normalize(mean, std)])
   # get normalized image
   img_normalized = transform_norm(img).float()
   img_normalized = img_normalized.unsqueeze_(0)
   # input = Variable(image_tensor)
   img_normalized = img_normalized.to(device)
   # print(img_normalized.shape)
   with torch.no_grad():
      model.eval()  
      output =model(img_normalized)
     # print(output)
      index = output.data.cpu().numpy().argmax()
      classes = train_ds.classes
      class_name = classes[index]
      return class_name