Update app.py

app.py

@@ -1,67 +1,114 @@
-import json
-import torch
 import gradio as gr
-import numpy as np
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
 from PIL import Image
-from torchvision import transforms
+import numpy as np
 
-from model import FoodIngredientClassifier
+# -----------------------------
+# Device
+# -----------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# -----------------------------
+# Model definition (same as training)
+# -----------------------------
+class FoodIngredientClassifier(nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+        self.backbone = models.resnet50(pretrained=False)
+        num_features = self.backbone.fc.in_features
+        self.backbone.fc = nn.Sequential(
+            nn.Dropout(0.5),
+            nn.Linear(num_features, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, num_classes)
+        )
 
-# Load classes
-with open("mlb_classes.json", "r") as f:
-    classes = json.load(f)
+    def forward(self, x):
+        return self.backbone(x)
 
-NUM_CLASSES = len(classes)
+# -----------------------------
+# Load checkpoint
+# -----------------------------
+checkpoint = torch.load("best_model.pth", map_location=device)
 
-# Load model
-model = FoodIngredientClassifier(NUM_CLASSES)
-checkpoint = torch.load(
-    "best_model.pth",
-    map_location=DEVICE,
-    weights_only=False
-)
+mlb = checkpoint["mlb"]
+class_names = mlb.classes_
+num_classes = len(class_names)
+
+model = FoodIngredientClassifier(num_classes)
 model.load_state_dict(checkpoint["model_state_dict"])
-model.to(DEVICE)
+model.to(device)
 model.eval()
 
-# Image preprocessing
+print(f"Loaded model with {num_classes} ingredient classes")
+
+# -----------------------------
+# Image transforms
+# -----------------------------
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
+        [0.485, 0.456, 0.406],
+        [0.229, 0.224, 0.225]
     )
 ])
 
+# -----------------------------
+# Utility
+# -----------------------------
+def clean_name(name):
+    return name.replace("_", " ").title()
+
+# -----------------------------
+# Prediction function
+# -----------------------------
 def predict(image, threshold=0.5):
+    if image is None:
+        return {}
+
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
     image = image.convert("RGB")
-    x = transform(image).unsqueeze(0).to(DEVICE)
+    input_tensor = transform(image).unsqueeze(0).to(device)
 
     with torch.no_grad():
-        logits = model(x)
+        logits = model(input_tensor)
         probs = torch.sigmoid(logits).cpu().numpy()[0]
 
-    results = [
-        (classes[i], float(probs[i]))
-        for i in np.where(probs > threshold)[0]
-    ]
+    results = {}
+    for idx, prob in enumerate(probs):
+        if prob >= threshold:
+            results[clean_name(class_names[idx])] = float(prob)
+
+    # Sort by confidence
+    results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
 
-    results.sort(key=lambda x: x[1], reverse=True)
-    return results[:10]
+    return results
 
-demo = gr.Interface(
+# -----------------------------
+# Gradio UI
+# -----------------------------
+iface = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(type="pil", label="Upload Food Image"),
-        gr.Slider(0.1, 0.9, value=0.5, label="Confidence Threshold")
+        gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Confidence Threshold")
     ],
     outputs=gr.Label(label="Detected Ingredients"),
-    title="🍽️ Food Ingredient Detector",
-    description="Upload a food image and detect likely ingredients using a ResNet50 multi-label model."
+    title="Food Ingredient Detection (Multi-Label)",
+    description=(
+        "Upload a food image to detect **multiple ingredients at once**. "
+        "This model is trained with multi-label classification using ResNet-50."
+    ),
+    theme=gr.themes.Soft(),
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()