File size: 730 Bytes
a189a79 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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'))
|