Update app.py

app.py

@@ -227,26 +227,31 @@ 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."
 
-iface = gr.Interface(
-    fn=inference_pipeline,
-    inputs=[
-        gr.Image(type="numpy", label="Upload chest X-ray (RGB or grayscale)"),
-        gr.Slider(minimum=0.1, maximum=0.9, step=0.01, value=0.5, label="Bacterial threshold (thresh_b)"),
-        gr.Slider(minimum=0.1, maximum=0.9, step=0.01, value=0.5, label="Viral threshold (thresh_v)"),
-        gr.Slider(minimum=0.1, maximum=0.9, step=0.01, value=0.5, label="Segmentation mask threshold (seg_thresh)")
-    ],
-    outputs=[
-        gr.Label(num_top_classes=1, label="Prediction"),
-        gr.Number(label="Bacterial Probability"),
-        gr.Number(label="Viral Probability"),
-        gr.Image(type="pil", label="Masked Image (input × mask)"),
-        gr.Image(type="pil", label="Segmentation Overlay (red mask)")
-    ],
-    title=title,
-    description=desc
-)
-with gr.Blocks() as samples_block:
-    gr.Markdown("### Test Samples")
+with gr.Blocks(title=title) as demo:
+    
+    gr.Markdown(f"## {title}\n{desc}")
+
+    with gr.Row():
+        with gr.Column():
+            img_in = gr.Image(type="numpy", label="Upload chest X-ray")
+            thresh_b = gr.Slider(0.1, 0.9, 0.5, step=0.01, label="Bacterial threshold")
+            thresh_v = gr.Slider(0.1, 0.9, 0.5, step=0.01, label="Viral threshold")
+            seg_thresh = gr.Slider(0.1, 0.9, 0.5, step=0.01, label="Segmentation mask threshold")
+
+        with gr.Column():
+            pred_out = gr.Label(num_top_classes=1, label="Prediction")
+            pb_out = gr.Number(label="Bacterial Probability")
+            pv_out = gr.Number(label="Viral Probability")
+            masked_img_out = gr.Image(type="pil", label="Masked Image")
+            overlay_out = gr.Image(type="pil", label="Segmentation Overlay")
+
+    img_in.change(
+        inference_pipeline,
+        inputs=[img_in, thresh_b, thresh_v, seg_thresh],
+        outputs=[pred_out, pb_out, pv_out, masked_img_out, overlay_out]
+    )
+
+    gr.Markdown("## Test Samples")
 
     with gr.Row():
         with gr.Column():
@@ -266,6 +271,6 @@ with gr.Blocks() as samples_block:
         with gr.Column():
             gr.Markdown("**BACTERIAL**")
 
-
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()
+