#!/usr/bin/env python3 """ UNIStainNet Interactive Demo — Hugging Face Spaces Virtual staining of H&E histopathology images to IHC (HER2, Ki67, ER, PR). Supports ZeroGPU (HF Pro) for live inference, falls back to gallery-only on CPU. """ import json import os import time from pathlib import Path import gradio as gr import numpy as np import torch import torch.nn.functional as F import torchvision.transforms as T import torchvision.transforms.functional as TF from PIL import Image from huggingface_hub import hf_hub_download # ── ZeroGPU support ────────────────────────────────────────────────── try: import spaces HAS_SPACES = True except ImportError: spaces = None HAS_SPACES = False GPU_AVAILABLE = torch.cuda.is_available() def _gpu_decorator(duration=60): if HAS_SPACES and hasattr(spaces, "GPU"): return spaces.GPU(duration=duration) return lambda fn: fn # ── Constants ──────────────────────────────────────────────────────── STAIN_NAMES = ["HER2", "Ki67", "ER", "PR"] GALLERY_DIR = Path(__file__).parent / "gallery" TARGET_SIZE = 512 MODEL_REPO = os.environ.get("MODEL_REPO", "faceless-void/UNIStainNet") CHECKPOINT_FILENAME = "mist_multistain_last.ckpt" NO_GPU_MSG = ( "GPU is not available on this Space (requires HF Pro for ZeroGPU). " "Please use the **Gallery** tab to browse pre-computed results, " "or run the app locally with a GPU: `python app.py`" ) # ── Lazy imports (avoid crash if no GPU) ───────────────────────────── _model_cache = {"model": None, "uni_model": None, "spatial_pool_size": 32} def _get_checkpoint_path(): local_path = Path(__file__).parent / "checkpoints" / CHECKPOINT_FILENAME if local_path.exists(): return str(local_path) return hf_hub_download(repo_id=MODEL_REPO, filename=CHECKPOINT_FILENAME) def _load_models(): """Load models onto GPU. Only called when GPU is confirmed available.""" from src.models.trainer import UNIStainNetTrainer import timm if _model_cache["model"] is None: ckpt_path = _get_checkpoint_path() print(f"Loading UNIStainNet from {ckpt_path} ...") model = UNIStainNetTrainer.load_from_checkpoint(ckpt_path, strict=False) model = model.cuda().eval() _model_cache["model"] = model _model_cache["spatial_pool_size"] = getattr(model.hparams, "uni_spatial_size", 32) print("Loading UNI ViT-L/16 ...") uni_model = timm.create_model( "hf-hub:MahmoodLab/uni", pretrained=True, init_values=1e-5, dynamic_img_size=True, ) uni_model = uni_model.cuda().eval() _model_cache["uni_model"] = uni_model print(" Models loaded") else: _model_cache["model"] = _model_cache["model"].cuda() _model_cache["uni_model"] = _model_cache["uni_model"].cuda() return _model_cache["model"], _model_cache["uni_model"], _model_cache["spatial_pool_size"] # ── Preprocessing ──────────────────────────────────────────────────── def preprocess_he(pil_image, target_size=TARGET_SIZE): w, h = pil_image.size short = min(w, h) left = (w - short) // 2 top = (h - short) // 2 pil_image = pil_image.crop((left, top, left + short, top + short)) if short != target_size: pil_image = pil_image.resize((target_size, target_size), Image.BICUBIC) return pil_image def pil_to_tensor(pil_image): t = TF.to_tensor(pil_image) t = TF.normalize(t, [0.5] * 3, [0.5] * 3) return t.unsqueeze(0) def tensor_to_pil(tensor): t = ((tensor[0].cpu() + 1) / 2).clamp(0, 1) return TF.to_pil_image(t) def extract_uni_features(uni_model, he_tensor_01, spatial_pool_size=32): from src.data.mist_dataset import STAIN_TO_LABEL uni_transform = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) B = he_tensor_01.shape[0] num_crops, patches_per_side = 4, 14 crop_h = he_tensor_01.shape[2] // num_crops crop_w = he_tensor_01.shape[3] // num_crops sub_crops = [] for i in range(num_crops): for j in range(num_crops): sub = he_tensor_01[:, :, i*crop_h:(i+1)*crop_h, j*crop_w:(j+1)*crop_w] sub = F.interpolate(sub, size=(224, 224), mode="bicubic", align_corners=False) sub = torch.stack([uni_transform(s) for s in sub]) sub_crops.append(sub) all_crops = torch.stack(sub_crops, dim=1).reshape(B * 16, 3, 224, 224).cuda() with torch.no_grad(): all_feats = uni_model.forward_features(all_crops) patch_tokens = all_feats[:, 1:, :] patch_tokens = patch_tokens.reshape(B, num_crops, num_crops, patches_per_side, patches_per_side, 1024) full_size = num_crops * patches_per_side full_grid = patch_tokens.permute(0, 1, 3, 2, 4, 5).reshape(B, full_size, full_size, 1024) S = spatial_pool_size if S < full_size: grid_bchw = full_grid.permute(0, 3, 1, 2) pooled = F.adaptive_avg_pool2d(grid_bchw, S) result = pooled.permute(0, 2, 3, 1) else: result = full_grid return result.reshape(B, S * S, 1024) # ── Inference functions ────────────────────────────────────────────── @_gpu_decorator(duration=60) def _generate_single_gpu(image, stain, guidance_scale): """GPU inference for single stain.""" from src.data.mist_dataset import STAIN_TO_LABEL model, uni_model, spatial_pool_size = _load_models() he_pil = preprocess_he(image) he_tensor = pil_to_tensor(he_pil).cuda() he_01 = ((he_tensor + 1) / 2).clamp(0, 1) uni_feats = extract_uni_features(uni_model, he_01, spatial_pool_size).cuda() labels = torch.tensor([STAIN_TO_LABEL[stain]], device="cuda", dtype=torch.long) with torch.no_grad(): gen = model.generate(he_tensor, uni_feats, labels, guidance_scale=guidance_scale) return tensor_to_pil(gen) @_gpu_decorator(duration=120) def _generate_all_gpu(image, guidance_scale): """GPU inference for all 4 stains.""" from src.data.mist_dataset import STAIN_TO_LABEL model, uni_model, spatial_pool_size = _load_models() he_pil = preprocess_he(image) he_tensor = pil_to_tensor(he_pil).cuda() he_01 = ((he_tensor + 1) / 2).clamp(0, 1) uni_feats = extract_uni_features(uni_model, he_01, spatial_pool_size).cuda() results = {} for stain in STAIN_NAMES: labels = torch.tensor([STAIN_TO_LABEL[stain]], device="cuda", dtype=torch.long) with torch.no_grad(): gen = model.generate(he_tensor, uni_feats, labels, guidance_scale=guidance_scale) results[stain] = tensor_to_pil(gen) return he_pil, results def generate_single_stain(image, stain, guidance_scale): """Wrapper with GPU availability check.""" if image is None: 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"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 raise def generate_all_stains(image, guidance_scale): """Wrapper with GPU availability check.""" if image is None: 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"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): return None, None, None, None, None, NO_GPU_MSG raise # ── Gallery ────────────────────────────────────────────────────────── def load_gallery(): meta_path = GALLERY_DIR / "metadata.json" if not meta_path.exists(): return None with open(meta_path) as f: return json.load(f) def _make_gallery_label(key, entry): """Create a human-readable label for a gallery entry.""" source = entry.get("source", "") parts = key.split("_") if source == "BCI": 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: 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 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 gen_pr = Image.open(base / entry["gen_pr"]).convert("RGB") if "gen_pr" in entry else None gt_stain = entry.get("gt_stain", "Unknown") info = f"**Ground truth stain:** {gt_stain}" return he, gt, gen_her2, gen_ki67, gen_er, gen_pr, info # ── Build Gradio App ───────────────────────────────────────────────── 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") as demo: # ── Header ──────────────────────────────────────────────────── gr.HTML("""

UNIStainNet

Virtual Immunohistochemistry Staining from H&E

Generate HER2, Ki67, ER, and PR stains from a single H&E breast tissue image using one unified deep learning model.

Breast Cancer Biomarkers HER2 / Ki67 / ER / PR One Model, 4 Stains
""") # ── 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.Markdown( "Browse pre-computed virtual staining results -- **no GPU required**. " "Each example shows the H&E input, ground truth IHC, and all 4 generated stains from our unified model." ) gallery_dropdown = gr.Dropdown( choices=gallery_display_names, value=gallery_display_names[0] if gallery_display_names else None, label="Select Example", ) gallery_info_box = gr.Markdown(value="") gr.Markdown("### Input & Ground Truth") 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) gr.Markdown("### Generated IHC Stains (all from the same H&E)") 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) 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'{NO_GPU_MSG}
' ) else: gr.Markdown( "Upload an H&E image and select a target IHC stain to generate." ) with gr.Accordion("Image upload guidelines", open=False): gr.Markdown( "- **Tissue type:** H&E-stained breast cancer tissue\n" "- **Magnification:** 20x recommended (trained on BCI and MIST datasets)\n" "- **Size:** Images are center-cropped and resized to 512x512 internally\n" "- **Format:** PNG, JPEG, or TIFF\n" "- **Best results:** Regions with invasive carcinoma; " "adipose or stromal tissue may produce lower quality output" ) with gr.Row(): with gr.Column(scale=1): 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", info="1.0 = standard generation, higher = stronger stain signal (CFG)", ) 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=380) generate_btn.click( fn=generate_single_stain, inputs=[input_image, stain_choice, guidance_slider], outputs=[output_image, gen_time], ) # ── Tab 3: Cross-Stain ─────────────────────────────────────── with gr.Tab("Cross-Stain Comparison", id="cross-stain"): if not GPU_AVAILABLE and not HAS_SPACES: gr.HTML( '
' f'{NO_GPU_MSG}
' ) else: gr.Markdown( "Generate **all 4 IHC stains** from a single H&E input." ) with gr.Accordion("Image upload guidelines", open=False): gr.Markdown( "- **Tissue type:** H&E-stained breast cancer tissue\n" "- **Magnification:** 20x recommended (trained on BCI and MIST datasets)\n" "- **Size:** Images are center-cropped and resized to 512x512 internally\n" "- **Format:** PNG, JPEG, or TIFF\n" "- **Best results:** Regions with invasive carcinoma; " "adipose or stromal tissue may produce lower quality output" ) with gr.Row(): 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.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) 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, inputs=[cross_input, cross_guidance], outputs=[cross_he_out, cross_her2, cross_ki67, cross_er, cross_pr, cross_time], ) # ── Tab 4: About ───────────────────────────────────────────── with gr.Tab("About", id="about"): gr.Markdown( """ ## 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. ### Links - **Paper**: [arXiv:2603.12716](https://arxiv.org/abs/2603.12716) - **Code**: [github.com/facevoid/UNIStainNet](https://github.com/facevoid/UNIStainNet) - **Project Page**: [facevoid.github.io/UNIStainNet](https://facevoid.github.io/UNIStainNet/) ### 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("""

UNIStainNet | arXiv | GitHub | Project Page

""") if __name__ == "__main__": demo.launch()