Update app.py

app.py

@@ -3,8 +3,16 @@
 # + SAFER PHONE MODE + MASK POST-PROCESSING + MASK SANITY FAILSAFE
 # + 3-STATE AGREEMENT (LOW / SCREEN+ / TB+)
 #
-# Requirements:
-#   pip install gradio timm torchvision opencv-python pillow transformers einops torch
+# Requirements (requirements.txt):
+#   gradio
+#   torch
+#   torchvision
+#   timm
+#   opencv-python
+#   pillow
+#   transformers
+#   einops
+#   open_clip_torch
 #
 # HF Spaces notes:
 # - weights are expected in ./weights/
@@ -32,19 +40,15 @@ from PIL import Image
 # USER CONFIG (HF Spaces friendly)
 # ============================================================
 
-# ---- Friendly model names for UI ----
 MODEL_NAME_TBNET = "TBNet (CNN model)"
 MODEL_NAME_RADIO = "RADIO (visual model)"
 
-# ---- Default TB/Lung weights ----
 DEFAULT_TB_WEIGHTS = "weights/best.pt"
 DEFAULT_LUNG_WEIGHTS = "weights/lung_unet_mont_shenzhen.pt"
 
-# ---- RADIO config (same env as TB) ----
 RADIO_HF_REPO = "nvidia/C-RADIOv4-SO400M"
 RADIO_REVISION = "c0457f5dc26ca145f954cd4fc5bb6114e5705ad8"
 
-# Your trained heads stored in this Space repo
 RADIO_RAW_HEAD_PATH = "weights/best_raw.pt"
 RADIO_MASKED_HEAD_PATH = "weights/best_masked.pt"
 
@@ -55,17 +59,14 @@ RADIO_THR_RED = 0.23
 RADIO_MASKED_MIN_COV = 0.15
 RADIO_GATE_DEFAULT = 0.21
 
-# ---- Consensus logic thresholds ----
 TBNET_SCREEN_THR = 0.30
 RADIO_SCREEN_THR = RADIO_THR_SCREEN
 
-# ---- Mask fail-safes ----
-FAIL_COV = 0.10     # <10% -> segmentation fail
-WARN_COV = 0.18     # <18% -> warn
-FAILSAFE_ON_BAD_MASK = True  # fail-safe on suspicious/cropped masks
+FAIL_COV = 0.10
+WARN_COV = 0.18
+FAILSAFE_ON_BAD_MASK = True
 
-# ---- Device policy ----
-FORCE_CPU = True  # HF CPU space: keep True
+FORCE_CPU = True
 DEVICE = torch.device("cpu" if FORCE_CPU else ("cuda" if torch.cuda.is_available() else "cpu"))
 
 
@@ -118,6 +119,63 @@ CLINICAL_GUIDANCE = (
 )
 
 
+# ============================================================
+# USER-FRIENDLY SUMMARY BUILDER
+# ============================================================
+def overall_summary(tb_band: Optional[str],
+                    tb_prob: Optional[float],
+                    radio_primary: Optional[float],
+                    radio_band: Optional[str],
+                    consensus_label: str,
+                    q_score: float,
+                    cov: float,
+                    warnings: List[str]) -> str:
+    # Overall label from agreement (keeps it simple for users)
+    if tb_prob is None:
+        overall_title = "INDETERMINATE — NEEDS REVIEW"
+        icon = "⚠️"
+    else:
+        if "AGREE: LOW" in consensus_label:
+            overall_title = "LOW TB LIKELIHOOD"
+            icon = "✅"
+        elif "AGREE: TB+" in consensus_label:
+            overall_title = "TB FEATURES SUSPECTED"
+            icon = "🚩"
+        elif "AGREE: SCREEN+" in consensus_label:
+            overall_title = "SCREEN-POSITIVE — REVIEW RECOMMENDED"
+            icon = "⚠️"
+        else:
+            overall_title = "INDETERMINATE — REVIEW RECOMMENDED"
+            icon = "⚠️"
+
+    reliability = "Good" if (q_score >= 70 and cov >= WARN_COV) else "Limited"
+    rel_icon = "🟢" if reliability == "Good" else "🟡"
+
+    warn_line = "None" if not warnings else f"{len(warnings)} note(s) below"
+
+    tb_prob_str = "N/A" if tb_prob is None else f"{tb_prob:.4f}"
+    radio_str = "N/A" if radio_primary is None else f"{radio_primary:.4f}"
+
+    return f"""
+## {icon} Overall screening result: **{overall_title}**
+
+**Reliability:** {rel_icon} **{reliability}** (Quality: {q_score:.0f}/100 • Lung coverage: {cov*100:.1f}% • Notes: {warn_line})
+
+### What this means
+- This is a **screening support tool**, not a diagnosis.
+- Two models analyze the same image: a **CNN model** (TBNet) and a **visual model** (RADIO).
+
+### Model agreement
+- **{consensus_label}**
+- {MODEL_NAME_TBNET} probability: **{tb_prob_str}**
+- {MODEL_NAME_RADIO} probability: **{radio_str}** {f"(band={radio_band})" if radio_band else ""}
+
+### What to do next
+- If you have symptoms/risk factors, seek clinician/radiologist review.
+- If TB is clinically suspected, consider **CBNAAT/GeneXpert** and sputum testing regardless of AI output.
+"""
+
+
 # ============================================================
 # LUNG U-NET (INFERENCE)
 # ============================================================
@@ -332,12 +390,6 @@ def apply_clahe(gray_u8: np.ndarray) -> np.ndarray:
     return clahe.apply(gray_u8)
 
 def phone_preprocess(gray_u8: np.ndarray) -> np.ndarray:
-    """
-    Safer phone preprocessing:
-    - only crop if border artifacts suggest a framed/screenshot input
-    - only apply CLAHE if underexposed or low sharpness
-    - crop sanity check to avoid destroying clean digital CXRs
-    """
     sharp = laplacian_sharpness(gray_u8)
     lo_clip, _hi_clip = exposure_scores(gray_u8)
     border = border_fraction(gray_u8)
@@ -362,23 +414,21 @@ def cam_entropy(cam: np.ndarray) -> float:
 def detect_diffuse_risk(prob_tb: float, cam_up: np.ndarray, quality_score: float) -> bool:
     if quality_score < 55:
         return False
-    # Only apply diffuse-risk heuristic in near-threshold negatives
     if prob_tb < 0.05:
         return False
     ent = cam_entropy(cam_up)
     return (prob_tb < TBNET_SCREEN_THR) and (ent > 6.5)
 
 def confidence_band(prob_tb: float, quality_score: float, diffuse: bool):
-    # Very low probability should not be forced to YELLOW just because attention is diffuse
     if prob_tb < 0.01 and quality_score >= 45:
-        return ("GREEN", "✅ Very low TB signal detected by the CNN model.")
+        return ("GREEN", "✅ Very low TB signal detected.")
     if quality_score < 55:
-        return ("YELLOW", "⚠️ Image quality is low; treat the result as indeterminate.")
+        return ("YELLOW", "⚠️ Image quality is low; treat as indeterminate.")
     if diffuse:
-        return ("YELLOW", "⚠️ Attention is non-focal; treat the result as indeterminate.")
+        return ("YELLOW", "⚠️ Attention is non-focal; treat as indeterminate.")
     if prob_tb >= TBNET_SCREEN_THR:
         return ("YELLOW", "⚠️ Screening-positive range; review recommended.")
-    return ("GREEN", "✅ No strong TB signal detected by the CNN model.")
+    return ("GREEN", "✅ No strong TB signal detected.")
 
 def make_mask_overlay(gray_u8: np.ndarray, mask_u8: np.ndarray) -> np.ndarray:
     base = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2RGB)
@@ -441,14 +491,14 @@ def mask_sanity_warnings(mask_full_u8: np.ndarray) -> List[str]:
 
 def recommendation_for_band(band: Optional[str]) -> str:
     if band in (None, "YELLOW"):
-        return "✅ Recommendation: Radiologist/clinician review is recommended (result is indeterminate)."
+        return "Radiologist/clinician review is recommended (result is indeterminate)."
     if band == "RED":
-        return "✅ Recommendation: Urgent clinical review + microbiological confirmation (CBNAAT/GeneXpert, sputum)."
-    return "✅ Recommendation: If symptoms/risk factors exist, clinical correlation is advised."
+        return "Urgent clinical review + microbiological confirmation (CBNAAT/GeneXpert, sputum) recommended."
+    return "If symptoms/risk factors exist, clinical correlation is advised."
 
 
 # ============================================================
-# AGREEMENT LOGIC (TBNet vs RADIO) — 3-state
+# AGREEMENT LOGIC (TBNet vs RADIO)
 # ============================================================
 def tbnet_state(tb_prob: float, tb_band: str) -> str:
     if tb_band == "RED":
@@ -475,7 +525,6 @@ def build_consensus(
     if tb_prob is None or tb_band is None:
         return ("N/A", f"{MODEL_NAME_TBNET} not available (lung segmentation failed / fail-safe).")
 
-    # PRIMARY = masked if available else raw
     if radio_masked is not None:
         radio_primary = radio_masked
         radio_used = "MASKED"
@@ -484,12 +533,11 @@ def build_consensus(
         radio_used = "RAW"
 
     if radio_primary is None:
-        return ("TBNet only", f"{MODEL_NAME_RADIO} not available → {MODEL_NAME_TBNET} probability={tb_prob:.4f} (band={tb_band}).")
+        return ("TBNet only", f"{MODEL_NAME_RADIO} not available → {MODEL_NAME_TBNET}={tb_prob:.4f} (band={tb_band}).")
 
     t = tbnet_state(tb_prob, tb_band)
     r = radio_state_from_prob(radio_primary)
-
-    rb = f" (RADIO band={radio_band})" if radio_band else ""
+    rb = f" (band={radio_band})" if radio_band else ""
 
     if t == r:
         return (
@@ -497,11 +545,10 @@ def build_consensus(
             f"Both models indicate **{t}**. {MODEL_NAME_TBNET}={tb_prob:.4f}, {MODEL_NAME_RADIO}({radio_used})={radio_primary:.4f}{rb}."
         )
 
-    # strong disagreement: one says SCREEN+/TB+ and the other says LOW
     if (t in ("SCREEN+", "TB+") and r == "LOW") or (r in ("SCREEN+", "TB+") and t == "LOW"):
         return (
             "DISAGREE",
-            f"Models disagree: {MODEL_NAME_TBNET} suggests **{t}** (band={tb_band}) vs {MODEL_NAME_RADIO} suggests **{r}** ({radio_used})={radio_primary:.4f}{rb}."
+            f"Models disagree: {MODEL_NAME_TBNET} suggests **{t}** vs {MODEL_NAME_RADIO} suggests **{r}** ({radio_used})={radio_primary:.4f}{rb}."
         )
 
     return (
@@ -537,7 +584,6 @@ class ModelBundle:
             load_tb_weights(tb, tb_weights, self.device)
             tb.eval()
             self.tb = tb
-            # EfficientNet in timm has conv_head on effb0
             self.cammer = GradCAM(tb, tb.backbone.conv_head)
             self.tb_path = tb_weights
             self.backbone = backbone
@@ -634,7 +680,6 @@ def radio_overlay_heatmap(rgb_u8: np.ndarray, heatmap01: np.ndarray, alpha: floa
     img = rgb_u8.astype(np.float32) / 255.0
     hm = np.clip(heatmap01, 0, 1).astype(np.float32)
     out = img.copy()
-    # subtle red overlay
     out[..., 0] = np.clip(out[..., 0] * (1 - alpha) + hm * alpha, 0, 1)
     return (out * 255).astype(np.uint8)
 
@@ -647,7 +692,6 @@ def radio_predict_from_arrays(gray_vis_u8: np.ndarray,
     RADIO_BUNDLE.load(device=device)
     dtype = torch.float16 if device.type == "cuda" else torch.float32
 
-    # ---------- RAW ----------
     raw_rgb = cv2.cvtColor(gray_vis_u8, cv2.COLOR_GRAY2RGB)
     px = RADIO_BUNDLE.processor(
         images=Image.fromarray(raw_rgb),
@@ -668,7 +712,6 @@ def radio_predict_from_arrays(gray_vis_u8: np.ndarray,
         alpha=0.35
     )
 
-    # ---------- MASKED (optional) ----------
     masked_prob = None
     masked_overlay = None
     masked_ran = False
@@ -697,7 +740,6 @@ def radio_predict_from_arrays(gray_vis_u8: np.ndarray,
             alpha=0.35
         )
 
-    # ---------- PRIMARY = masked if available else raw ----------
     prob_primary = masked_prob if masked_prob is not None else prob_raw
 
     if prob_primary >= RADIO_THR_RED:
@@ -731,7 +773,7 @@ def analyze_one_image(
     tb_weights: str,
     lung_weights: str,
     backbone: str,
-    threshold: float,  # kept for UI compatibility; main logic uses TBNET_SCREEN_THR
+    threshold: float,
     phone_mode: bool,
     img_size: int = 224,
     fail_cov: float = FAIL_COV,
@@ -754,8 +796,6 @@ def analyze_one_image(
         mask256 = torch.sigmoid(mask_logits)[0, 0].cpu().numpy()
 
     mask256_bin = (mask256 > 0.5).astype(np.uint8)
-
-    # post-process: keep 2 lungs, close, fill holes
     mask256_bin = keep_top_k_components(mask256_bin, k=2)
     k = max(3, int(0.02 * 256))
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
@@ -765,7 +805,6 @@ def analyze_one_image(
     coverage = float(mask256_bin.mean())
     mask_full = cv2.resize(mask256_bin, (gray_vis.shape[1], gray_vis.shape[0]), interpolation=cv2.INTER_NEAREST)
 
-    # fail-safe if coverage too low
     if coverage < fail_cov:
         overlay_rgb = cv2.cvtColor(cv2.resize(gray_vis, (img_size, img_size)), cv2.COLOR_GRAY2RGB)
         return {
@@ -774,13 +813,13 @@ def analyze_one_image(
             "pred": "INDETERMINATE",
             "band": "YELLOW",
             "band_text": (
-                "⚠️ The lung segmentation looks unreliable, so the TB screening score is disabled for safety.\n\n"
-                "Please use a clearer frontal chest X-ray (PA/AP) or seek radiologist review."
+                "⚠️ Lung segmentation looks unreliable, so the TBNet screening score is disabled for safety.\n\n"
+                "Please use a clearer standard frontal CXR (PA/AP) or seek radiologist review."
             ),
             "quality_score": float(q_score),
             "diffuse_risk": False,
             "warnings": (
-                [f"Lung segmentation coverage is too small ({coverage*100:.1f}%)."]
+                [f"Lung segmentation coverage too small ({coverage*100:.1f}%)."]
                 + (["Phone/WhatsApp mode enabled; artifacts possible."] if phone_mode else [])
                 + q_warn
             ),
@@ -795,7 +834,6 @@ def analyze_one_image(
             "overlay_clean": overlay_rgb,
         }
 
-    # fail-safe if mask looks like single lung / cropped
     sanity = mask_sanity_warnings(mask_full.astype(np.uint8))
     if FAILSAFE_ON_BAD_MASK and sanity:
         overlay_rgb = cv2.cvtColor(cv2.resize(gray_vis, (img_size, img_size)), cv2.COLOR_GRAY2RGB)
@@ -806,8 +844,8 @@ def analyze_one_image(
             "band": "YELLOW",
             "band_text": (
                 "⚠️ The image appears cropped/non-standard (mask sanity check). "
-                "TB screening score is disabled for safety.\n\n"
-                "Please use a standard frontal chest X-ray (PA/AP) or seek radiologist review."
+                "TBNet screening score is disabled for safety.\n\n"
+                "Please use a standard frontal CXR (PA/AP) or seek radiologist review."
             ),
             "quality_score": float(q_score),
             "diffuse_risk": False,
@@ -843,7 +881,7 @@ def analyze_one_image(
     cam_up = cam_u8.astype(np.float32) / 255.0
 
     diffuse = detect_diffuse_risk(prob_tb, cam_up, q_score)
-    band_base, _band_hint = confidence_band(prob_tb, q_score, diffuse)
+    band_base, _ = confidence_band(prob_tb, q_score, diffuse)
 
     allow_red = (prob_tb >= 0.70 and q_score >= 55 and (not diffuse) and coverage >= warn_cov)
     band = "RED" if allow_red else band_base
@@ -851,14 +889,12 @@ def analyze_one_image(
     pred = REPORT_LABELS[band]["title"]
     band_text = REPORT_LABELS[band]["summary"]
 
-    # Extra helpful line if YELLOW but probability is very low
     if band == "YELLOW" and prob_tb < 0.05:
         band_text = (
             "⚠️ TB probability is very low, but the result is marked **indeterminate** because reliability is limited.\n\n"
             + band_text
         )
 
-    # Build Grad-CAM overlay
     heat = cv2.applyColorMap((cam_up * 255).astype(np.uint8), cv2.COLORMAP_JET)
     overlay_clean = cv2.addWeighted(rgb, 0.65, heat, 0.35, 0)
 
@@ -946,14 +982,13 @@ def run_analysis(
             img_size=224,
         )
 
-        # -------------------------
         # RADIO (optional)
-        # -------------------------
         radio_text = f"{MODEL_NAME_RADIO} is disabled."
         radio_raw_overlay = None
         radio_masked_overlay = None
         radio_raw_val: Optional[float] = None
         radio_masked_val: Optional[float] = None
+        radio_primary_val: Optional[float] = None
         radio_band: Optional[str] = None
 
         radio_raw_str = ""
@@ -968,8 +1003,8 @@ def run_analysis(
                     device=BUNDLE.device,
                     gate_threshold=float(radio_gate),
                 )
-
                 radio_raw_val = float(r["prob_raw"])
+                radio_primary_val = float(r["prob_primary"])
                 radio_masked_val = None if r["masked_prob"] is None else float(r["masked_prob"])
                 radio_band = str(r["band"])
 
@@ -977,7 +1012,8 @@ def run_analysis(
                 radio_masked_str = "" if radio_masked_val is None else f"{radio_masked_val:.4f}"
 
                 radio_text = (
-                    f"**{MODEL_NAME_RADIO} result:** {r['pred']} | RAW={radio_raw_val:.4f}"
+                    f"**{MODEL_NAME_RADIO} result:** {r['pred']} | "
+                    f"PRIMARY={radio_primary_val:.4f} | RAW={radio_raw_val:.4f}"
                     + (f" | MASKED={radio_masked_val:.4f}" if radio_masked_val is not None else "")
                     + f" | Band={radio_band}"
                 )
@@ -985,15 +1021,11 @@ def run_analysis(
                 radio_masked_overlay = r["masked_overlay"]
             except Exception as e:
                 radio_text = f"{MODEL_NAME_RADIO} error: {type(e).__name__}: {e}"
-                radio_raw_str = ""
-                radio_masked_str = ""
                 radio_raw_val = None
                 radio_masked_val = None
+                radio_primary_val = None
                 radio_band = None
 
-        # -------------------------
-        # Agreement between models
-        # -------------------------
         consensus_label, consensus_detail = build_consensus(
             tb_prob=out["prob"],
             tb_band=out["band"],
@@ -1002,9 +1034,7 @@ def run_analysis(
             radio_band=radio_band,
         )
 
-        # -------------------------
         # Table row
-        # -------------------------
         prob_str = "" if out["prob"] is None else f"{out['prob']:.4f}"
         cov_str = f"{out.get('lung_coverage', 0.0) * 100:.1f}%"
 
@@ -1021,9 +1051,7 @@ def run_analysis(
             consensus_label,
         ])
 
-        # -------------------------
-        # Visual outputs
-        # -------------------------
+        # Gallery
         orig_rgb = cv2.cvtColor(cv2.resize(out["orig_gray"], (512, 512)), cv2.COLOR_GRAY2RGB)
         vis_rgb = cv2.cvtColor(cv2.resize(out["vis_gray"], (512, 512)), cv2.COLOR_GRAY2RGB)
         mask_overlay = cv2.resize(out["mask_overlay"], (512, 512))
@@ -1044,36 +1072,56 @@ def run_analysis(
         if radio_masked_overlay is not None:
             gallery_items.append((cv2.resize(radio_masked_overlay, (512, 512)), f"{name} • RADIO MASKED heatmap"))
 
-        # -------------------------
-        # Details panel (user-friendly)
-        # -------------------------
+        # Details (dashboard style)
+        summary_md = overall_summary(
+            tb_band=out.get("band"),
+            tb_prob=out.get("prob"),
+            radio_primary=radio_primary_val,
+            radio_band=radio_band,
+            consensus_label=consensus_label,
+            q_score=float(out["quality_score"]),
+            cov=float(out.get("lung_coverage", 0.0)),
+            warnings=out.get("warnings", []),
+        )
+
         warn_txt = "\n".join([f"- {w}" for w in out["warnings"]]) if out["warnings"] else "- None"
         tb_line = "N/A (disabled by fail-safe)" if out["prob"] is None else f"{out['prob']:.4f}"
         rec_line = recommendation_for_band(out.get("band"))
 
         details_md.append(
-            f"""### {name}
+            f"""{summary_md}
+
+---
+
+<details>
+<summary><b>{MODEL_NAME_TBNET} details</b></summary>
+
+- **Result:** {out['pred']} ({out['band']})
+- **Recommendation:** {rec_line}
+- **Probability (screening score):** {tb_line}
+- **Attention pattern:** {"Diffuse / non-focal" if out["diffuse_risk"] else "Focal / localized"}
 
-**{MODEL_NAME_TBNET} result:** **{out['pred']}**
-{rec_line}
+</details>
 
-**What this means**
-{out['band_text']}
+<details>
+<summary><b>{MODEL_NAME_RADIO} details</b></summary>
 
-**Why it decided this**
-- {MODEL_NAME_TBNET} probability: {tb_line}
-- Image quality: {out['quality_score']:.0f}/100
-- Lung mask coverage: {out.get('lung_coverage', 0.0) * 100:.1f}%
-- Attention pattern (TBNet): {"Diffuse / non-focal" if out["diffuse_risk"] else "Focal / localized"}
+{radio_text}
+
+</details>
+
+<details>
+<summary><b>Image quality & segmentation</b></summary>
+
+- **Quality score:** {out['quality_score']:.0f}/100
+- **Lung mask coverage:** {out.get('lung_coverage', 0.0) * 100:.1f}%
 
 **Notes that may affect reliability**
 {warn_txt}
 
-**{MODEL_NAME_RADIO} output**
-{radio_text}
+</details>
 
-**Agreement between models (TBNet vs RADIO):** **{consensus_label}**
-- {consensus_detail}
+---
 
 **Clinical guidance**
 {CLINICAL_GUIDANCE}
@@ -1096,10 +1144,10 @@ def build_ui():
     """
 
     with gr.Blocks(title="TB X-ray Assistant (TBNet + RADIO)", css=css) as demo:
-        gr.Markdown('<div class="title">TB X-ray Assistant (Auto Lung Mask • Research Use)</div>')
+        gr.Markdown('<div class="title">TB X-ray Assistant (Research Use)</div>')
         gr.Markdown(
-            f"<div class='subtitle'>Lung U-Net masking → <b>{MODEL_NAME_TBNET}</b> + Grad-CAM • "
-            f"Optional <b>{MODEL_NAME_RADIO}</b> (C-RADIOv4 + heads) • Agreement summary</div>"
+            f"<div class='subtitle'>Auto lung mask → <b>{MODEL_NAME_TBNET}</b> + Grad-CAM • "
+            f"Optional <b>{MODEL_NAME_RADIO}</b> (C-RADIOv4 + heads) • User-friendly summary</div>"
         )
 
         with gr.Row():
@@ -1121,7 +1169,6 @@ def build_ui():
                     label="Phone/WhatsApp Mode (safe: conditional crop + conditional CLAHE)"
                 )
 
-                # RADIO
                 use_radio = gr.Checkbox(value=False, label=f"Enable {MODEL_NAME_RADIO}")
                 radio_gate = gr.Slider(
                     0.10, 0.40, value=RADIO_GATE_DEFAULT, step=0.01,