import torch
from torchvision import transforms
from PIL import Image, ImageFile
import pandas as pd
import os
import math
from model import ConvolutionalNet
from collections import Counter
from vector_dict import vector_dict

ImageFile.LOAD_TRUNCATED_IMAGES = True

model = ConvolutionalNet()
model.load_state_dict(torch.load('model.pt'))

transform = transforms.Compose([
  transforms.ToTensor(),
  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
  transforms.Resize((256, 256))
])

def get_prediction(path):
    img = Image.open(path)
    with torch.no_grad():
      pred = model(transform(img))
    return vector_dict[torch.max(pred, 1)[1].item()]

print(get_prediction('data/test/Afghanistan/39841.png'))