Refine Gradio App.

app.py

@@ -1,34 +1,34 @@
 import gradio as gr
-import torch
-from torchvision import transforms
-from torchvision import models
-from PIL import Image
-import requests
-
-# Load a pre-trained model
-model = models.resnet50(pretrained=True)
-model.eval()
-
-# Preprocess the input image
-def preprocess(image):
-    transform = transforms.Compose([
-        transforms.Resize(256),
-        transforms.CenterCrop(224),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-    ])
-    return transform(image).unsqueeze(0)
-
-
-def predict(image):
-    image = Image.fromarray(image.astype('uint8'), 'RGB')
-    preprocessed_img = preprocess(image)
-    outputs = model(preprocessed_img)
-    _, predicted = torch.max(outputs, 1)
-    return predicted.item()
+from src.config import Config
+from src.predict import predict
+from gradio.components import Image, Label, Button
 
 # Create a Gradio interface
 image = gr.inputs.Image()
 label = gr.outputs.Label(num_top_classes=1)
 
-gr.Interface(predict, image, label, capture_session=True).launch()
+with gr.Blocks(theme='gradio/monochrome') as demo:
+    with gr.Row():
+        with open('src/assets/header.md', 'r', encoding='utf-8') as f:
+            header = gr.Markdown(f.read())
+            
+    with gr.Row(align="center", align_items="center"):
+        input_image = Image(
+            label="Image",
+            interactive=True
+
+        )
+        output = Label(
+            num_top_classes=5,
+            label="Prediction"
+        )
+
+
+    with gr.Row():
+        predict_button = Button(value="Predict Animal", label="Predict", info="Click to make a prediction.")
+        predict_button.click(fn=predict, inputs=input_image, outputs=output)
+
+
+
+
+demo.launch()

src/__init__.py

@@ -0,0 +1,3 @@
+from .config import Config
+from .utils import *
+from .predict import predict

src/assets/header.md

@@ -0,0 +1,8 @@
+# 🦁 Wildlife Classification 🐘
+
+This simple application uses a pre-trained model to classify images of wildlife animals. This project is driven by the ambition to protect wildlife and to raise awareness of animal mistreatment. 
+### > *Species-appropriate is only the freedom.* (R. M. Lehmann)
+
+---
+Further information can be found in the github repository [here](https://github.com/ilyii/wildlife).
+

src/config.py

@@ -0,0 +1,21 @@
+import torch
+import torchvision
+
+from src.models import MobileNetV3Large
+
+
+class Config:
+    # - - - DATA - - - #
+    classes = [
+        'african buffalo', 'african elephant', 'alligator', 'alligators', 'alpine ibex', 'anaconda', 'antelope', 'badger', 'bat', 'bear', 'beavers', 'bee', 'beetle', 'bighorn sheep', 'bison', 'black bear', 'blue whale', 'boar', 'bottlenose dolphin', 'bowhead whale', 'butterfly', 'capybara', 'cat', 'caterpillar', 'catfish', 'chamois', 'cheetah', 'chicken', 'chimpanzee', 'clownfish', 'cockroach', 'common dolphin', 'cow', 'coyote', 'crab', 'crayfish', 'crocodile', 'crocodiles', 'crow', 'deer', 'dog', 'dolphin', 'donkey', 'dragonfly', 'duck', 'eagle', 'elephant', 'elk', 'false killer whale', 'flamingo', 'fly', 'fox', 'frogs', 'gazelle', 'geese', 'giraffe', 'goat', 'golden eagle', 'goldfish', 'goose', 'gorilla', 'grasshopper', 'gray whale', 'hamster', 'hare', 'hedgehog', 'himalayan tahr', 'hippopotamus', 'hornbill', 'horse', 'hummingbird', 'humpback dolphin', 'humpback whale', 'hyena', 'irrawaddy dolphin', 'jaguar', 'jellyfish', 'kangaroo', 'koala', 'ladybugs', 'leopard', 'lion', 'lizard', 'lobster', 'lynx', 'macaw', 'marmot', 'minke whale', 'monkey', 'moose', 'mosquito', 'moth', 'mountain goat', 'mouse', 'octopus', 'okapi', 'orangutan', 'orca', 'ostrich', 'otter', 'owl', 'ox', 'oyster', 'panda', 'panther', 'parrot', 'pelecaniformes', 'penguin', 'pig', 'pigeon', 'pika', 'poison dart frog', 'porcupine', 'possum', 'prairie dog', 'rabbit', 'raccoon', 'rat', 'red panda', 'reindeer', 'rhino', 'rhinoceros', 'right whale', 'rocky mountain elk', 'salmon', 'sandpiper', 'seahorse', 'seal', 'sei whale', 'shark', 'sheep', 'sloth', 'snails', 'snake', 'snow leopard', 'sparrow', 'sperm whale', 'spotted hyena', 'squid', 'squirrel', 'starfish', 'swan', 'tapir', 'tiger', 'toads', 'toucan', 'turkey', 'turtle', 'warthog', 'water snakes', 'whale', 'wildebeest', 'wolf', 'wolverine', 'wombat', 'woodpecker', 'zebra'
+        ]
+
+    num_classes = len(classes)
+
+    # - - - PARAMS - - - #
+    modelpath = './src/models/mobilenetv3_large_100_224.pt'
+    imgsize = 224
+    model = MobileNetV3Large(modelpath, num_classes)
+    
+
+

src/models/__init__.py

@@ -0,0 +1 @@
+from .mobilenet import *

src/models/mobilenet.py

@@ -0,0 +1,17 @@
+import torch
+import torchvision
+from torch import nn
+
+
+class MobileNetV3Large(nn.Module):
+
+    def __init__(self, ckpt, num_classes) -> None:
+        super().__init__()
+        self.model = torchvision.models.mobilenet_v3_large(pretrained=True)
+        self.model.classifier[3] = nn.Linear(self.model.classifier[3].in_features, num_classes)
+        self.model.load_state_dict(torch.load(ckpt))
+        self.model.to('cpu')
+        self.model.eval()
+
+    def forward(self, x):
+        return self.model(x)

src/predict.py

@@ -0,0 +1,32 @@
+import torch
+import gradio as gr
+from torchvision import transforms
+from torchvision import models
+from PIL import Image
+
+from src.config import Config
+from src.utils import id2label
+
+
+def predict(image):
+    if image is None:
+        raise gr.Error("No image found. Please upload an image to predict.")
+
+    transform = transforms.Compose([
+        transforms.Resize(int(Config.imgsize * 1.143)),
+        transforms.CenterCrop(Config.imgsize),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    model = Config.model
+    try:
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
+    except:
+        raise gr.Error("Image could not be converted to RGB. Please try another image.")
+    preprocessed_img = transform(image).unsqueeze(0)
+    outputs = model(preprocessed_img)
+    outputs = torch.softmax(outputs, dim=1)
+    preds = torch.topk(outputs, 5)
+    labels = {id2label(preds.indices[0][i].item()).capitalize(): preds.values[0][i].item() for i in range(5)}
+    return labels
+

src/utils.py

@@ -0,0 +1,4 @@
+from src.config import Config 
+
+def id2label(id:int):
+    return Config.classes[id]