Pre-trained model giving vague results

I am trying to implement this for a single image and not getting any mean value below 5.0. The good quality images also at times return low values.

I am sharing the main.py file, please check if anything is wrong with the code.


import argparse
import os

import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import torch
from torch import no_grad
import torch.autograd as autograd
import torch.optim as optim

import torchvision.transforms as transforms
import torchvision.datasets as dsets
import torchvision.models as models

import torch.nn.functional as F

from model import *

import cv2
file_name = 'bad'
filename = '/home/shayan/Projects/NIMA/images/'+file_name+'.jpg'

image = cv2.imread(filename)
image = cv2.resize(image,(224,224))

img_arr = image.transpose(2, 0, 1) # C x H x W
img_arr = np.expand_dims(img_arr,axis = 0)
print(img_arr.shape)

img_tensor = torch.from_numpy(img_arr)
img_tensor = img_tensor.type('torch.FloatTensor')
print(img_tensor.shape,img_tensor.size)

cuda = torch.cuda.is_available()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if cuda:
    print("Device: GPU")
else:
    print("Device: CPU")

base_model = models.vgg16(pretrained=True)
model = NIMA(base_model)

model.load_state_dict(torch.load("/home/shayan/Projects/NIMA/epoch-12.pkl", map_location=lambda storage, loc: storage))
print("Successfully loaded model")

with torch.no_grad():

    model.eval()

output = model(img_tensor)
output = output.view(10, 1)

predicted_mean, predicted_std = 0.0, 0.0
for i, elem in enumerate(output, 1):
    predicted_mean += i * elem
for j, elem in enumerate(output, 1):
    predicted_std += elem * (j - predicted_mean) ** 2
print("________________")
print(u"({}) \u00B1{}".format(round(float(predicted_mean),2), round(float(predicted_std), 2)))

yunxiaoshi / Neural-IMage-Assessment

Pre-trained model giving vague results #12