Update app.py

app.py

@@ -4,6 +4,7 @@ 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
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -13,13 +14,13 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 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),
@@ -32,14 +33,16 @@ class FoodIngredientClassifier(nn.Module):
             nn.Dropout(0.3),
             nn.Linear(512, num_classes)
         )
-
+        
     def forward(self, x):
         return self.backbone(x)
 
-
 # -------------------------
 # LOAD CHECKPOINT
 # -------------------------
+# Add sklearn MultiLabelBinarizer to safe globals before loading
+torch.serialization.add_safe_globals([MultiLabelBinarizer])
+
 checkpoint = torch.load("model.pth", map_location=DEVICE)
 
 mlb = checkpoint["mlb"]
@@ -70,27 +73,26 @@ transform = transforms.Compose([
 def predict(image):
     image = image.convert("RGB")
     img_tensor = transform(image).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]
-
+    
     results = sorted(
         [(ing, float(conf)) for ing, conf in zip(ingredients, confidences)],
         key=lambda x: x[1],
         reverse=True
     )
-
+    
     if not results:
         return {"No ingredient detected": 1.0}
-
+    
     return {k: v for k, v in results}
 
-
 # -------------------------
 # GRADIO INTERFACE
 # -------------------------