Update app.py

app.py

@@ -1,26 +1,21 @@
-import torch
 import gradio as gr
+import torch
+import torch.nn as nn
+from torchvision import transforms, models
 from PIL import Image
 import numpy as np
-from torchvision import transforms, models
-import torch.nn as nn
-from sklearn.preprocessing import MultiLabelBinarizer  # Add this import
+import pickle
+import os
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# ============================================================================
+# MODEL DEFINITION (must match training code exactly)
+# ============================================================================
 
-# -------------------------
-# 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
-        )
-
+        self.backbone = models.vit_b_16(weights=None)
         num_features = self.backbone.heads.head.in_features
-
         self.backbone.heads = nn.Sequential(
             nn.Dropout(0.5),
             nn.Linear(num_features, 1024),
@@ -38,72 +33,103 @@ class FoodIngredientClassifier(nn.Module):
         return self.backbone(x)
 
 
-# -------------------------
-# LOAD CHECKPOINT - FIXED VERSION
-# -------------------------
-# Add the sklearn class to safe globals before loading
-torch.serialization.add_safe_globals([MultiLabelBinarizer])
+# ============================================================================
+# LOAD MODEL
+# ============================================================================
 
-checkpoint = torch.load("model.pth", map_location=DEVICE)
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+THRESHOLD = 0.5
 
-mlb = checkpoint["mlb"]
-num_classes = len(mlb.classes_)
+def load_model():
+    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()
+    return model, mlb
 
-model = FoodIngredientClassifier(num_classes)
-model.load_state_dict(checkpoint["model_state_dict"])
-model.to(DEVICE)
-model.eval()
+model, mlb = load_model()
 
-threshold = checkpoint.get("optimal_threshold", 0.5)
+# ============================================================================
+# TRANSFORM
+# ============================================================================
 
-# -------------------------
-# 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)
+# ============================================================================
+# INFERENCE FUNCTION
+# ============================================================================
+
+def predict(image, threshold):
+    if image is None:
+        return "Please upload an image."
+
+    img = image.convert("RGB")
+    img_tensor = transform(img).unsqueeze(0).to(DEVICE)
 
     with torch.no_grad():
         output = model(img_tensor)
         probs = torch.sigmoid(output).cpu().numpy()[0]
 
     pred_indices = np.where(probs > threshold)[0]
-    ingredients = mlb.classes_[pred_indices]
-    confidences = probs[pred_indices]
+
+    if len(pred_indices) == 0:
+        return "No ingredients detected above the threshold. Try lowering it."
 
     results = sorted(
-        [(ing, float(conf)) for ing, conf in zip(ingredients, confidences)],
+        zip(mlb.classes_[pred_indices], probs[pred_indices]),
         key=lambda x: x[1],
         reverse=True
     )
 
-    if not results:
-        return {"No ingredient detected": 1.0}
-
-    return {k: v for k, v in results}
-
+    output_lines = ["### 🍽️ Detected Ingredients\n"]
+    for ingredient, confidence in results:
+        bar = "█" * int(confidence * 20)
+        output_lines.append(f"**{ingredient}** — {confidence:.1%}  `{bar}`")
+
+    return "\n\n".join(output_lines)
+
+
+# ============================================================================
+# GRADIO UI
+# ============================================================================
+
+with gr.Blocks(title="Food Ingredient Detector") as demo:
+    gr.Markdown("""
+    # 🥗 Food Ingredient Detector
+    Upload a photo of food and the model will identify its ingredients.
+    Built with a **ViT-B/16** backbone fine-tuned for multi-label ingredient classification.
+    """)
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Food Image")
+            threshold_slider = gr.Slider(
+                minimum=0.1, maximum=0.9, value=0.5, step=0.05,
+                label="Detection Threshold",
+                info="Lower = more ingredients detected, higher = only confident predictions"
+            )
+            predict_btn = gr.Button("🔍 Detect Ingredients", variant="primary")
+
+        with gr.Column():
+            output = gr.Markdown(label="Results")
+
+    predict_btn.click(
+        fn=predict,
+        inputs=[image_input, threshold_slider],
+        outputs=output
+    )
 
-# -------------------------
-# GRADIO INTERFACE
-# -------------------------
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=10),
-    title="Food Ingredient Classifier",
-    description="Upload a food image to detect ingredients."
-)
+    gr.Examples(
+        examples=[],  # Add example image paths here if you include sample images
+        inputs=image_input
+    )
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch()