evaluate.py

import copy

import torch
import torchvision.models as models
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tqdm import tqdm
# Define transformations for the images
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# Load the dataset
dataset = datasets.ImageFolder(root='../../Downloads/orc_vs_elf', transform=transform)

print(dataset.classes)
# Create a DataLoader
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=0)

# Create an instance of the model
model = models.resnet18(pretrained=False)
model.fc = nn.Linear(512, len(dataset.classes))

# Load the model state dictionary
model.load_state_dict(torch.load('model-elf.pth'))

def imshow_(inp, title=None):
    """Imshow for Tensor."""
    inp = inp .permute(1, 2, 0).numpy()
    print(inp.shape)
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)

    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)
    plt.show()

imageNames = ['../../Downloads/elf_test.jpg']
for imageName in imageNames:
    image = Image.open(imageName)
    transform = composed = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
    x = transform(image)
    z=model(x.unsqueeze_(0))
    _,yhat=torch.max(z.data, 1)
    # print(yhat)
    prediction = "Orc"
    if yhat == 1:
        prediction = "Elf"
    if yhat == 2:
        prediction = "Other"
    imshow_(transform(image),imageName+": Prediction = "+prediction)