viewer/colormaps.py

"""
Anatomy label tables and color assignments for the legend overlay.

CT labels are loaded directly from the upstream config so we stay in sync.
A deterministic golden-ratio HSV walk maps each label ID to a distinguishable RGB.
"""
from __future__ import annotations

import colorsys
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
UPSTREAM_LABEL_DICT = ROOT / "repos" / "NV-Generate-CTMR" / "configs" / "label_dict.json"


def _load_ct_labels() -> dict[int, str]:
    if not UPSTREAM_LABEL_DICT.exists():
        return {}
    raw: dict[str, int] = json.loads(UPSTREAM_LABEL_DICT.read_text())
    return {int(v): k for k, v in raw.items() if not k.startswith("dummy")}


def _color_for(idx: int) -> tuple[int, int, int]:
    """Deterministic high-contrast color from an integer ID."""
    golden = 0.61803398875
    h = (idx * golden) % 1.0
    s = 0.55 + 0.35 * ((idx * 7) % 5) / 4.0
    v = 0.70 + 0.25 * ((idx * 11) % 4) / 3.0
    r, g, b = colorsys.hsv_to_rgb(h, s, v)
    return int(r * 255), int(g * 255), int(b * 255)


CT_LABELS_BY_ID: dict[int, str] = _load_ct_labels()


def color_table(label_ids: list[int]) -> list[tuple[int, int, int, int]]:
    """RGBA palette for niivue label overlay — index 0 is background (transparent)."""
    if not label_ids:
        return []
    max_id = max(label_ids)
    palette: list[tuple[int, int, int, int]] = [(0, 0, 0, 0)]
    for i in range(1, max_id + 1):
        if i in label_ids:
            r, g, b = _color_for(i)
            palette.append((r, g, b, 200))
        else:
            palette.append((0, 0, 0, 0))
    return palette


def legend_html(used_labels: dict[int, str]) -> str:
    """Compact horizontal legend showing only labels present in the generated mask."""
    if not used_labels:
        return ""
    rows: list[str] = []
    for lbl, name in sorted(used_labels.items()):
        r, g, b = _color_for(lbl)
        rows.append(
            f'<div class="nv-legend-chip">'
            f'<span class="nv-swatch" style="background:rgb({r},{g},{b})"></span>'
            f'<span class="nv-label">{name}</span></div>'
        )
    count = len(used_labels)
    return (
        '<div class="nv-legend">'
        f'<div class="nv-legend-title">Classes present in mask <span class="nv-legend-count">{count}</span></div>'
        '<div class="nv-legend-grid">' + "".join(rows) + "</div>"
        "</div>"
    )