Professional UI overhaul: custom CSS, badges, human-readable gallery, auto-load, section headings

app.py

@@ -48,6 +48,90 @@ NO_GPU_MSG = (
     "or run the app locally with a GPU: `python app.py`"
 )
 
+# ── Custom CSS ───────────────────────────────────────────────────────
+CUSTOM_CSS = """
+.header-container {
+    text-align: center;
+    padding: 1.5rem 1rem 0.5rem 1rem;
+}
+.header-container h1 {
+    font-size: 1.8rem;
+    font-weight: 700;
+    margin-bottom: 0.3rem;
+}
+.header-subtitle {
+    text-align: center;
+    color: #555;
+    font-size: 0.95rem;
+    margin-bottom: 0.8rem;
+}
+.badge-row {
+    display: flex;
+    justify-content: center;
+    gap: 0.6rem;
+    flex-wrap: wrap;
+    margin-bottom: 1rem;
+}
+.badge {
+    display: inline-block;
+    padding: 0.25rem 0.75rem;
+    border-radius: 999px;
+    font-size: 0.8rem;
+    font-weight: 600;
+    background: #e8eaf6;
+    color: #3949ab;
+}
+.badge-green {
+    background: #e8f5e9;
+    color: #2e7d32;
+}
+.badge-purple {
+    background: #f3e5f5;
+    color: #7b1fa2;
+}
+.badge-orange {
+    background: #fff3e0;
+    color: #e65100;
+}
+.gpu-notice {
+    background: #fff8e1;
+    border: 1px solid #ffe082;
+    border-radius: 8px;
+    padding: 0.75rem 1rem;
+    margin-bottom: 1rem;
+    font-size: 0.9rem;
+    color: #6d4c00;
+}
+.section-heading {
+    font-size: 1.05rem;
+    font-weight: 600;
+    color: #333;
+    margin-bottom: 0.5rem;
+    border-bottom: 2px solid #e0e0e0;
+    padding-bottom: 0.3rem;
+}
+.gallery-info {
+    background: #f5f5f5;
+    border-radius: 8px;
+    padding: 0.6rem 1rem;
+    margin-bottom: 0.8rem;
+    font-size: 0.88rem;
+    color: #555;
+}
+.about-section {
+    max-width: 800px;
+    margin: 0 auto;
+    padding: 1rem;
+}
+footer {
+    text-align: center;
+    padding: 1rem;
+    color: #999;
+    font-size: 0.8rem;
+}
+"""
+
+
 # ── Lazy imports (avoid crash if no GPU) ─────────────────────────────
 _model_cache = {"model": None, "uni_model": None, "spatial_pool_size": 32}
 
@@ -188,13 +272,13 @@ def _generate_all_gpu(image, guidance_scale):
 def generate_single_stain(image, stain, guidance_scale):
     """Wrapper with GPU availability check."""
     if image is None:
-        return None, "No image uploaded"
+        return None, "Please upload an H&E image first."
     if not GPU_AVAILABLE and not HAS_SPACES:
         return None, NO_GPU_MSG
     try:
         t0 = time.time()
         result = _generate_single_gpu(image, stain, guidance_scale)
-        return result, f"{time.time() - t0:.2f}s"
+        return result, f"Generated in {time.time() - t0:.2f}s"
     except RuntimeError as e:
         if "NVIDIA" in str(e) or "CUDA" in str(e) or "cuda" in str(e):
             return None, NO_GPU_MSG
@@ -204,13 +288,13 @@ def generate_single_stain(image, stain, guidance_scale):
 def generate_all_stains(image, guidance_scale):
     """Wrapper with GPU availability check."""
     if image is None:
-        return None, None, None, None, None, "No image uploaded"
+        return None, None, None, None, None, "Please upload an H&E image first."
     if not GPU_AVAILABLE and not HAS_SPACES:
         return None, None, None, None, None, NO_GPU_MSG
     try:
         t0 = time.time()
         he_pil, results = _generate_all_gpu(image, guidance_scale)
-        elapsed = f"{time.time() - t0:.2f}s"
+        elapsed = f"Generated all 4 stains in {time.time() - t0:.2f}s"
         return he_pil, results["HER2"], results["Ki67"], results["ER"], results["PR"], elapsed
     except RuntimeError as e:
         if "NVIDIA" in str(e) or "CUDA" in str(e) or "cuda" in str(e):
@@ -228,10 +312,30 @@ def load_gallery():
         return json.load(f)
 
 
-def show_gallery(name, gallery):
-    if not name or not gallery or name not in gallery:
-        return None, None, None, None, None, None
-    entry = gallery[name]
+def _make_gallery_label(key, entry):
+    """Create a human-readable label for a gallery entry."""
+    source = entry.get("source", "")
+    gt_stain = entry.get("gt_stain", "")
+    # Extract a short sample ID from the key
+    parts = key.split("_")
+    if source == "BCI":
+        # e.g. BCI_HER2_3+_00277_test_3+ -> "BCI - HER2 3+ (#00277)"
+        her2_class = parts[2] if len(parts) > 2 else ""
+        sample_id = parts[3] if len(parts) > 3 else ""
+        return f"BCI - HER2 {her2_class} (#{sample_id})"
+    else:
+        # e.g. MIST_Ki67_10M2102916_10_20 -> "MIST - Ki67 (10M2102916)"
+        stain = parts[1] if len(parts) > 1 else ""
+        sample_id = parts[2] if len(parts) > 2 else ""
+        return f"MIST - {stain} ({sample_id})"
+
+
+def show_gallery(display_name, gallery, name_map):
+    """Show a gallery example by its display name."""
+    key = name_map.get(display_name)
+    if not key or not gallery or key not in gallery:
+        return None, None, None, None, None, None, ""
+    entry = gallery[key]
     base = GALLERY_DIR / "images"
     he = Image.open(base / entry["he"]).convert("RGB") if "he" in entry else None
     gt = Image.open(base / entry["gt"]).convert("RGB") if "gt" in entry else None
@@ -239,42 +343,121 @@ def show_gallery(name, gallery):
     gen_ki67 = Image.open(base / entry["gen_ki67"]).convert("RGB") if "gen_ki67" in entry else None
     gen_er = Image.open(base / entry["gen_er"]).convert("RGB") if "gen_er" in entry else None
     gen_pr = Image.open(base / entry["gen_pr"]).convert("RGB") if "gen_pr" in entry else None
-    return he, gt, gen_her2, gen_ki67, gen_er, gen_pr
+
+    source = entry.get("source", "Unknown")
+    gt_stain = entry.get("gt_stain", "Unknown")
+    info = f"**Dataset:** {source} | **Ground truth stain:** {gt_stain}"
+    return he, gt, gen_her2, gen_ki67, gen_er, gen_pr, info
 
 
 # ── Build Gradio App ─────────────────────────────────────────────────
 
-gallery = load_gallery()
-gallery_names = list(gallery.keys()) if gallery else []
-
-with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
-    gr.Markdown(
-        """
-        # UNIStainNet: Foundation-Model-Guided Virtual Staining
-        **H&E → IHC (HER2, Ki67, ER, PR)**  | 
-        Single unified model  |  42M parameters  | 
-        UNI spatial conditioning
-        """
-    )
-
-    # ── Tab 1: Single Stain ──────────────────────────────────────
-    with gr.Tab("Virtual Staining"):
-        if not GPU_AVAILABLE and not HAS_SPACES:
-            gr.Markdown(
-                f"> **Note:** {NO_GPU_MSG}"
+gallery_data = load_gallery()
+gallery_name_map = {}  # display_name -> key
+gallery_display_names = []
+if gallery_data:
+    for key, entry in gallery_data.items():
+        label = _make_gallery_label(key, entry)
+        gallery_name_map[label] = key
+        gallery_display_names.append(label)
+
+with gr.Blocks(title="UNIStainNet — Virtual IHC Staining", css=CUSTOM_CSS) as demo:
+
+    # ── Header ──────────────────────────────────────────────────────
+    gr.HTML("""
+    <div class="header-container">
+        <h1>UNIStainNet</h1>
+        <p style="font-size:1.1rem; color:#555; margin-top:0.2rem;">
+            Foundation-Model-Guided Virtual Staining of H&amp;E to IHC
+        </p>
+    </div>
+    <div class="header-subtitle">
+        Translate H&amp;E histopathology images into immunohistochemistry (IHC) stains
+        for breast cancer biomarkers using a single unified deep learning model.
+    </div>
+    <div class="badge-row">
+        <span class="badge">42M Parameters</span>
+        <span class="badge-green badge">4 IHC Stains</span>
+        <span class="badge-purple badge">UNI Foundation Model</span>
+        <span class="badge-orange badge">Single Forward Pass</span>
+    </div>
+    """)
+
+    # ── Tab 1: Gallery (default — works without GPU) ────────────
+    with gr.Tab("Gallery", id="gallery"):
+        if not gallery_display_names:
+            gr.Markdown("No pre-computed gallery available.")
+        else:
+            gr.HTML("""
+            <div class="gallery-info">
+                Browse pre-computed virtual staining results — <strong>no GPU required</strong>.
+                Each example shows the H&amp;E input, ground truth IHC, and all 4 generated stains from our unified model.
+            </div>
+            """)
+            with gr.Row():
+                gallery_dropdown = gr.Dropdown(
+                    choices=gallery_display_names,
+                    value=gallery_display_names[0] if gallery_display_names else None,
+                    label="Select Example",
+                    scale=3,
+                )
+                gallery_info_box = gr.Markdown(value="", scale=2)
+
+            gr.HTML('<p class="section-heading">Input & Ground Truth</p>')
+            with gr.Row():
+                gal_he = gr.Image(type="pil", label="H&E Input", height=280, show_download_button=False)
+                gal_gt = gr.Image(type="pil", label="Ground Truth IHC", height=280, show_download_button=False)
+
+            gr.HTML('<p class="section-heading">Generated IHC Stains (all from the same H&E)</p>')
+            with gr.Row():
+                gal_her2 = gr.Image(type="pil", label="Generated HER2", height=280, show_download_button=False)
+                gal_ki67 = gr.Image(type="pil", label="Generated Ki67", height=280, show_download_button=False)
+                gal_er = gr.Image(type="pil", label="Generated ER", height=280, show_download_button=False)
+                gal_pr = gr.Image(type="pil", label="Generated PR", height=280, show_download_button=False)
+
+            def _show_gallery_wrapper(display_name):
+                return show_gallery(display_name, gallery_data, gallery_name_map)
+
+            gallery_dropdown.change(
+                fn=_show_gallery_wrapper,
+                inputs=[gallery_dropdown],
+                outputs=[gal_he, gal_gt, gal_her2, gal_ki67, gal_er, gal_pr, gallery_info_box],
             )
+
+            # Auto-load first example
+            demo.load(
+                fn=lambda: _show_gallery_wrapper(gallery_display_names[0]) if gallery_display_names else (None,) * 7,
+                outputs=[gal_he, gal_gt, gal_her2, gal_ki67, gal_er, gal_pr, gallery_info_box],
+            )
+
+    # ── Tab 2: Single Stain ──────────────────────────────────────
+    with gr.Tab("Virtual Staining", id="inference"):
+        if not GPU_AVAILABLE and not HAS_SPACES:
+            gr.HTML(f'<div class="gpu-notice">{NO_GPU_MSG}</div>')
+        else:
+            gr.HTML("""
+            <div class="gallery-info">
+                Upload an H&amp;E image and select a target IHC stain. The model generates
+                the virtual stain in a single forward pass (~1 second on GPU).
+            </div>
+            """)
         with gr.Row():
             with gr.Column(scale=1):
-                input_image = gr.Image(type="pil", label="Upload H&E Image", height=400)
-                stain_choice = gr.Radio(choices=STAIN_NAMES, value="HER2", label="Target IHC Stain")
+                input_image = gr.Image(type="pil", label="Upload H&E Image", height=380)
+                stain_choice = gr.Radio(
+                    choices=STAIN_NAMES, value="HER2",
+                    label="Target IHC Stain",
+                    info="Select which immunohistochemistry marker to generate",
+                )
                 guidance_slider = gr.Slider(
                     minimum=1.0, maximum=3.0, step=0.1, value=1.0,
-                    label="Guidance Scale (1.0 = no CFG)",
+                    label="Guidance Scale",
+                    info="1.0 = standard generation, higher = stronger stain signal (CFG)",
                 )
-                generate_btn = gr.Button("Generate", variant="primary")
+                generate_btn = gr.Button("Generate", variant="primary", size="lg")
                 gen_time = gr.Textbox(label="Status", interactive=False)
             with gr.Column(scale=1):
-                output_image = gr.Image(type="pil", label="Generated IHC", height=400)
+                output_image = gr.Image(type="pil", label="Generated IHC", height=380)
 
         generate_btn.click(
             fn=generate_single_stain,
@@ -282,30 +465,36 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
             outputs=[output_image, gen_time],
         )
 
-    # ── Tab 2: Cross-Stain ───────────────────────────────────────
-    with gr.Tab("Cross-Stain Comparison"):
+    # ── Tab 3: Cross-Stain ───────────────────────────────────────
+    with gr.Tab("Cross-Stain Comparison", id="cross-stain"):
         if not GPU_AVAILABLE and not HAS_SPACES:
-            gr.Markdown(
-                f"> **Note:** {NO_GPU_MSG}"
-            )
-        gr.Markdown(
-            "Generate **all 4 IHC stains** from a single H&E input. "
-            "Demonstrates the unified multi-stain capability."
-        )
+            gr.HTML(f'<div class="gpu-notice">{NO_GPU_MSG}</div>')
+        else:
+            gr.HTML("""
+            <div class="gallery-info">
+                Generate <strong>all 4 IHC stains</strong> from a single H&amp;E input.
+                This demonstrates the unified multi-stain capability of UNIStainNet.
+            </div>
+            """)
         with gr.Row():
-            cross_input = gr.Image(type="pil", label="Upload H&E Image", height=350)
-            cross_guidance = gr.Slider(
-                minimum=1.0, maximum=3.0, step=0.1, value=1.0, label="Guidance Scale",
-            )
-        cross_btn = gr.Button("Generate All Stains", variant="primary")
-        cross_time = gr.Textbox(label="Status", interactive=False)
+            with gr.Column(scale=1):
+                cross_input = gr.Image(type="pil", label="Upload H&E Image", height=300)
+            with gr.Column(scale=1):
+                cross_guidance = gr.Slider(
+                    minimum=1.0, maximum=3.0, step=0.1, value=1.0,
+                    label="Guidance Scale",
+                    info="1.0 = standard generation, higher = stronger stain signal",
+                )
+                cross_btn = gr.Button("Generate All 4 Stains", variant="primary", size="lg")
+                cross_time = gr.Textbox(label="Status", interactive=False)
 
+        gr.HTML('<p class="section-heading">Results</p>')
         with gr.Row():
-            cross_he_out = gr.Image(type="pil", label="H&E Input", height=300)
-            cross_her2 = gr.Image(type="pil", label="HER2", height=300)
-            cross_ki67 = gr.Image(type="pil", label="Ki67", height=300)
-            cross_er = gr.Image(type="pil", label="ER", height=300)
-            cross_pr = gr.Image(type="pil", label="PR", height=300)
+            cross_he_out = gr.Image(type="pil", label="H&E Input", height=250)
+            cross_her2 = gr.Image(type="pil", label="HER2", height=250)
+            cross_ki67 = gr.Image(type="pil", label="Ki67", height=250)
+            cross_er = gr.Image(type="pil", label="ER", height=250)
+            cross_pr = gr.Image(type="pil", label="PR", height=250)
 
         cross_btn.click(
             fn=generate_all_stains,
@@ -313,69 +502,76 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
             outputs=[cross_he_out, cross_her2, cross_ki67, cross_er, cross_pr, cross_time],
         )
 
-    # ── Tab 3: Gallery ───────────────────────────────────────────
-    with gr.Tab("Gallery"):
-        if not gallery_names:
-            gr.Markdown("No pre-computed gallery available.")
-        else:
-            gr.Markdown(
-                "Pre-computed examples — no GPU required. "
-                "Select an example to view the H&E input, ground truth, and generated IHC stains."
-            )
-            gallery_dropdown = gr.Dropdown(
-                choices=gallery_names,
-                value=gallery_names[0] if gallery_names else None,
-                label="Select Example",
-            )
-            with gr.Row():
-                gal_he = gr.Image(type="pil", label="H&E Input", height=300)
-                gal_gt = gr.Image(type="pil", label="Ground Truth IHC", height=300)
-            with gr.Row():
-                gal_her2 = gr.Image(type="pil", label="Generated HER2", height=300)
-                gal_ki67 = gr.Image(type="pil", label="Generated Ki67", height=300)
-                gal_er = gr.Image(type="pil", label="Generated ER", height=300)
-                gal_pr = gr.Image(type="pil", label="Generated PR", height=300)
-
-            gallery_dropdown.change(
-                fn=lambda name: show_gallery(name, gallery),
-                inputs=[gallery_dropdown],
-                outputs=[gal_he, gal_gt, gal_her2, gal_ki67, gal_er, gal_pr],
-            )
-
     # ── Tab 4: About ─────────────────────────────────────────────
-    with gr.Tab("About"):
+    with gr.Tab("About", id="about"):
         gr.Markdown(
             """
-            ## UNIStainNet
+            <div class="about-section">
+
+            ## UNIStainNet: Foundation-Model-Guided Virtual Staining
 
-            A SPADE-UNet generator conditioned on dense spatial tokens from a frozen
-            [UNI](https://github.com/mahmoodlab/UNI) pathology foundation model (ViT-L/16).
+            UNIStainNet is a deep learning model for **virtual immunohistochemistry (IHC) staining**
+            from standard hematoxylin & eosin (H&E) histopathology images. It translates routine H&E
+            slides into IHC stains for four clinically important breast cancer biomarkers:
+            **HER2**, **Ki67**, **ER**, and **PR**.
 
-            **Key features:**
-            - Dense UNI spatial conditioning (32x32 = 1,024 tokens)
-            - Misalignment-aware loss suite for consecutive-section training pairs
-            - Single unified model serves 4 IHC markers (HER2, Ki67, ER, PR)
-            - 42M generator parameters, single forward pass inference
+            ### Why Virtual Staining?
 
-            ### Architecture
+            - **Tissue conservation** — eliminates the need for additional serial sections
+            - **Faster turnaround** — results in seconds instead of hours/days
+            - **Cost reduction** — one H&E slide replaces multiple IHC tests for screening
+            - **Consistency** — no batch-to-batch staining variability
+
+            ### How It Works
+
+            The model uses a **SPADE-UNet generator** conditioned on dense spatial features from a
+            frozen [UNI](https://github.com/mahmoodlab/UNI) pathology foundation model (ViT-L/16,
+            pretrained on 100M+ histopathology patches). A FiLM-based stain embedding allows a
+            **single unified model** to generate all 4 IHC stains.
 
             | Component | Details |
             |-----------|---------|
-            | Generator | SPADE-UNet with UNI spatial conditioning + FiLM stain embeddings |
-            | Discriminator | Multi-scale PatchGAN (512 + 256) with spectral norm |
-            | UNI Features | 4x4 sub-crop tiling → UNI ViT-L/16 → 32x32 spatial tokens |
-            | Parameters | 42M (generator), UNI frozen (303M) |
-
-            ### Results (MIST, unified model)
-
-            | Stain | FID ↓ | KID×1k ↓ | Pearson-R ↑ | DAB KL ↓ |
-            |-------|-------|-----------|-------------|----------|
-            | HER2  | 34.5  | 2.2       | 0.929       | 0.166    |
-            | Ki67  | 27.2  | 1.8       | 0.927       | 0.119    |
-            | ER    | 29.2  | 1.8       | 0.949       | 0.182    |
-            | PR    | 29.0  | 1.1       | 0.943       | 0.171    |
+            | **Generator** | SPADE-UNet with UNI spatial conditioning + FiLM stain embeddings |
+            | **Foundation Model** | UNI ViT-L/16 (frozen, 303M parameters) |
+            | **Spatial Tokens** | 4x4 sub-crop tiling of H&E input, yielding 32x32 = 1,024 tokens |
+            | **Generator Parameters** | 42M |
+            | **Inference** | Single forward pass (~1 second on GPU) |
+
+            ### Quantitative Results (MIST Dataset, Unified Model)
+
+            | Stain | FID | KID x1k | Pearson-R | DAB KL |
+            |-------|-----|---------|-----------|--------|
+            | HER2  | 34.5 | 2.2   | 0.929     | 0.166  |
+            | Ki67  | 27.2 | 1.8   | 0.927     | 0.119  |
+            | ER    | 29.2 | 1.8   | 0.949     | 0.182  |
+            | PR    | 29.0 | 1.1   | 0.943     | 0.171  |
+
+            ### Key Innovations
+
+            - **Dense UNI spatial conditioning**: Unlike prior methods that use global image features,
+              UNIStainNet extracts spatially-resolved features at 32x32 resolution, enabling the generator
+              to leverage fine-grained morphological context from the pathology foundation model.
+            - **Misalignment-aware training**: Because H&E and IHC are cut from consecutive tissue sections
+              (not the same section), there are inherent spatial shifts. Our loss suite (perceptual loss,
+              DAB intensity supervision, unconditional discriminator) is designed to handle this misalignment.
+            - **Classifier-free guidance (CFG)**: 10% class dropout and 10% UNI dropout during training
+              enables tunable generation strength at inference time.
+
+            ### Disclaimer
+
+            This is a **research tool** for exploratory analysis. It is not intended for clinical diagnosis
+            and has not undergone regulatory validation. Generated stains should not be used for treatment decisions.
+
+            </div>
             """
         )
 
+    # ── Footer ───────────────────────────────────────────────────
+    gr.HTML("""
+    <footer>
+        UNIStainNet | Built with Gradio
+    </footer>
+    """)
+
 if __name__ == "__main__":
     demo.launch(theme=gr.themes.Soft())