Commit 641b581f authored by fhuerliman's avatar fhuerliman
Browse files

Upload New File

parent dbd7a26f
# Load Libraries
import numpy as np
import torch
from PIL import Image
print('================')
print('Painter oracle (classify image according to painter label)')
print('================')
class_names = ['auguste-renoir', 'claude-monet', 'vincent-van-gogh']
print('labels: ', class_names)
device = torch.device("cpu")
# Load the model (re-trained ResNet-18 for painter classification)
model_path = 'painter_classifier_model_ResNet-18.pt'
print('model: ', model_path)
model_ft = torch.load(model_path, map_location=torch.device('cpu'))
def process_image(image_path):
#image = Image.open(image_path)
image = image_path
# Resize
img = image.resize((256, 256))
# Center crop
width = 256
height = 256
new_width = 224
new_height = 224
left = (width - new_width) / 2
top = (height - new_height) / 2
right = (width + new_width) / 2
bottom = (height + new_height) / 2
img = img.crop((left, top, right, bottom))
# Convert to numpy, transpose color dimension and normalize
img = np.array(img).transpose((2, 0, 1)) / 256
# Standardization
means = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
stds = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
img = img - means
img = img / stds
img_tensor = torch.Tensor(img)
return img_tensor
im_path = 'monet.jpg'
print('image: ',im_path)
im = Image.open(im_path)
img_tensor = process_image(im)
img_tensor = img_tensor.view(1,3,224,224)
prediction = model_ft(img_tensor.to(device))
print("-------------------------------------")
print("Prediction of Category:",class_names[prediction.argmax()])
print("Probability is:",round(torch.nn.functional.softmax(prediction,dim=1).max().item()*100,2),"%")
print("-------------------------------------")
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment