app.py

import torch
from PIL import Image
from torchvision.transforms import v2
import torchvision.models as models
import os
import gradio as gr

# Load the pre-trained model that was used in the training script. In this case, it was the ResNet18 model
model = models.resnet18()

# Modify the final fully connected layer to have 73 output classes, same as in the training script
num_ftrs = model.fc.in_features
model.fc = torch.nn.Linear(num_ftrs, 73)

# Load the saved state dictionary (Device agnostic; inference for image classification is decent even on CPU)
model.load_state_dict(torch.load('Dogrun2.pth', map_location=torch.device('cpu')))

# Set the model to evaluation mode
model.eval()

# Automatically detect the available device (CPU or GPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Transforms the images to how they were tested for the model to read for inference. Keep Exactly the same as the transformation for the test and valid sets. No randomizing here!
transforms_test = v2.Compose([
    v2.Resize((224, 224), antialias=True),
    v2.CenterCrop((224, 224)),
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

labels = ['Afghan-Hound', 'Airedale-Terrier', 'Akita', 'Alaskan-Malamute', 'American-Foxhound', 'American-Hairless-Terrier', 'American-Water-Spaniel', 'Basenji', 'Basset-Hound', 'Beagle', 'Bearded-Collie', 'Belgian-Malinois', 'Belgian-Sheepdog', 'Bernese-Mountain-Dog', 'Bichon-Frise', 'Bloodhound', 'Bluetick-Coonhound', 'Border-Collie', 'Borzoi', 'Boston-Terrier', 'Boxer', 'Bull-Terrier', 'Bulldog', 'Bullmastiff', 'Cairn-Terrier', 'Cane-Corso', 'Cavalier-King-Charles-Spaniel', 'Chihuahua', 'Chinese-Crested', 'Chinese-Shar-Pei', 'Chow-Chow', 'Clumber-Spaniel', 'Cockapoo', 'Cocker-Spaniel', 'Collie', 'Dachshund', 'Dalmatian', 'Doberman-Pinscher', 'French-Bulldog', 'German-Shepherd', 'German-Shorthaired-Pointer', 'Golden-Retriever', 'Great-Dane', 'Great-Pyrenees', 'Greyhound', 'Irish-Water-Spaniel', 'Irish-Wolfhound', 'Japanese-Chin', 'Komondor', 'Labradoodle', 'Labrador-Retriever', 'Lhasa-Apso', 'Maltese', 'Miniature-Schnauzer', 'Newfoundland', 'Norwegian-Elkhound', 'Pekingese', 'Pembroke-Welsh-Corgi', 'Pomeranian', 'Poodle', 'Pug', 'Rhodesian-Ridgeback', 'Rottweiler', 'Saint-Bernard', 'Samoyed', 'Scottish-Terrier', 'Shiba-Inu', 'Shih-Tzu', 'Siberian-Husky', 'Staffordshire-Bull-Terrier', 'Vizsla', 'Xoloitzcuintli', 'Yorkshire-Terrier']

# I have the breed_nicknames dictionary set because some breeds aren't recognized that much by the official breed, such as the ones below.
breed_nicknames = {
    'Xoloitzcuintli': ' (Mexican Hairless)',
    'Staffordshire-Bull-Terrier': ' (Pitbull)',
    'Pembroke-Welsh-Corgi': ' (Corgi)',
}

def predict(input_img):
    transformed_img = transforms_test(input_img)
    transformed_img = transformed_img.unsqueeze(0).to(device)
    logits = model(transformed_img)
    output_softmax = torch.nn.functional.softmax(logits, dim=1)
    topk_values, topk_indices = torch.topk(output_softmax, 3)
    topk_indices = topk_indices.tolist()[0]
    topk_labels = [labels[index] for index in topk_indices]
    topk_probs = topk_values[0].tolist()
    result = "\n".join([f"{label:<20}: {prob*100:.2f}%" for label, prob in zip(topk_labels, topk_probs)])
    return result

gradio_app = gr.Interface(
    predict,
    inputs=gr.Image(label="Please select a clear image of your good dog to upload.", sources=['upload'], type="pil"),
    outputs=[gr.Label(label="Result")],
    title="What Breed is Your Dog??",
)


if __name__ == "__main__":
    gradio_app.launch()