initial_commit

app.py

@@ -0,0 +1,62 @@
+import time
+import gradio as gr
+from pathlib import Path
+from models import *
+
+class_idx_to_names = {
+    0: "pizza", 
+    1: "steak", 
+    2: "sushi"
+}
+
+examples = [[str(path)] for path in Path(r"examples").glob("*")]
+
+def predict_one(model, transforms, image, device, class_idx_to_names):
+    model.eval()
+    model = model.to(device)
+    with torch.inference_mode():
+            
+        start_time = time.perf_counter()
+        image_transformed = transforms(image).unsqueeze(dim = 0).to(device)
+        
+        y_logits = model(image_transformed)
+        
+        y_preds = torch.softmax(y_logits, dim = 1)
+        y_probs = torch.argmax(y_preds, dim = 1)
+        
+        end_time = time.perf_counter()
+
+        predictions = {class_idx_to_names[index]: x.item() for index, x in enumerate(y_preds[0])}
+
+    return predictions, end_time - start_time
+
+def predict(image, model_choice):
+
+    if model_choice is None or model_choice == "effnet_b2":
+        model, transforms = get_effnet_b2()
+    else:
+        model, transforms = get_vit_16_base_transformer()
+    
+    predictions, time_taken = predict_one(model, transforms, image, "cpu", class_idx_to_names)
+    return predictions, time_taken
+
+
+title = "Food Recognition 🍕🍕"
+desc = "A dual model app ft. EfficientNetB2 Feature Extractor and VisionTransformer."
+article = '''
+## Stats on different Models
+---
+| Model Name      | Train Loss | Test Loss | Train Accuracy | Test Accuracy | Num Parameters | Model Size |
+|-----------------|------------|-----------|----------------|---------------|----------------|------------|
+| EfficientNet_b2 | 0.340270   | 0.301134  | 0.906250       | 0.953409      | 7705221        | 29.91 MB   |
+| ViT_Base_16     | 0.040448   | 0.055140  | 0.995833       | 0.981250      | 85800963       | 327.39 MB  |
+'''
+demo = gr.Interface(fn = predict, 
+                    inputs = [gr.Image(type = "pil", label = "upload an Jpeg or Png"), gr.Radio(["effnet_b2", "ViT (Vision Transformer)"], label = "choose model (default on effnet)")],
+                    outputs = [gr.Label(num_top_classes=3, label = "predictions"), gr.Number(label = "Prediction Time in seconds")], 
+                    examples = examples, 
+                    title = title,
+                    description=desc,
+                    article=article)
+
+demo.launch(debug = False)

models.py

@@ -0,0 +1,17 @@
+import torch
+
+from torchvision.models import vit_b_16, ViT_B_16_Weights
+from torchvision.models import efficientnet_b2, EfficientNet_B2_Weights
+
+
+def get_vit_16_base_transformer():
+    vit_b_16_model = torch.load(r"models\vit_16_base_custom_head_3_classes.pth")
+    vit_b_16_transforms = ViT_B_16_Weights.DEFAULT.transforms()
+
+    return vit_b_16_model, vit_b_16_transforms
+
+def get_effnet_b2():
+    eff_net_b2_model = torch.load(r"models\eff_netb2_custom_head_3_classes.pth")
+    eff_net_b2_transforms = EfficientNet_B2_Weights.DEFAULT.transforms()
+
+    return eff_net_b2_model, eff_net_b2_transforms

models/eff_netb2_custom_head_3_classes.pth

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

models/vit_16_base_custom_head_3_classes.pth

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

requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+gradio