changes

app.py

@@ -0,0 +1,48 @@
+import gradio as gr
+from model import create_effnetb2_model
+import os
+import torch
+from torch import nn
+from typing import List, Dict, Tuple
+from timeit import default_timer as timer
+
+class_names = ['pizza', 'steak', 'sushi']
+model, transforms = create_effnetb2_model(num_classes=len(class_names))
+
+ckpt = torch.load('effnet_ckpt.tar', map_location='cpu')
+model.load_state_dict(ckpt['model_state_dict'])
+model.to('cpu')
+
+
+def predict(img) -> Tuple[Dict, float]:
+    start = timer()
+    img = transforms(img)
+    img = img.unsqueeze(0)
+    img = img.to('cpu')
+    model.to('cpu')
+    model.eval()
+    with torch.inference_mode():
+        pred_logits = model(img)
+        pred_probs = nn.Softmax(dim=1)(pred_logits).squeeze(0)
+        pred_probs_dict = {class_names[i]: pred_probs[i].item() for i in range(len(class_names))}
+    end = timer()
+    return pred_probs_dict, round(end - start, 4)
+
+
+examples_dir = 'examples'
+examples = [[os.path.join(examples_dir, f)] for f in os.listdir(examples_dir)]
+
+import gradio as gr
+title = "Pizza, Steak, Sushi Classifier 🍕🥩🍣"
+description = "This efficientnetb2 model classifies images of pizza, steak, and sushi."
+article = "Created for practice using [Gradio](https://www.gradio.app/)"
+demo = gr.Interface(fn=predict,
+                    inputs=gr.Image(type="pil", label="Image of Pizza, Steak, or Sushi"),
+                    outputs=[gr.Label(label="Predictions", num_top_classes=len(class_names)),
+                             gr.Number(label="Prediction Time (s)")],
+                    examples=examples,
+                    title=title,
+                    description=description,
+                    article=article)
+demo.launch(debug=False,
+            share=True)

effnet_ckpt.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d650522724426fc86d8c3f630cc399b26ce5df5c35679be1b21e0aa94493498c
+size 31275834

model.py

@@ -0,0 +1,12 @@
+from torch import nn
+from torchvision.models import efficientnet_b2, EfficientNet_B2_Weights
+
+
+def create_effnetb2_model(num_classes=3):
+    weights_effnetb2 = EfficientNet_B2_Weights.DEFAULT
+    transforms_effnetb2 = weights_effnetb2.transforms()
+    model_effnetb2 = efficientnet_b2(weights=weights_effnetb2)
+    for param in model_effnetb2.parameters():
+        param.requires_grad = False
+    model_effnetb2.classifier[1] = nn.Linear(in_features=1408, out_features=num_classes)
+    return model_effnetb2, transforms_effnetb2

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.2.0
+torchvision>=0.17.0
+gradio>=4.26.0