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

# Setup class names
class_names= ['alpine sea holly',
 'anthurium',
 'artichoke',
 'azalea',
 'ball moss',
 'balloon flower',
 'barbeton daisy',
 'bearded iris',
 'bee balm',
 'bird of paradise',
 'bishop of llandaff',
 'black-eyed susan',
 'blackberry lily',
 'blanket flower',
 'bolero deep blue',
 'bougainvillea',
 'bromelia',
 'buttercup',
 'californian poppy',
 'camellia',
 'canna lily',
 'canterbury bells',
 'cape flower',
 'carnation',
 'cautleya spicata',
 'clematis',
 "colt's foot",
 'columbine',
 'common dandelion',
 'corn poppy',
 'cyclamen',
 'daffodil',
 'desert-rose',
 'english marigold',
 'fire lily',
 'foxglove',
 'frangipani',
 'fritillary',
 'garden phlox',
 'gaura',
 'gazania',
 'geranium',
 'giant white arum lily',
 'globe thistle',
 'globe-flower',
 'grape hyacinth',
 'great masterwort',
 'hard-leaved pocket orchid',
 'hibiscus',
 'hippeastrum',
 'japanese anemone',
 'king protea',
 'lenten rose',
 'lotus lotus',
 'love in the mist',
 'magnolia',
 'mallow',
 'marigold',
 'mexican aster',
 'mexican petunia',
 'monkshood',
 'moon orchid',
 'morning glory',
 'orange dahlia',
 'osteospermum',
 'oxeye daisy',
 'passion flower',
 'pelargonium',
 'peruvian lily',
 'petunia',
 'pincushion flower',
 'pink primrose',
 'pink-yellow dahlia',
 'poinsettia',
 'primula',
 'prince of wales feathers',
 'purple coneflower',
 'red ginger',
 'rose',
 'ruby-lipped cattleya',
 'siam tulip',
 'silverbush',
 'snapdragon',
 'spear thistle',
 'spring crocus',
 'stemless gentian',
 'sunflower',
 'sweet pea',
 'sweet william',
 'sword lily',
 'thorn apple',
 'tiger lily',
 'toad lily',
 'tree mallow',
 'tree poppy',
 'trumpet creeper',
 'wallflower',
 'water lily',
 'watercress',
 'wild pansy',
 'windflower',
 'yellow iris'
 ]

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

# Create EffNetB2 model
effnetb2, effnetb2_transforms = create_effnetb2_model(
    num_classes=102, # len(class_names) would also work
)

# Load saved weights
effnetb2.load_state_dict(
    torch.load(
        f="pretrained_effnetb2_feature_extractor_fl102.pth",
        map_location=torch.device("cpu"),  # load to CPU
    )
)

### 3. Predict function ###

# Create 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 ###

title = "Flofi Demo"
description = "An EfficientNetB2 feature extractor computer vision model to classify images of 102 flower species."
article = "Created by Haydar Uçar."

# 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
                    title=title,
                    description=description,
                    article=article)

# Launch the demo!
demo.launch() # generate a publically shareable URL?