Update app.py

app.py

@@ -14,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),
@@ -33,14 +33,15 @@ class FoodIngredientClassifier(nn.Module):
             nn.Dropout(0.3),
             nn.Linear(512, num_classes)
         )
-        
+
     def forward(self, x):
         return self.backbone(x)
 
+
 # -------------------------
-# LOAD CHECKPOINT
+# LOAD CHECKPOINT - FIXED VERSION
 # -------------------------
-# Add sklearn MultiLabelBinarizer to safe globals before loading
+# Add the sklearn class to safe globals before loading
 torch.serialization.add_safe_globals([MultiLabelBinarizer])
 
 checkpoint = torch.load("model.pth", map_location=DEVICE)
@@ -73,26 +74,27 @@ 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
 # -------------------------