Update app.py

app.py

@@ -171,33 +171,48 @@ def classify_masked_tensor(masked_img_tensor, thresh_b=0.5, thresh_v=0.5):
 
 def inference_pipeline(img, thresh_b=0.5, thresh_v=0.5, seg_thresh=0.5):
     """
-    Returns: label, bacterial_prob, viral_prob, masked_image (PIL), mask (PIL)
+    Returns: label, bacterial_prob, viral_prob, masked_image (PIL), mask_overlay (PIL)
     """
+
     pil = Image.fromarray(img.astype('uint8'), 'RGB')
-    masked_tensor, mask_np, masked_pil = infer_mask_and_mask_image(pil, threshold=seg_thresh)
-    pb, pv, label = classify_masked_tensor(masked_tensor, thresh_b=thresh_b, thresh_v=thresh_v)
+
+    masked_tensor, mask_np, masked_pil = infer_mask_and_mask_image(
+        pil, threshold=seg_thresh
+    )
+
+    pb, pv, pred_label = classify_masked_tensor(
+        masked_tensor, 
+        thresh_b=thresh_b, 
+        thresh_v=thresh_v
+    )
+
+    # Convert mask to PIL
     mask_vis = (mask_np * 255).astype(np.uint8)
     mask_pil = Image.fromarray(mask_vis).convert("L")
-    display_orig = pil.resize((300,300))
-    overlay = Image.new("RGBA", display_orig.size)
-    overlay.paste(display_orig.convert("RGBA"))
-    # red mask with alpha
-    red_mask = Image.fromarray(np.zeros((300,300,3), dtype=np.uint8))
-    red_mask = Image.fromarray(np.stack([mask_vis, np.zeros_like(mask_vis), np.zeros_like(mask_vis)], axis=2))
-    red_mask = red_mask.convert("RGBA")
-    # apply alpha where mask is 1
+
+    # Resize original for overlay
+    display_orig = pil.resize((300, 300))
+
+    # Create red mask overlay
+    red_mask = np.zeros((300, 300, 3), dtype=np.uint8)
+    red_mask = np.stack([mask_vis, np.zeros_like(mask_vis), np.zeros_like(mask_vis)], axis=2)
+    red_mask = Image.fromarray(red_mask).convert("RGBA")
+
     alpha = (mask_np * 120).astype(np.uint8)
     red_mask.putalpha(Image.fromarray(alpha))
+
     blended = Image.alpha_composite(display_orig.convert("RGBA"), red_mask)
-    # return values
+
+
     return (
-        pred_label,
-        float(prob_bact),
-        float(prob_viral),
-        masked_image,          # PIL image
-        overlay_image          # PIL image
+        pred_label,      
+        float(pb),      
+        float(pv),        
+        masked_pil,        
+        blended            
     )
 
+
 title = "Chest X-ray: UNet segmentation + 2 binary classifiers"
 desc = "Pipeline: UNet -> mask lungs -> two binary classifiers (Normal vs Bacterial, Normal vs Viral). " \
        "If both classifiers fire, the stronger probability is chosen (fallback). Thresholds adjustable."