Update app.py

app.py

@@ -190,13 +190,13 @@ def inference_pipeline(img, thresh_b=0.5, thresh_v=0.5, seg_thresh=0.5):
     red_mask.putalpha(Image.fromarray(alpha))
     blended = Image.alpha_composite(display_orig.convert("RGBA"), red_mask)
     # return values
-    return {
-        "Prediction": label,
-        "Bacterial Probability": float(pb),
-        "Viral Probability": float(pv),
-        "Masked Image": masked_pil,
-        "Segmentation Overlay": blended
-    }
+    return (
+        pred_label,
+        float(prob_bact),
+        float(prob_viral),
+        masked_image,          # PIL image
+        overlay_image          # PIL image
+    )
 
 title = "Chest X-ray: UNet segmentation + 2 binary classifiers"
 desc = "Pipeline: UNet -> mask lungs -> two binary classifiers (Normal vs Bacterial, Normal vs Viral). " \