Update app.py

app.py

@@ -12,12 +12,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
 # -----------------------------
-# Model definition (same as training)
+# Model definition (MATCHES TRAINING)
 # -----------------------------
 class FoodIngredientClassifier(nn.Module):
     def __init__(self, num_classes):
         super().__init__()
-        self.backbone = models.resnet50(pretrained=False)
+        self.backbone = models.resnet50(weights=None)
         num_features = self.backbone.fc.in_features
         self.backbone.fc = nn.Sequential(
             nn.Dropout(0.5),
@@ -31,7 +31,7 @@ class FoodIngredientClassifier(nn.Module):
         return self.backbone(x)
 
 # -----------------------------
-# Load checkpoint
+# Load checkpoint (PyTorch 2.6+ safe)
 # -----------------------------
 checkpoint = torch.load(
     "best_model.pth",
@@ -39,7 +39,6 @@ checkpoint = torch.load(
     weights_only=False
 )
 
-
 mlb = checkpoint["mlb"]
 class_names = mlb.classes_
 num_classes = len(class_names)
@@ -70,11 +69,11 @@ def clean_name(name):
     return name.replace("_", " ").title()
 
 # -----------------------------
-# Prediction function
+# Prediction function (BULLETPROOF)
 # -----------------------------
-def predict(image, threshold=0.00005):
+def predict(image, threshold):
     if image is None:
-        return {}
+        return {"error": "No image provided"}
 
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image)
@@ -86,10 +85,20 @@ def predict(image, threshold=0.00005):
         logits = model(input_tensor)
         probs = torch.sigmoid(logits).cpu().numpy()[0]
 
-    results = {}
-    for idx, prob in enumerate(probs):
-        if prob >= threshold:
-            results[clean_name(class_names[idx])] = float(prob)
+    # Threshold-based predictions
+    results = {
+        clean_name(class_names[i]): float(probs[i])
+        for i in range(len(probs))
+        if probs[i] >= threshold
+    }
+
+    # 🔒 Fallback: always return top 5
+    if not results:
+        top_idx = np.argsort(probs)[-5:][::-1]
+        results = {
+            clean_name(class_names[i]): float(probs[i])
+            for i in top_idx
+        }
 
     # Sort by confidence
     results = dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
@@ -103,17 +112,23 @@ iface = gr.Interface(
     fn=predict,
     inputs=[
         gr.Image(type="pil", label="Upload Food Image"),
-        gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Confidence Threshold")
+        gr.Slider(
+            minimum=0.00001,
+            maximum=0.5,
+            value=0.05,
+            step=0.01,
+            label="Confidence Threshold"
+        )
     ],
     outputs=gr.JSON(label="Detected Ingredients"),
     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."
+        "Upload a food image to detect **multiple ingredients at once**.\n\n"
+        "This is a **multi-label ResNet-50 model** using sigmoid outputs."
     ),
     theme=gr.themes.Soft(),
     allow_flagging="never"
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch(enable_queue=True)