Update app.py

app.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
+import html
 import traceback
 
 import gradio as gr
-import matplotlib.pyplot as plt
 import pandas as pd
 from PIL import Image
 
@@ -11,144 +11,261 @@ from src.config import CLASS_DISPLAY_NAMES, CLASS_NAMES, ENSEMBLE_MEMBERS
 from src.modeling import diagnose_checkpoints, load_ensemble, predict, weighted_ensemble_cam
 
 
-DISCLAIMER = """
-<div class="notice">
-  <b>Research prototype only.</b> This Space is for demonstrating a trained MRI image classifier from the submitted research workflow. It is not a medical device and must not be used for diagnosis, treatment, or patient triage.
-</div>
-"""
-
-
 CUSTOM_CSS = """
 :root {
+  --bg-1: #eef2ff;
+  --bg-2: #f8fafc;
+  --ink: #101936;
+  --muted: #53617d;
+  --line: rgba(148,163,184,.28);
+  --card: rgba(255,255,255,.84);
+  --primary: #5b4ff5;
+  --primary-2: #6d5dfc;
+  --shadow: 0 18px 54px rgba(15, 23, 42, .13);
   --radius-xl: 26px;
   --radius-lg: 18px;
-  --glass: rgba(255,255,255,0.80);
-  --stroke: rgba(148,163,184,0.28);
-  --shadow: 0 18px 60px rgba(15, 23, 42, .16);
 }
+
 .gradio-container {
-  max-width: 1220px !important;
+  max-width: 1240px !important;
   margin: auto !important;
   background:
-    radial-gradient(circle at 10% 8%, rgba(99,102,241,.18), transparent 30%),
-    radial-gradient(circle at 90% 0%, rgba(14,165,233,.17), transparent 28%),
-    linear-gradient(180deg, #f8fafc 0%, #eef2ff 52%, #f8fafc 100%) !important;
+    radial-gradient(circle at 8% 3%, rgba(99,102,241,.18), transparent 30%),
+    radial-gradient(circle at 92% 0%, rgba(14,165,233,.16), transparent 28%),
+    linear-gradient(180deg, var(--bg-1) 0%, var(--bg-2) 56%, #f8fafc 100%) !important;
   font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif !important;
+  color: var(--ink) !important;
 }
-.hero {
-  padding: 34px 36px;
+
+.hero-card {
+  padding: 38px 38px 34px;
   border-radius: var(--radius-xl);
-  background: linear-gradient(135deg, rgba(15,23,42,.96), rgba(30,64,175,.88) 52%, rgba(2,132,199,.84));
-  color: white;
+  background: rgba(255,255,255,.88);
+  border: 1px solid rgba(219,226,238,.95);
   box-shadow: var(--shadow);
-  border: 1px solid rgba(255,255,255,.16);
-  margin-bottom: 20px;
+  backdrop-filter: blur(16px);
+  margin: 10px 0 24px;
 }
-.hero h1 {
-  margin: 0 0 10px 0 !important;
-  font-size: clamp(2rem, 5vw, 3.8rem) !important;
+
+.hero-card h1 {
+  margin: 0 !important;
+  color: var(--ink);
+  font-size: clamp(2.5rem, 5vw, 4.15rem) !important;
   line-height: 1.02 !important;
-  letter-spacing: -0.05em !important;
+  letter-spacing: -0.055em !important;
+  font-weight: 900 !important;
 }
-.hero p {
-  font-size: 1.02rem;
-  color: rgba(255,255,255,.86);
-  max-width: 880px;
-  margin: 0;
+
+.hero-card h2 {
+  margin: 8px 0 0 !important;
+  color: #273250;
+  font-size: clamp(1.45rem, 2.5vw, 2rem) !important;
+  line-height: 1.18 !important;
+  font-weight: 780 !important;
+  letter-spacing: -0.03em !important;
 }
-.badges {
+
+.class-chip-row {
   display: flex;
   flex-wrap: wrap;
-  gap: 10px;
-  margin-top: 20px;
-}
-.badge {
-  border: 1px solid rgba(255,255,255,.22);
-  background: rgba(255,255,255,.12);
-  color: #fff;
-  padding: 8px 12px;
-  border-radius: 999px;
-  font-weight: 650;
-  font-size: .88rem;
+  gap: 12px;
+  margin-top: 30px;
+  padding-top: 24px;
+  border-top: 1px solid rgba(148,163,184,.35);
 }
-.notice {
-  padding: 14px 16px;
+
+.class-chip {
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  min-height: 46px;
+  padding: 0 22px;
   border-radius: 16px;
-  background: rgba(254, 243, 199, .82);
-  border: 1px solid rgba(245, 158, 11, .34);
-  color: #78350f;
-  margin: 10px 0 18px 0;
+  background: rgba(255,255,255,.72);
+  border: 1px solid rgba(148,163,184,.32);
+  color: #273250;
+  font-size: .98rem;
+  font-weight: 760;
+  box-shadow: 0 8px 22px rgba(15,23,42,.04);
+}
+
+.glass-card, .result-card, .prob-card {
+  border-radius: var(--radius-xl) !important;
+  background: var(--card) !important;
+  border: 1px solid rgba(219,226,238,.95) !important;
+  box-shadow: var(--shadow) !important;
+  backdrop-filter: blur(16px);
+}
+
+.glass-card {
+  padding: 24px !important;
+}
+
+.card-title {
+  margin: 0 0 18px;
+  font-size: 1.25rem;
+  font-weight: 850;
+  letter-spacing: -.025em;
+  color: var(--ink);
 }
+
 .result-card {
-  padding: 22px;
-  border-radius: var(--radius-xl);
-  background: var(--glass);
-  border: 1px solid var(--stroke);
-  box-shadow: 0 14px 44px rgba(15,23,42,.11);
-  backdrop-filter: blur(12px);
+  padding: 30px 34px;
+  margin-bottom: 24px;
 }
+
 .pred-title {
-  font-size: 1.05rem;
-  text-transform: uppercase;
-  letter-spacing: .08em;
-  color: #475569;
-  font-weight: 800;
-  margin-bottom: 4px;
+  font-size: 1.02rem;
+  color: var(--ink);
+  font-weight: 780;
+  margin-bottom: 24px;
 }
+
 .pred-label {
-  font-size: clamp(2rem, 4vw, 3.4rem);
-  line-height: 1;
-  font-weight: 900;
-  letter-spacing: -0.045em;
-  color: #0f172a;
+  font-size: clamp(2.7rem, 4.6vw, 4rem);
+  line-height: .98;
+  font-weight: 920;
+  letter-spacing: -0.055em;
+  color: var(--ink);
 }
+
 .pred-sub {
-  margin-top: 12px;
-  color: #475569;
-  font-size: 1rem;
+  margin-top: 22px;
+  color: var(--muted);
+  font-size: 1.08rem;
+}
+
+.pred-sub b {
+  color: var(--primary);
 }
+
 .metric-grid {
   display: grid;
   grid-template-columns: repeat(3, minmax(0, 1fr));
-  gap: 12px;
-  margin-top: 18px;
+  gap: 16px;
+  margin-top: 30px;
 }
+
 .metric {
-  padding: 14px;
+  min-height: 84px;
+  padding: 18px 14px;
   border-radius: 16px;
-  background: rgba(255,255,255,.76);
-  border: 1px solid rgba(148,163,184,.28);
+  background: rgba(255,255,255,.72);
+  border: 1px solid rgba(148,163,184,.26);
+  text-align: center;
 }
-.metric .k { color: #64748b; font-size: .8rem; font-weight: 800; text-transform: uppercase; letter-spacing: .07em; }
-.metric .v { color: #0f172a; font-size: 1.35rem; font-weight: 900; margin-top: 4px; }
-.status-good { color: #166534; font-weight: 800; }
-.status-bad { color: #991b1b; font-weight: 800; }
-.footer-note { color: #64748b; font-size: .92rem; }
-button.primary, .primary button {
+
+.metric .k {
+  color: #34415f;
+  font-size: .93rem;
+  font-weight: 650;
+}
+
+.metric .v {
+  color: var(--ink);
+  font-size: 1.28rem;
+  font-weight: 850;
+  margin-top: 10px;
+}
+
+.metric .v.confidence {
+  color: var(--primary);
+}
+
+.prob-card {
+  padding: 30px 34px;
+}
+
+.prob-title {
+  margin: 0 0 22px;
+  color: var(--ink);
+  font-size: 1.25rem;
+  font-weight: 850;
+  letter-spacing: -.025em;
+}
+
+.prob-item {
+  padding: 20px 28px 24px;
+  border-radius: 18px;
+  background: rgba(255,255,255,.74);
+  border: 1px solid rgba(148,163,184,.24);
+  margin-bottom: 22px;
+}
+
+.prob-head {
+  display: flex;
+  justify-content: space-between;
+  align-items: baseline;
+  gap: 20px;
+  margin-bottom: 18px;
+}
+
+.prob-label {
+  color: var(--ink);
+  font-size: 1.08rem;
+  font-weight: 820;
+}
+
+.prob-percent {
+  color: var(--ink);
+  font-size: 1.03rem;
+  font-weight: 820;
+}
+
+.prob-item.top .prob-percent {
+  color: var(--primary);
+}
+
+.prob-track {
+  height: 10px;
+  border-radius: 999px;
+  background: #e8edf5;
+  overflow: hidden;
+}
+
+.prob-fill {
+  height: 100%;
+  border-radius: 999px;
+  background: linear-gradient(90deg, var(--primary), var(--primary-2));
+}
+
+#run_button button {
   border-radius: 16px !important;
+  min-height: 48px !important;
   font-weight: 850 !important;
+  background: linear-gradient(135deg, var(--primary), var(--primary-2)) !important;
 }
-.block, .form, .panel, .tabitem, .gr-box {
+
+.footer-note {
+  color: #64748b;
+  font-size: .92rem;
+  margin-top: 18px;
+}
+
+.status-good { color: #166534; font-weight: 800; }
+.status-bad { color: #991b1b; font-weight: 800; }
+
+.block, .form, .panel, .tabitem, .gr-box, .gradio-container .wrap {
   border-radius: 18px !important;
 }
-@media (max-width: 800px) {
-  .hero { padding: 24px 20px; }
+
+@media (max-width: 900px) {
+  .hero-card { padding: 26px 20px; }
   .metric-grid { grid-template-columns: 1fr; }
+  .prob-card, .result-card { padding: 24px 20px; }
 }
 """
 
 
-HERO_HTML = """
-<div class="hero">
-  <h1>Brain MRI Ensemble Classifier</h1>
-  <p>Upload a brain MRI image and run the selected lightweight weighted-soft ensemble from the research notebook: EfficientNet-B0 + MobileNetV3-Small.</p>
-  <div class="badges">
-    <span class="badge">4 classes</span>
-    <span class="badge">Weighted soft voting</span>
-    <span class="badge">EfficientNet-B0</span>
-    <span class="badge">MobileNetV3-Small</span>
-    <span class="badge">Optional heatmap</span>
-  </div>
+CLASS_CHIPS_HTML = "".join(
+    f"<span class='class-chip'>{html.escape(CLASS_DISPLAY_NAMES[name])}</span>" for name in CLASS_NAMES
+)
+
+HERO_HTML = f"""
+<div class="hero-card">
+  <h1>LCVC DeepFuse</h1>
+  <h2>Brain MRI Ensemble Classifier</h2>
+  <div class="class-chip-row">{CLASS_CHIPS_HTML}</div>
 </div>
 """
 
@@ -156,7 +273,7 @@ HERO_HTML = """
 def _status_markdown() -> str:
     ok, _df, message = diagnose_checkpoints()
     cls = "status-good" if ok else "status-bad"
-    return f"<div class='{cls}'>{message}</div>"
+    return f"<div class='{cls}'>{html.escape(message).replace(chr(10), '<br>')}</div>"
 
 
 def _model_table() -> pd.DataFrame:
@@ -178,64 +295,101 @@ def _research_metrics_table() -> pd.DataFrame:
 
 
 def _deployed_members_table() -> pd.DataFrame:
-    rows = []
-    for m in ENSEMBLE_MEMBERS:
-        rows.append(
+    return pd.DataFrame(
+        [
             {
                 "member": m["display_name"],
                 "weight": f"{m['weight']:.8f}",
                 "checkpoint": f"models/{m['checkpoint_file']}",
             }
-        )
-    return pd.DataFrame(rows)
+            for m in ENSEMBLE_MEMBERS
+        ]
+    )
+
 
+def _pct(value: float) -> str:
+    return f"{100.0 * float(value):.2f}%"
 
-def _probability_plot(prob_df: pd.DataFrame):
-    fig, ax = plt.subplots(figsize=(7.5, 4.2))
-    labels = prob_df["class"].tolist()[::-1]
-    values = prob_df["probability"].tolist()[::-1]
-    ax.barh(labels, values)
-    ax.set_xlim(0, 1)
-    ax.set_xlabel("Probability")
-    ax.set_title("Ensemble class probabilities")
-    ax.grid(axis="x", alpha=0.28)
-    for idx, value in enumerate(values):
-        ax.text(min(value + 0.015, 0.98), idx, f"{value*100:.1f}%", va="center", fontweight="bold")
-    fig.tight_layout()
-    return fig
+
+def _empty_prediction_card() -> str:
+    return """
+    <div class="result-card">
+      <div class="pred-title">Top Prediction</div>
+      <div class="pred-label">Waiting</div>
+      <div class="pred-sub">Upload an MRI image, then run prediction.</div>
+      <div class="metric-grid">
+        <div class="metric"><div class="k">Input size</div><div class="v">—</div></div>
+        <div class="metric"><div class="k">Model votes</div><div class="v">3</div></div>
+        <div class="metric"><div class="k">Confidence</div><div class="v confidence">—</div></div>
+      </div>
+    </div>
+    """
 
 
-def _prediction_card(label: str, confidence: float, image: Image.Image, heatmap_enabled: bool) -> str:
+def _prediction_card(label: str, confidence: float, image: Image.Image) -> str:
     width, height = image.size if image is not None else (0, 0)
-    heatmap_text = "Generated" if heatmap_enabled else "Off"
+    label = html.escape(label)
+    confidence_text = _pct(confidence)
     return f"""
     <div class="result-card">
-      <div class="pred-title">Top prediction</div>
+      <div class="pred-title">Top Prediction</div>
       <div class="pred-label">{label}</div>
-      <div class="pred-sub">Weighted-soft ensemble confidence: <b>{confidence*100:.2f}%</b></div>
+      <div class="pred-sub">Confidence: <b>{confidence_text}</b></div>
       <div class="metric-grid">
-        <div class="metric"><div class="k">Input size</div><div class="v">{width}×{height}</div></div>
+        <div class="metric"><div class="k">Input size</div><div class="v">{width} × {height}</div></div>
         <div class="metric"><div class="k">Model votes</div><div class="v">3</div></div>
-        <div class="metric"><div class="k">Heatmap</div><div class="v">{heatmap_text}</div></div>
+        <div class="metric"><div class="k">Confidence</div><div class="v confidence">{confidence_text}</div></div>
       </div>
     </div>
     """
 
 
+def _probabilities_card(probabilities: dict[str, float] | None = None, top_class: str | None = None) -> str:
+    probabilities = probabilities or {name: 0.0 for name in CLASS_NAMES}
+
+    rows: list[tuple[str, str, float]] = []
+    for class_name in CLASS_NAMES:
+        display = CLASS_DISPLAY_NAMES[class_name]
+        probability = float(probabilities.get(class_name, 0.0))
+        rows.append((class_name, display, probability))
+
+    # Highest probability first, but labels are always the real four class names.
+    rows.sort(key=lambda item: item[2], reverse=True)
+
+    items = []
+    for class_name, display, probability in rows:
+        percent = max(0.0, min(100.0, probability * 100.0))
+        top_class_name = " top" if class_name == top_class else ""
+        items.append(
+            f"""
+            <div class="prob-item{top_class_name}">
+              <div class="prob-head">
+                <span class="prob-label">{html.escape(display)}</span>
+                <span class="prob-percent">{percent:.2f}%</span>
+              </div>
+              <div class="prob-track"><div class="prob-fill" style="width:{percent:.4f}%"></div></div>
+            </div>
+            """
+        )
+
+    return f"""
+    <div class="prob-card">
+      <div class="prob-title">Class Probabilities</div>
+      {''.join(items)}
+    </div>
+    """
+
+
 def run_prediction(image: Image.Image, make_heatmap: bool):
     if image is None:
         raise gr.Error("Upload an MRI image first.")
 
     try:
         result = predict(image)
-        prob_df = result.probability_df.copy()
-        prob_df["probability"] = prob_df["probability"].map(lambda x: round(float(x), 6))
-        plot = _probability_plot(result.probability_df)
-        heatmap = None
-        if make_heatmap:
-            heatmap = weighted_ensemble_cam(image, result.predicted_class)
-        card = _prediction_card(result.predicted_display, result.confidence, image, make_heatmap and heatmap is not None)
-        return card, prob_df, result.member_df, plot, heatmap
+        heatmap = weighted_ensemble_cam(image, result.predicted_class) if make_heatmap else None
+        prediction_html = _prediction_card(result.predicted_display, result.confidence, image)
+        probabilities_html = _probabilities_card(result.probabilities, result.predicted_class)
+        return prediction_html, probabilities_html, heatmap
     except FileNotFoundError as exc:
         raise gr.Error(str(exc)) from exc
     except Exception as exc:
@@ -248,90 +402,75 @@ def warmup_status() -> str:
     if not ok:
         return message
     try:
-        # Load once so the first user prediction is faster.
         load_ensemble()
         return "✅ Checkpoints found and ensemble loaded successfully."
     except Exception as exc:
         return f"❌ Checkpoints were found, but model loading failed: {exc}"
 
 
-with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="sky"), title="MRI Ensemble Classifier") as demo:
+with gr.Blocks(
+    css=CUSTOM_CSS,
+    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="sky"),
+    title="LCVC DeepFuse",
+) as demo:
     gr.HTML(HERO_HTML)
-    gr.HTML(DISCLAIMER)
-
-    with gr.Tabs():
-        with gr.Tab("Predict"):
-            with gr.Row(equal_height=False):
-                with gr.Column(scale=5):
-                    image_input = gr.Image(
-                        label="Upload MRI image",
-                        type="pil",
-                        height=390,
-                        sources=["upload", "clipboard"],
-                    )
-                    with gr.Row():
-                        heatmap_toggle = gr.Checkbox(
-                            value=False,
-                            label="Generate ensemble heatmap",
-                            info="Slower on CPU. Use after the basic prediction works.",
-                        )
-                        run_button = gr.Button("Run ensemble prediction", variant="primary", elem_classes="primary")
-                with gr.Column(scale=7):
-                    prediction_html = gr.HTML(
-                        "<div class='result-card'><div class='pred-title'>Waiting for image</div><div class='pred-label'>—</div><div class='pred-sub'>Upload an MRI image and run the selected ensemble.</div></div>",
-                        label="Prediction",
-                    )
-                    probabilities_output = gr.Dataframe(
-                        label="Class probabilities",
-                        headers=["class", "probability", "percent"],
-                        interactive=False,
-                    )
-            with gr.Row():
-                probability_plot = gr.Plot(label="Probability chart")
-                heatmap_output = gr.Image(label="Optional ensemble heatmap", type="pil")
-            member_output = gr.Dataframe(label="Per-member predictions", interactive=False)
-
-            run_button.click(
-                fn=run_prediction,
-                inputs=[image_input, heatmap_toggle],
-                outputs=[prediction_html, probabilities_output, member_output, probability_plot, heatmap_output],
-            )
-
-        with gr.Tab("Model status"):
-            gr.Markdown("### Checkpoint status")
-            status_md = gr.Markdown(_status_markdown())
-            status_table = gr.Dataframe(value=_model_table(), interactive=False, label="Required files")
-            refresh_btn = gr.Button("Refresh status")
-            load_btn = gr.Button("Test-load ensemble", variant="secondary")
-            load_status = gr.Textbox(label="Load result", interactive=False)
-            refresh_btn.click(fn=_status_markdown, inputs=None, outputs=status_md)
-            refresh_btn.click(fn=_model_table, inputs=None, outputs=status_table)
-            load_btn.click(fn=warmup_status, inputs=None, outputs=load_status)
-
-        with gr.Tab("Research summary"):
-            gr.Markdown(
-                """
-### Selected ensemble
-
-The deployed model is the selected **`lightweight_effnet_mobilenet | optimized_val_ce_weighted_soft`** ensemble. It uses only the non-zero-weight members from the ablation result. The zero-weight EfficientNet/MobileNet members are not loaded because they do not affect weighted-soft inference.
-
-Class order used at inference: **glioma, meningioma, notumor, pituitary**.
-                """
-            )
-            gr.Dataframe(value=_deployed_members_table(), label="Deployed members", interactive=False)
-            gr.Dataframe(value=_research_metrics_table(), label="Reported evaluation metrics", interactive=False)
-            gr.Markdown(
-                """
-### Practical interpretation
-
-High validation/test scores from the research split do not make this a clinical diagnostic tool. Before any real-world medical use, the model would need independent external validation, bias checks, clinical review, calibration review, privacy/security review, and regulatory evaluation.
-                """
-            )
-
-    gr.Markdown(
-        "<div class='footer-note'>Built for Hugging Face Spaces with Gradio. Put all required checkpoint files inside the repository's <code>models/</code> folder.</div>"
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=5, min_width=360):
+            with gr.Group(elem_classes=["glass-card"]):
+                gr.HTML('<h3 class="card-title">Upload MRI Image</h3>')
+                image_input = gr.Image(
+                    label="Upload MRI Image",
+                    show_label=False,
+                    type="pil",
+                    height=430,
+                    sources=["upload", "clipboard"],
+                )
+                heatmap_toggle = gr.Checkbox(
+                    value=False,
+                    label="Generate optional heatmap",
+                    info=None,
+                )
+                run_button = gr.Button("Run Prediction", variant="primary", elem_id="run_button")
+
+            with gr.Group(elem_classes=["glass-card"]):
+                gr.HTML('<h3 class="card-title">Optional Heatmap</h3>')
+                heatmap_output = gr.Image(
+                    label="Optional Heatmap",
+                    show_label=False,
+                    type="pil",
+                    height=390,
+                )
+
+        with gr.Column(scale=7, min_width=420):
+            prediction_html = gr.HTML(_empty_prediction_card())
+            probabilities_html = gr.HTML(_probabilities_card())
+
+    run_button.click(
+        fn=run_prediction,
+        inputs=[image_input, heatmap_toggle],
+        outputs=[prediction_html, probabilities_html, heatmap_output],
     )
 
+    with gr.Accordion("Model status and research details", open=False):
+        status_md = gr.HTML(_status_markdown())
+        status_table = gr.Dataframe(value=_model_table(), interactive=False, label="Required checkpoint files")
+        with gr.Row():
+            refresh_btn = gr.Button("Refresh status")
+            load_btn = gr.Button("Test-load ensemble")
+        load_status = gr.Textbox(label="Load result", interactive=False)
+
+        gr.Markdown("### Deployed members")
+        gr.Dataframe(value=_deployed_members_table(), interactive=False, label="Selected non-zero-weight members")
+        gr.Markdown("### Reported evaluation metrics")
+        gr.Dataframe(value=_research_metrics_table(), interactive=False, label="Research split metrics")
+
+        refresh_btn.click(fn=_status_markdown, inputs=None, outputs=status_md)
+        refresh_btn.click(fn=_model_table, inputs=None, outputs=status_table)
+        load_btn.click(fn=warmup_status, inputs=None, outputs=load_status)
+
+    gr.HTML('<div class="footer-note">Research prototype only — not for clinical diagnosis.</div>')
+
 
 if __name__ == "__main__":
     demo.queue(max_size=16).launch()