tuning/render.py

"""Shared step-visual renderers — the single source of truth for how a pipeline stage
is drawn, used by both the Gradio app and the static HTML export.

Lifted from `scripts/audit_pipeline.py` (b64/overlay/draw_outlines/class_panel/
lidar_panel) so the interactive and shareable views render identically, plus an
`iou_overlay` for scoring a recipe mask against a hand-drawn gold mask.
"""

from __future__ import annotations

import base64
import io

import numpy as np
from PIL import Image, ImageDraw

# ADE20K names for the EoMT cascade primary (outdoor-relevant subset).
ADE_NAMES = {0: "wall", 1: "building", 2: "sky", 4: "tree", 6: "road", 9: "grass",
             11: "sidewalk", 13: "earth", 17: "plant", 21: "water", 20: "car", 25: "?",
             29: "field", 46: "sand", 52: "path", 94: "land"}
# Empirical meanings for the incumbent mask2former's generic LABEL_0..7.
M2F_NAMES = {0: "background", 1: "open (grass+dirt)", 2: "street-edge band", 3: "pavement",
             4: "canopy", 6: '"water" (fires on flat turf)', 7: "roofs (as cropland)"}

PALETTE = [(80, 200, 60), (0, 110, 40), (220, 60, 60), (150, 110, 70), (235, 220, 120),
           (170, 120, 40), (60, 130, 235), (120, 120, 130), (200, 60, 200), (230, 130, 30),
           (90, 200, 200), (200, 120, 200), (255, 180, 40)]


def b64(img: Image.Image, max_w: int = 900, quality: int = 82) -> str:
    """A data-URI JPEG, downscaled to `max_w` — embeds directly in shareable HTML."""
    if img.width > max_w:
        img = img.resize((max_w, round(img.height * max_w / img.width)), Image.LANCZOS)
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "JPEG", quality=quality)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def overlay(base: Image.Image, mask: np.ndarray, color, alpha: float = 0.55) -> Image.Image:
    arr = np.asarray(base.convert("RGB")).astype(np.float32)
    arr[mask] = arr[mask] * (1 - alpha) + np.array(color, np.float32) * alpha
    return Image.fromarray(arr.astype(np.uint8))


def draw_outlines(base: Image.Image, outlines, color, width: int = 4) -> Image.Image:
    """`outlines` is a list of (xs, ys) coord-list pairs (polygon_pixel_outlines)."""
    img = base.convert("RGB").copy()
    d = ImageDraw.Draw(img)
    for xs, ys in outlines:
        pts = [(float(x), float(y)) for x, y in zip(xs, ys, strict=False)]
        if len(pts) > 1:
            d.line(pts + [pts[0]], fill=color, width=width)
    return img


def class_panel(base: Image.Image, pred: np.ndarray, names: dict,
                min_frac: float = 0.003) -> tuple[Image.Image, list]:
    """Per-class colored overlay + legend chips (name, %, color) for present classes."""
    arr = np.asarray(base.convert("RGB")).astype(np.float32) * 0.45
    legend = []
    codes = [c for c in np.unique(pred) if (pred == c).sum() / pred.size >= min_frac]
    for i, c in enumerate(sorted(codes, key=lambda c: -(pred == c).sum())):
        col = PALETTE[i % len(PALETTE)]
        arr[pred == c] += np.array(col, np.float32) * 0.55
        legend.append({"name": names.get(int(c), f"class {c}"),
                       "color": "#{:02x}{:02x}{:02x}".format(*col),
                       "pct": round(float((pred == c).sum() / pred.size * 100), 1)})
    return Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8)), legend


def lidar_panel(base: Image.Image, viz: dict) -> Image.Image:
    """Ground points colored: green = counted lawn, red = removed (hardscape/roof)."""
    img = base.convert("RGB").copy()
    gpx, gpy = viz.get("ground_px"), viz.get("ground_py")
    if gpx is None:
        return img
    d = ImageDraw.Draw(img)
    is_lawn = viz.get("lawn_mask")
    for k in range(len(gpx)):
        x, y = float(gpx[k]), float(gpy[k])
        col = (60, 220, 60) if (is_lawn is not None and is_lawn[k]) else (235, 60, 60)
        d.ellipse([x - 2, y - 2, x + 2, y + 2], fill=col)
    return img


def iou_overlay(base: Image.Image, recipe_mask: np.ndarray, gold_mask: np.ndarray,
                alpha: float = 0.5) -> Image.Image:
    """True-positive (green) / false-positive (red, recipe-only) / false-negative
    (blue, gold-only) overlay — a visual read on where a recipe misses the truth."""
    arr = np.asarray(base.convert("RGB")).astype(np.float32)
    tp = recipe_mask & gold_mask
    fp = recipe_mask & ~gold_mask
    fn = ~recipe_mask & gold_mask
    for mask, color in ((tp, (60, 220, 60)), (fp, (235, 60, 60)), (fn, (70, 130, 235))):
        arr[mask] = arr[mask] * (1 - alpha) + np.array(color, np.float32) * alpha
    return Image.fromarray(arr.astype(np.uint8))