app.py

"""
Hugging Face Space: FastAPI + Gradio (Uvicorn loads ``app:app``).

Visualization matches ``tools/infer.py`` when ``--output`` is set: same
``init_detector`` / ``patch_config`` / ``inference_detector`` / ``save_result``
calls (including ``font_size=6``, ``thickness=1``).
"""
import functools
import os
import sys
import tempfile
from collections import OrderedDict

# -----------------------------------------------------------------------------
# Repo bootstrap (must run before ``nuhtc`` imports)
# -----------------------------------------------------------------------------
proj_path = "/workspace"
sys.path.insert(0, f"{proj_path}/thirdparty/mmdetection")
sys.path.insert(0, proj_path)

import torch
import uvicorn
import gradio as gr
import numpy as np
from fastapi import FastAPI
from mmcv import Config
from mmdet.apis import inference_detector
from PIL import Image

from nuhtc.apis.inference import init_detector, save_result
from nuhtc.utils import patch_config

# -----------------------------------------------------------------------------
# Paths & environment
# -----------------------------------------------------------------------------
CONFIG = f"{proj_path}/models/htc_lite_PanNuke_infer.py"
CHECKPOINT = f"{proj_path}/models/pannuke.pth"
DEFAULT_SCORE = 0.35
DEMO_IMG_DIR = f"{proj_path}/demo/imgs"

# Class key → PanNuke-style full name; insertion order is summary / count order
# (indices follow ``model.CLASSES``: T, I, C, D, E).
CLASS_ORDER = OrderedDict(
    [
        ("T", "Neoplastic"),
        ("E", "Non-Neoplastic Epithelial"),
        ("I", "Inflammatory"),
        ("C", "Connective"),
        ("D", "Dead"),
    ]
)

# -----------------------------------------------------------------------------
# Gradio layout (CSS + center flow icon HTML)
# -----------------------------------------------------------------------------
# Center cue between Input and Segmentation (rounded “forward” chevron).
_FLOW_DIVIDER_HTML = """
<div class="nuhtc-flow-wrap" aria-hidden="true">
  <svg width="72" height="72" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
    <path
      d="M2 9.5C2 8.67157 2.67157 8 3.5 8H12V5.41421C12 4.52331 13.0771 4.07714 13.7071 4.70711L20.2929 11.2929C20.6834 11.6834 20.6834 12.3166 20.2929 12.7071L13.7071 19.2929C13.0771 19.9229 12 19.4767 12 18.5858V16H3.5C2.67157 16 2 15.3284 2 14.5V9.5Z"
      fill="#e4e4e7"/>
  </svg>
</div>
"""

_THUMB_CSS = """
.nuhtc-wrap {
  max-width: 1120px;
  margin: 0 auto;
  padding: 0.35rem 0.6rem 1.75rem;
  background: transparent;
}
.nuhtc-header-wrap { text-align: center; padding: 1rem 0.5rem 1.2rem; }
.nuhtc-header-wrap h2 {
  margin: 0 0 0.4rem 0 !important;
  font-size: 1.6rem !important;
  font-weight: 650 !important;
  letter-spacing: -0.025em !important;
  color: #0f172a !important;
  border: none !important;
}
.nuhtc-header-wrap .markdown, .nuhtc-header-wrap .prose { font-size: 0.94rem; color: #64748b; line-height: 1.5; }
.nuhtc-card {
  background: #fff;
  border-radius: 14px;
  border: 1px solid #e8ecf3;
  box-shadow: 0 1px 2px rgba(15, 23, 42, 0.04), 0 10px 28px rgba(15, 23, 42, 0.07);
  padding: 1.05rem 1.2rem 1.15rem;
  margin-bottom: 0.95rem;
}
.nuhtc-core-row { align-items: flex-start !important; }
/* Drop default grey panel behind Gradio Image so only the card + photo read as UI. */
.nuhtc-img-panel .form,
.nuhtc-img-panel .block,
.nuhtc-img-panel .image-container,
.nuhtc-img-panel [class*="image-container"],
.nuhtc-img-panel .wrap {
  background: transparent !important;
  background-color: transparent !important;
  border-color: transparent !important;
  box-shadow: none !important;
}
.nuhtc-img-panel > div {
  background: transparent !important;
  background-color: transparent !important;
}
.nuhtc-img-panel img {
  border-radius: 0 !important;
  box-shadow: 0 8px 24px rgba(15, 23, 42, 0.11), 0 2px 8px rgba(15, 23, 42, 0.06) !important;
}
.nuhtc-img-panel .image-container,
.nuhtc-img-panel [class*="image-container"],
.nuhtc-img-panel .wrap { border-radius: 0 !important; overflow: hidden; }
/* Middle column: vertically centered in the row vs. top-aligned image columns. */
.nuhtc-arrow-col {
  min-width: 80px;
  max-width: 100px;
  align-self: center !important;
  display: flex !important;
  justify-content: center !important;
}
.nuhtc-flow-wrap {
  display: flex;
  align-items: center;
  justify-content: center;
  padding: 0.15rem;
  min-height: 0;
}
.nuhtc-flow-wrap svg { display: block; }
/* Wraps hint (optional) + thumbnails|button row; separator sits above hint or row. */
.nuhtc-demo-section {
  margin-top: 0.65rem;
  padding-top: 0.55rem;
  border-top: 1px solid #eef1f6;
}
.nuhtc-demo-hint-wrap { margin-bottom: 0.4rem; }
.nuhtc-demo-hint-wrap .markdown p {
  margin: 0 !important;
  font-size: 0.8rem !important;
  color: #64748b !important;
  text-align: left !important;
}
.nuhtc-demo-btn-row {
  align-items: center !important;
  justify-content: center !important;
}
/* Demo column: left-align hint + Examples (Gradio body lives in .form / .block). */
.nuhtc-demo-btn-row .nuhtc-demo-strip-wrap,
.nuhtc-demo-btn-row .nuhtc-demo-strip-wrap .form,
.nuhtc-demo-btn-row .nuhtc-demo-strip-wrap .block {
  display: flex !important;
  flex-direction: column !important;
  align-items: flex-start !important;
  width: 100% !important;
}
.nuhtc-demo-strip-wrap {
  min-width: 0;
  text-align: left;
}
.nuhtc-demo-strip-wrap .markdown { width: 100%; max-width: 100%; }
.nuhtc-demo-strip-wrap .markdown,
.nuhtc-demo-strip-wrap .prose { text-align: left !important; }
.nuhtc-demo-strip-wrap .markdown p { margin: 0 0 0.4rem 0 !important; font-size: 0.8rem !important; color: #64748b !important; }
.nuhtc-demo-strip {
  width: 100%;
  max-width: 100%;
  display: flex !important;
  flex-direction: column !important;
  align-items: flex-start !important;
  overflow-x: auto;
  overflow-y: hidden;
  padding: 0.2rem 0 0.35rem;
  -webkit-overflow-scrolling: touch;
}
.nuhtc-demo-strip .form,
.nuhtc-demo-strip .block {
  display: flex !important;
  flex-direction: column !important;
  align-items: flex-start !important;
  width: 100%;
}
.nuhtc-demo-strip::-webkit-scrollbar { height: 7px; }
.nuhtc-demo-strip::-webkit-scrollbar-thumb { background: #cbd5e1; border-radius: 6px; }
.nuhtc-demo-strip .grid-wrap {
  display: flex !important;
  flex-direction: row !important;
  flex-wrap: nowrap !important;
  justify-content: flex-start !important;
  gap: 0.65rem !important;
  width: max-content !important;
}
.nuhtc-demo-strip img {
  max-height: 86px !important;
  width: auto !important;
  max-width: none !important;
  border-radius: 9px !important;
  object-fit: cover;
  box-shadow: 0 2px 10px rgba(15, 23, 42, 0.09);
}
.nuhtc-demo-strip button {
  padding: 4px 6px !important;
  min-height: 0 !important;
  border-radius: 9px !important;
}
.nuhtc-bottom-row { align-items: center !important; }
.nuhtc-slider-col { padding-right: 0.5rem; flex-grow: 1; }
.nuhtc-slider-col label, .nuhtc-slider-col .label-wrap { font-weight: 500 !important; color: #334155 !important; }
.nuhtc-slider-col input[type="range"] { accent-color: #ea580c; }
/* Button column: same column + .form/.block pattern as .nuhtc-demo-strip, but cross-axis centered. */
.nuhtc-btn-col {
  display: flex !important;
  flex-direction: column !important;
  align-items: center !important;
  justify-content: center !important;
  padding: 0.25rem 0 0 0.5rem;
}
.nuhtc-demo-btn-row .nuhtc-btn-col {
  width: 100% !important;
  max-width: 100% !important;
  min-width: 0 !important;
  padding: 0 0 0 0.5rem !important;
}
.nuhtc-demo-btn-row .nuhtc-btn-col .form,
.nuhtc-demo-btn-row .nuhtc-btn-col .block {
  display: flex !important;
  flex-direction: column !important;
  align-items: center !important;
  width: 100% !important;
}
.nuhtc-demo-btn-row .nuhtc-run-btn {
  align-self: center !important;
}
.nuhtc-summary-row { width: 100%; margin-top: 0.75rem; padding-top: 0.65rem; border-top: 1px solid #eef1f6; }
.nuhtc-run-btn button, .nuhtc-run-btn button.primary {
  background: linear-gradient(168deg, #ffb86b 0%, #fb923c 42%, #ea580c 100%) !important;
  border: none !important;
  color: #fff !important;
  font-weight: 650 !important;
  font-size: 0.95rem !important;
  border-radius: 11px !important;
  padding: 0.55rem 1.25rem !important;
  box-shadow: 0 2px 8px rgba(234, 88, 12, 0.3), inset 0 1px 0 rgba(255,255,255,0.28) !important;
  transition: transform 0.12s ease, box-shadow 0.12s ease;
}
.nuhtc-run-btn button:hover {
  box-shadow: 0 4px 16px rgba(234, 88, 12, 0.38), inset 0 1px 0 rgba(255,255,255,0.32) !important;
  transform: translateY(-1px);
}
.nuhtc-summary textarea {
  border-radius: 11px !important;
  line-height: 1.58 !important;
  font-size: 0.91rem !important;
  padding: 0.72rem 0.9rem !important;
  background: linear-gradient(180deg, #fafbfc 0%, #f1f4f8 100%) !important;
  border: 1px solid #e2e8f0 !important;
  color: #1e293b !important;
  min-height: 5.6rem;
}
.nuhtc-summary label { font-weight: 500 !important; color: #475569 !important; }
"""


# -----------------------------------------------------------------------------
# Helpers & inference
#
# Order: inputs → bbox helpers → model → ``run_inference`` / ``infer`` → ``build_ui``.
# ``infer`` saves a temp PNG, calls ``run_inference`` (MMDet + ``save_result``), then
# ``summarize`` for the status line (same score threshold as the overlay).
# -----------------------------------------------------------------------------
def prepare_image(image) -> Image.Image:
    """Turn Gradio / numpy input into an RGB :class:`PIL.Image.Image`."""
    if isinstance(image, np.ndarray):
        pil_image = Image.fromarray(image)
    else:
        pil_image = image

    if pil_image.mode != "RGB":
        return pil_image.convert("RGB")
    return pil_image


def list_demo_paths(max_images: int = 6):
    """Sorted demo PNG paths as ``[[path], ...]`` for ``gr.Examples``."""
    if not os.path.isdir(DEMO_IMG_DIR):
        return []

    rows = []
    for filename in sorted(os.listdir(DEMO_IMG_DIR)):
        if not filename.lower().endswith(".png"):
            continue
        full_path = os.path.join(DEMO_IMG_DIR, filename)
        rows.append([full_path])
        if len(rows) >= max_images:
            break
    return rows


def extract_bboxes(detection):
    """Take bbox branch from MMDet ``inference_detector`` output (often ``(bboxes, masks)``)."""
    if isinstance(detection, tuple) and len(detection) > 0:
        return detection[0]
    return detection


def count_boxes(bboxes, score_threshold: float) -> int:
    """How many detections pass ``score_threshold`` (MMDet boxes: last column is score)."""
    if bboxes is None:
        return 0
    arr = np.asarray(bboxes)
    if arr.size == 0:
        return 0
    if arr.ndim == 2 and arr.shape[1] >= 5:
        return int(np.sum(arr[:, 4] >= float(score_threshold)))
    return int(len(arr))


def format_summary(counts_dict, score_thr: float = 0.35) -> str:
    """Narrative English summary from per-class counts (PanNuke-style names)."""
    class_map = dict(CLASS_ORDER)
    total = sum(counts_dict.values())

    detected = [
        f"{class_map.get(k, k)} ({v})" for k, v in counts_dict.items() if v > 0
    ]
    not_detected = [class_map.get(k, k) for k, v in counts_dict.items() if v == 0]

    thr = float(score_thr)
    summary = f"Detection Summary (Score >= {thr:.2f}): "
    summary += f"A total of {total} nuclei were identified. "

    if detected:
        summary += f"{', '.join(detected)}. "

    if not_detected:
        summary += f"No {', '.join(not_detected)} nuclei were detected."

    return summary


def summarize(model, detection, score_threshold: float) -> str:
    """Per-class counts at ``score_threshold``, then narrative summary (same filter as overlay)."""
    class_names = list(getattr(model, "CLASSES", ("T", "I", "C", "D", "E")))
    class_to_index = {name: i for i, name in enumerate(class_names)}
    bbox_per_class = extract_bboxes(detection)

    if not isinstance(bbox_per_class, (list, tuple)):
        return f"Could not parse counts (result type: {type(detection).__name__})."

    thr = float(score_threshold)
    counts_dict = OrderedDict()
    for short_key in CLASS_ORDER:
        idx = class_to_index.get(short_key)
        if idx is None or idx >= len(bbox_per_class):
            counts_dict[short_key] = 0
        else:
            counts_dict[short_key] = count_boxes(bbox_per_class[idx], thr)

    return format_summary(counts_dict, score_thr=thr)


@functools.lru_cache(maxsize=1)
def get_model():
    """Load the NuHTC detector once (cached for the lifetime of the process).

    Uses the same config and weights as ``tools/infer.py`` in the Docker image,
    e.g. ``htc_lite_swin_pytorch_fpn_PanNuke_seasaw_CAS.py`` and ``pannuke.pth``.
    Override paths with ``NUHTC_CONFIG`` / ``NUHTC_CHECKPOINT`` if needed.
    """
    device = "cuda" if torch.cuda.is_available() else "cpu"

    cfg = Config.fromfile(CONFIG)
    cfg.work_dir = "./work_dirs"
    cfg = patch_config(cfg)

    model = init_detector(cfg, CHECKPOINT, device=device)
    model.eval()
    model.CLASSES = ("T", "I", "C", "D", "E")

    return model


def run_inference(model, image_path: str, score_threshold: float, output_dir: str):
    """Run MMDet on one image, persist outputs (incl. overlay), return paths + summary text."""
    detection = inference_detector(model, image_path)
    overlay_path = os.path.join(output_dir, os.path.basename(image_path))

    save_result(
        model,
        image_path,
        detection,
        score_thr=score_threshold,
        out_file=overlay_path,
        font_size=6,
        thickness=1,
    )

    summary = summarize(model, detection, score_threshold)
    return overlay_path, summary


def infer(image, score_threshold: float):
    """Gradio handler: load model, run one image, return visualization + status text."""
    if image is None:
        return None, "Please upload an image."

    model = get_model()
    thr = (
        float(score_threshold)
        if score_threshold is not None
        else DEFAULT_SCORE
    )
    rgb_image = prepare_image(image)

    with tempfile.TemporaryDirectory() as tmpdir:
        input_png = os.path.join(tmpdir, "input.png")
        rgb_image.save(input_png)

        overlay_path, summary = run_inference(model, input_png, thr, tmpdir)
        return Image.open(overlay_path).copy(), summary


def build_ui():
    """Build the Gradio UI; returns the :class:`gr.Blocks` root."""
    # Gradio 6+: ``css`` belongs on ``mount_gradio_app`` / ``launch``, not ``Blocks(...)``.
    with gr.Blocks() as blocks:
        with gr.Column(elem_classes=["nuhtc-wrap"]):
            with gr.Column(elem_classes=["nuhtc-header-wrap"]):
                gr.Markdown(
                    "## NuHTC\n\n"
                    "Nuclei instance segmentation on 40× H&E patches."
                )

            with gr.Column(elem_classes=["nuhtc-card"]):
                with gr.Row(equal_height=False, elem_classes=["nuhtc-core-row"]):
                    with gr.Column(scale=5, min_width=180, elem_classes=["nuhtc-img-panel"]):
                        inp = gr.Image(type="pil", label="Input")
                    with gr.Column(scale=2, min_width=80, elem_classes=["nuhtc-arrow-col"]):
                        gr.HTML(_FLOW_DIVIDER_HTML)
                    with gr.Column(scale=5, min_width=180, elem_classes=["nuhtc-img-panel"]):
                        out = gr.Image(type="pil", label="Segmentation")

                demo_rows = list_demo_paths(max_images=6)
                with gr.Column(elem_classes=["nuhtc-demo-section"]):
                    if demo_rows:
                        with gr.Column(elem_classes=["nuhtc-demo-hint-wrap"]):
                            gr.Markdown("Pick a thumbnail to load into Input.")
                    with gr.Row(equal_height=False, elem_classes=["nuhtc-demo-btn-row"]):
                        with gr.Column(scale=6, elem_classes=["nuhtc-demo-strip-wrap"]):
                            if demo_rows:
                                with gr.Column(elem_classes=["nuhtc-demo-strip"]):
                                    gr.Examples(
                                        examples=demo_rows,
                                        inputs=inp,
                                        cache_examples=False,
                                        label=" ",
                                    )
                        with gr.Column(scale=2, min_width=140, elem_classes=["nuhtc-btn-col"]):
                            btn = gr.Button(
                                "Run inference",
                                variant="primary",
                                elem_classes=["nuhtc-run-btn"],
                            )

            with gr.Column(elem_classes=["nuhtc-card"]):
                with gr.Row(equal_height=False, elem_classes=["nuhtc-bottom-row"]):
                    with gr.Column(scale=1, elem_classes=["nuhtc-slider-col"]):
                        score = gr.Slider(
                            0.05,
                            0.95,
                            value=DEFAULT_SCORE,
                            step=0.05,
                            label="Score threshold",
                        )
                with gr.Column(elem_classes=["nuhtc-summary-row"]):
                    log = gr.Textbox(
                        label="Summary",
                        lines=4,
                        max_lines=12,
                        elem_classes=["nuhtc-summary"],
                    )
            btn.click(infer, [inp, score], [out, log])
    return blocks


# -----------------------------------------------------------------------------
# Gradio UI + ASGI
# -----------------------------------------------------------------------------
demo = build_ui()
demo.queue()
app = gr.mount_gradio_app(FastAPI(), demo, path="/", css=_THUMB_CSS)

if __name__ == "__main__":

    port = int(os.environ.get("PORT", "7860"))
    uvicorn.run(app, host="0.0.0.0", port=port)