Fix: inline styles (no CSS dependency), remove GT images (licensing), muted badge colors

app.py

@@ -48,89 +48,6 @@ 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}
@@ -315,30 +232,26 @@ def load_gallery():
 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."""
+    """Show a gallery example by its display name. Only model-generated outputs shown."""
     key = name_map.get(display_name)
     if not key or not gallery or key not in gallery:
-        return None, None, None, None, None, None, ""
+        return None, None, None, None, None, ""
     entry = gallery[key]
     base = GALLERY_DIR / "images"
+    # Only show H&E input and model-generated stains (no ground truth — licensing)
     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
     gen_her2 = Image.open(base / entry["gen_her2"]).convert("RGB") if "gen_her2" in entry else None
     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
@@ -346,8 +259,8 @@ def show_gallery(display_name, gallery, name_map):
 
     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
+    info = f"**Source:** {source} | **Original IHC stain:** {gt_stain}"
+    return he, gen_her2, gen_ki67, gen_er, gen_pr, info
 
 
 # ── Build Gradio App ─────────────────────────────────────────────────
@@ -361,25 +274,29 @@ if gallery_data:
         gallery_name_map[label] = key
         gallery_display_names.append(label)
 
-with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
+with gr.Blocks(title="UNIStainNet -- Virtual IHC Staining") as demo:
 
-    # ── Header ──────────────────────────────────────────────────────
+    # ── Header (all inline styles for Gradio 6.x compatibility) ──
     gr.HTML("""
-    <div class="header-container">
-        <h1>UNIStainNet</h1>
-        <p style="font-size:1.1rem; color:#555; margin-top:0.2rem;">
+    <div style="text-align:center; padding:1.5rem 1rem 0.5rem 1rem;">
+        <h1 style="font-size:1.8rem; font-weight:700; margin-bottom:0.3rem;">UNIStainNet</h1>
+        <p style="font-size:1.05rem; color:#555; margin-top:0.2rem;">
             Foundation-Model-Guided Virtual Staining of H&amp;E to IHC
         </p>
     </div>
-    <div class="header-subtitle">
+    <p style="text-align:center; color:#555; font-size:0.95rem; margin-bottom:0.8rem;">
         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>
+    </p>
+    <div style="display:flex; justify-content:center; gap:0.6rem; flex-wrap:wrap; margin-bottom:1rem;">
+        <span style="display:inline-block; padding:0.25rem 0.75rem; border-radius:999px;
+                     font-size:0.8rem; font-weight:600; background:#dce3f9; color:#1a3a8a;">42M Parameters</span>
+        <span style="display:inline-block; padding:0.25rem 0.75rem; border-radius:999px;
+                     font-size:0.8rem; font-weight:600; background:#d4edda; color:#155724;">4 IHC Stains</span>
+        <span style="display:inline-block; padding:0.25rem 0.75rem; border-radius:999px;
+                     font-size:0.8rem; font-weight:600; background:#e8d5f5; color:#5b1a8a;">UNI Foundation Model</span>
+        <span style="display:inline-block; padding:0.25rem 0.75rem; border-radius:999px;
+                     font-size:0.8rem; font-weight:600; background:#f5ddc4; color:#7a3b10;">Single Forward Pass</span>
     </div>
     """)
 
@@ -388,12 +305,10 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
         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>
-            """)
+            gr.Markdown(
+                "Browse pre-computed virtual staining results -- **no GPU required**. "
+                "Each example shows the H&E input and all 4 generated IHC stains from our unified model."
+            )
             gallery_dropdown = gr.Dropdown(
                 choices=gallery_display_names,
                 value=gallery_display_names[0] if gallery_display_names else None,
@@ -401,15 +316,14 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
             )
             gallery_info_box = gr.Markdown(value="")
 
-            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)
-                gal_gt = gr.Image(type="pil", label="Ground Truth IHC", height=280)
+            gr.Markdown("### H&E Input")
+            gal_he = gr.Image(type="pil", label="H&E Input", height=300)
 
-            gr.HTML('<p class="section-heading">Generated IHC Stains (all from the same H&E)</p>')
+            gr.Markdown("### Generated IHC Stains")
             with gr.Row():
                 gal_her2 = gr.Image(type="pil", label="Generated HER2", height=280)
                 gal_ki67 = gr.Image(type="pil", label="Generated Ki67", height=280)
+            with gr.Row():
                 gal_er = gr.Image(type="pil", label="Generated ER", height=280)
                 gal_pr = gr.Image(type="pil", label="Generated PR", height=280)
 
@@ -419,26 +333,28 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
             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],
+                outputs=[gal_he, 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],
+                fn=lambda: _show_gallery_wrapper(gallery_display_names[0]) if gallery_display_names else (None,) * 6,
+                outputs=[gal_he, 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>')
+            gr.HTML(
+                '<div style="background:#fff8e1; border:1px solid #ffe082; border-radius:8px; '
+                'padding:0.75rem 1rem; margin-bottom:1rem; font-size:0.9rem; color:#6d4c00;">'
+                f'{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>
-            """)
+            gr.Markdown(
+                "Upload an H&E image and select a target IHC stain. "
+                "The model generates the virtual stain in a single forward pass (~1 second on GPU)."
+            )
         with gr.Row():
             with gr.Column(scale=1):
                 input_image = gr.Image(type="pil", label="Upload H&E Image", height=380)
@@ -466,14 +382,16 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
     # ── Tab 3: Cross-Stain ───────────────────────────────────────
     with gr.Tab("Cross-Stain Comparison", id="cross-stain"):
         if not GPU_AVAILABLE and not HAS_SPACES:
-            gr.HTML(f'<div class="gpu-notice">{NO_GPU_MSG}</div>')
+            gr.HTML(
+                '<div style="background:#fff8e1; border:1px solid #ffe082; border-radius:8px; '
+                'padding:0.75rem 1rem; margin-bottom:1rem; font-size:0.9rem; color:#6d4c00;">'
+                f'{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>
-            """)
+            gr.Markdown(
+                "Generate **all 4 IHC stains** from a single H&E input. "
+                "This demonstrates the unified multi-stain capability of UNIStainNet."
+            )
         with gr.Row():
             with gr.Column(scale=1):
                 cross_input = gr.Image(type="pil", label="Upload H&E Image", height=300)
@@ -486,7 +404,7 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
                 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>')
+        gr.Markdown("### Results")
         with gr.Row():
             cross_he_out = gr.Image(type="pil", label="H&E Input", height=250)
             cross_her2 = gr.Image(type="pil", label="HER2", height=250)
@@ -504,72 +422,68 @@ with gr.Blocks(title="UNIStainNet — Virtual IHC Staining") as demo:
     with gr.Tab("About", id="about"):
         gr.Markdown(
             """
-            <div class="about-section">
-
-            ## UNIStainNet: Foundation-Model-Guided Virtual Staining
-
-            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**.
-
-            ### Why Virtual Staining?
-
-            - **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 |
-            | **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>
+## UNIStainNet: Foundation-Model-Guided Virtual Staining
+
+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**.
+
+### Why Virtual Staining?
+
+- **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 |
+| **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.
             """
         )
 
     # ── Footer ───────────────────────────────────────────────────
     gr.HTML("""
-    <footer>
+    <p style="text-align:center; padding:1rem; color:#999; font-size:0.8rem;">
         UNIStainNet | Built with Gradio
-    </footer>
+    </p>
     """)
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft(), css=CUSTOM_CSS)
+    demo.launch()