fix: revert fix what

app.py

@@ -34,7 +34,7 @@ class WasteClassifier:
         img_tensor = self.transform(image).unsqueeze(0).to(self.device)
 
         with torch.no_grad():
-            outputs, seg_mask_logits = self.model(img_tensor)  # Handle both outputs
+            outputs, seg_mask = self.model(img_tensor)  # Handle both outputs
             probabilities = torch.nn.functional.softmax(outputs, dim=1)
 
         probs = probabilities[0].cpu().numpy()
@@ -42,9 +42,8 @@ class WasteClassifier:
         confidence = np.max(probs)
 
         # Process segmentation mask
-        seg_mask_probs = torch.sigmoid(seg_mask_logits)
         seg_mask = (
-            seg_mask_probs[0, 0].cpu().numpy().astype(np.float32)
+            seg_mask[0, 0].cpu().numpy().astype(np.float32)
         )
         # Get first image, first channel
         # seg_mask = (seg_mask >= 0.2).astype(np.float32)  # Threshold at 0.2