fix: maybe

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 = self.model(img_tensor)  # Handle both outputs
+            outputs, seg_mask_logits = self.model(img_tensor)  # Handle both outputs
             probabilities = torch.nn.functional.softmax(outputs, dim=1)
 
         probs = probabilities[0].cpu().numpy()
@@ -42,9 +42,11 @@ class WasteClassifier:
         confidence = np.max(probs)
 
         # Process segmentation mask
+        seg_mask_probs = torch.sigmoid(seg_mask_logits)
         seg_mask = (
-            torch.sigmoid(seg_mask)[0, 0].cpu().numpy().astype(np.float32)
-        )  # Get first image, first channel
+            seg_mask_probs[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
 
         # Resize mask back to original image size
@@ -77,7 +79,7 @@ def interface(classifier):
         mask = results["segmentation_mask"]
 
         overlay = image_np.copy()
-        overlay[mask < 0.5] = overlay[mask < 0.5] * 0
+        overlay[mask < 0.2] = overlay[mask < 0.2] * 0
 
         output_str = f"Predicted Class: {results['predicted_class']}\n"
         output_str += f"Confidence: {results['confidence']*100:.2f}%\n\n"