DimitrisKatos/AnimalClassification

1 2### 1. Imports and class names setup3import gradio as gr4import os 5import torch6import gradio as gr7import torchvision8 9from model import create_effnetb2_model10from timeit import default_timer as timer11from typing import Dict, Tuple12 13class_names = ['butterfly',  'cat',  'chicken',  'cow', 'dog',14               'elephant',  'horse',   'sheep',  'spider', 'squirrel']15 16### 2. Model and transforms prepartaion ###17effnetb2, effnetb2_transforms = create_effnetb2_model()18 19# Loade the save weights.20effnetb2.load_state_dict(torch.load(f = "effnetb2_model.pth",21                                    map_location = torch.device("cpu")))22 23### 3. Predict Function ###24effnetb2 = effnetb2.to('cpu')25def predict(img) -> Tuple[Dict, float]:26    """Transforms and performs a prediction on img and returns prediction and time taken.27    """28    29    # Start the timer30    start_time = timer()31    32    # Transform the target image and add a batch dimension33    img = effnetb2_transforms(img).unsqueeze(0)34    35    # Put model into evaluation mode and turn on inference mode36    effnetb2.eval()37    with torch.inference_mode():38        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities39        pred_probs = torch.softmax(effnetb2(img), dim=1)40    41    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)42    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}43    44    # Calculate the prediction time45    pred_time = round(timer() - start_time, 5)46    47    # Return the prediction dictionary and prediction time 48    return pred_labels_and_probs, pred_time49 50### 4. ### 51 52# Create title, description and article strings53title = "AnimalsClassification "54description = """An EfficientNetB2 feature extractor computer vision model to classify images of ten different animals.55                Curently the app can identify 10 diffferent animal species which is the following.56                1. Dog57                2. Cat58                3. Horse59                4. Butterfly60                5. Cow61                6. Chicken62                7. Sheep63                8. Squirrel64                9. Elephant65                10. Spider"""66article = "ModelDeployment"67 68# Create example list.69example_list = [["examples/" + example] for example in os.listdir('examples')]70 71# Create the Gradio demo72demo = gr.Interface(fn=predict, # mapping function from input to output73                    inputs=gr.Image(type="pil"), # what are the inputs?74                    outputs=[gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?75                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs76                    examples=example_list, 77                    title=title,78                    description=description,79                    article=article)80 81# Launch the demo!82demo.launch(debug=False, # print errors locally?83            share=True) # generate a publically shareable URL?84