app.py

### 1. Imports and class names setup ### 
import gradio as gr
import os
import torch

from model import create_effnetb2_model
from timeit import default_timer as timer
from typing import Tuple, Dict

### 2. Model and transforms preparation ###

class_names = ['Afghan_hound','African_hunting_dog','Airedale','American_Staffordshire_terrier','Appenzeller','Australian_terrier',
 'Bedlington_terrier','Bernese_mountain_dog','Blenheim_spaniel','Border_collie','Border_terrier','Boston_bull','Bouvier_des_Flandres',
 'Brabancon_griffon','Brittany_spaniel','Cardigan','Chesapeake_Bay_retriever','Chihuahua','Dandie_Dinmont','Doberman','English_foxhound',
 'English_setter','English_springer','EntleBucher','Eskimo_dog','French_bulldog','German_shepherd','German_short-haired_pointer','Gordon_setter',
 'Great_Dane','Great_Pyrenees','Greater_Swiss_Mountain_dog','Ibizan_hound','Irish_setter','Irish_terrier','Irish_water_spaniel','Irish_wolfhound',
 'Italian_greyhound','Japanese_spaniel','Kerry_blue_terrier','Labrador_retriever','Lakeland_terrier','Leonberg','Lhasa','Maltese_dog','Mexican_hairless',
 'Newfoundland','Norfolk_terrier','Norwegian_elkhound','Norwich_terrier','Old_English_sheepdog','Pekinese','Pembroke','Pomeranian','Rhodesian_ridgeback',
 'Rottweiler','Saint_Bernard','Saluki','Samoyed','Scotch_terrier','Scottish_deerhound','Sealyham_terrier','Shetland_sheepdog','Shih-Tzu','Siberian_husky',
 'Staffordshire_bullterrier','Sussex_spaniel','Tibetan_mastiff','Tibetan_terrier','Walker_hound','Weimaraner','Welsh_springer_spaniel','West_Highland_white_terrier',
 'Yorkshire_terrier','affenpinscher','basenji','basset','beagle','black-and-tan_coonhound','bloodhound','bluetick','borzoi','boxer','briard','bull_mastiff',
 'cairn','chow','clumber','cocker_spaniel','collie','curly-coated_retriever','dhole','dingo','flat-coated_retriever','giant_schnauzer', 'golden_retriever',
 'groenendael','keeshond','kelpie','komondor','kuvasz','malamute','malinois','miniature_pinscher','miniature_poodle','miniature_schnauzer','otterhound',
 'papillon','pug','redbone','schipperke','silky_terrier','soft-coated_wheaten_terrier','standard_poodle','standard_schnauzer','toy_poodle','toy_terrier',
 'vizsla','whippet','wire-haired_fox_terrier']

# Create EffNetB2 model
effnetb2, effnetb2_transforms = create_effnetb2_model(
    num_classes=120,
)

# Load saved weights
effnetb2.load_state_dict(
    torch.load(
        f="DOG-BREED-CLASSIFICATION-MODEL.pth",
        map_location=torch.device("cpu"),  # load to CPU
    )
)


### 3. Predict function ###

def predict(img) -> Tuple[Dict, float]:
    """Transforms and performs a prediction on img and returns prediction and time taken.
    """
    # Start the timer
    start_time = timer()
    
    # Transform the target image and add a batch dimension
    img = effnetb2_transforms(img).unsqueeze(0)
    
    # Put model into evaluation mode and turn on inference mode
    effnetb2.eval()
    with torch.inference_mode():
        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
        pred_probs = torch.softmax(effnetb2(img), dim=1)
    
    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
    
    # Calculate the prediction time
    pred_time = round(timer() - start_time, 5)
    
    # Return the prediction dictionary and prediction time 
    return pred_labels_and_probs, pred_time


### 4. Gradio app ###

# Create title, description and article strings
title = "🐶DOG BREED CLASSIFICATION APP🐶"
description = "computer vision model to classify different breeds of dogs from image"
article = "Created at [Google Colab](https://colab.research.google.com/drive/1pxA88oSjeoXS4Kt9fKAWIBaHlXwJ0vQr#scrollTo=e7MgaOxqEVJn)."

example_list = [["examples/" + example] for example in os.listdir("examples")]

# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs=gr.Image(type="pil"), # what are the inputs?
                    outputs=[gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?
                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
                    examples=example_list, 
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch()