Update app.py

app.py

@@ -2,24 +2,71 @@ import torch
 import gradio as gr
 from PIL import Image
 import numpy as np
-from torchvision import transforms
-
-from model import FoodIngredientClassifier
-from utils import load_model_and_mlb
+from torchvision import transforms, models
+import torch.nn as nn
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load model + label binarizer
-model, mlb, threshold = load_model_and_mlb("model.pth", DEVICE)
+# -------------------------
+# MODEL DEFINITION
+# -------------------------
+class FoodIngredientClassifier(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+
+        self.backbone = models.vit_b_16(
+            weights=models.ViT_B_16_Weights.DEFAULT
+        )
+
+        num_features = self.backbone.heads.head.in_features
+
+        self.backbone.heads = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(num_features, 1024),
+            nn.BatchNorm1d(1024),
+            nn.ReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(1024, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
+
+    def forward(self, x):
+        return self.backbone(x)
+
+
+# -------------------------
+# LOAD CHECKPOINT
+# -------------------------
+checkpoint = torch.load("model.pth", map_location=DEVICE)
+
+mlb = checkpoint["mlb"]
+num_classes = len(mlb.classes_)
+
+model = FoodIngredientClassifier(num_classes)
+model.load_state_dict(checkpoint["model_state_dict"])
+model.to(DEVICE)
 model.eval()
 
+threshold = checkpoint.get("optimal_threshold", 0.5)
+
+# -------------------------
+# IMAGE TRANSFORM
+# -------------------------
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
+    transforms.Normalize(
+        [0.485, 0.456, 0.406],
+        [0.229, 0.224, 0.225]
+    )
 ])
 
+# -------------------------
+# PREDICTION FUNCTION
+# -------------------------
 def predict(image):
     image = image.convert("RGB")
     img_tensor = transform(image).unsqueeze(0).to(DEVICE)
@@ -43,6 +90,10 @@ def predict(image):
 
     return {k: v for k, v in results}
 
+
+# -------------------------
+# GRADIO INTERFACE
+# -------------------------
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),