samgeo3-lab/scripts/merge_multi_results.py

"""Merge multi-prompt SamGeo3 results into one overview (overlay + grid + HTML).

Works on the output directory from multi_prompt_segment.py
(e.g. output/dji_0016_compact with *_mask.png files).

Usage (from project root, venv active):
    python scripts/merge_multi_results.py --result-dir output/dji_0016_compact
    python scripts/merge_multi_results.py --result-dir output/dji_0016_compact --image data/DJI_....JPG
    python scripts/merge_multi_results.py --result-dir output/dji_0016_compact --max-side 2048
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image, ImageDraw, ImageFont

ROOT = Path(__file__).resolve().parents[1]

# Distinct RGB colors for class overlays (cycled if more prompts)
PALETTE = [
    (230, 25, 75),  # red
    (60, 180, 75),  # green
    (0, 130, 200),  # blue
    (245, 130, 48),  # orange
    (145, 30, 180),  # purple
    (70, 240, 240),  # cyan
    (240, 50, 230),  # magenta
    (210, 245, 60),  # lime
    (250, 190, 212),  # pink
    (0, 128, 128),  # teal
    (220, 190, 255),  # lavender
    (170, 110, 40),  # brown
    (255, 250, 200),  # beige
    (128, 0, 0),  # maroon
    (170, 255, 195),  # mint
    (128, 128, 0),  # olive
    (255, 215, 180),  # apricot
    (0, 0, 128),  # navy
    (128, 128, 128),  # gray
    (255, 225, 25),  # yellow
]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Merge multi-prompt segmentation results")
    p.add_argument(
        "--result-dir",
        type=Path,
        required=True,
        help="Directory with *_mask.png (and optional summary.json)",
    )
    p.add_argument(
        "--image",
        type=Path,
        default=None,
        help="Original image (else from summary.json or data/ fallback)",
    )
    p.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Where to write merged outputs (default: <result-dir>/merged)",
    )
    p.add_argument(
        "--max-side",
        type=int,
        default=2560,
        help="Max width/height of overview images (keeps memory/file size down)",
    )
    p.add_argument(
        "--alpha",
        type=float,
        default=0.45,
        help="Mask overlay opacity 0~1",
    )
    p.add_argument(
        "--grid-cols",
        type=int,
        default=4,
        help="Columns in thumbnail grid",
    )
    p.add_argument(
        "--thumb-size",
        type=int,
        default=480,
        help="Thumbnail max side for grid cells",
    )
    p.add_argument(
        "--no-grid",
        action="store_true",
        help="Skip per-class grid image",
    )
    p.add_argument(
        "--no-html",
        action="store_true",
        help="Skip HTML report",
    )
    return p.parse_args()


def load_summary(result_dir: Path) -> dict[str, Any] | None:
    path = result_dir / "summary.json"
    if not path.is_file():
        return None
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def _common_prefix(strings: list[str]) -> str:
    if not strings:
        return ""
    prefix = strings[0]
    for s in strings[1:]:
        while not s.startswith(prefix) and prefix:
            prefix = prefix[:-1]
        if not prefix:
            break
    return prefix


def discover_masks(result_dir: Path) -> list[dict[str, Any]]:
    """Return list of {prompt, mask_path, ann_path, n_from_scores} from files."""
    masks = sorted(result_dir.glob("*_mask.png"))
    bases = [mp.name[: -len("_mask.png")] for mp in masks]
    # e.g. common "DJI_..._0016_" then remainder is prompt tag (may contain _)
    cpref = _common_prefix(bases)
    # prefer cut at last underscore of common prefix so prompt is clean
    if cpref and not cpref.endswith("_"):
        # trim to last underscore so we don't eat prompt chars
        li = cpref.rfind("_")
        cpref = cpref[: li + 1] if li >= 0 else ""

    items: list[dict[str, Any]] = []
    for mp, base in zip(masks, bases):
        tag = base[len(cpref) :] if cpref and base.startswith(cpref) else base
        if not tag:
            # fallback: last underscore segment only
            tag = base.rsplit("_", 1)[-1]
        prompt = tag.replace("_", " ")
        ann = result_dir / f"{base}_ann.png"
        scores = result_dir / f"{base}_scores.npy"
        n_obj = None
        if scores.is_file():
            try:
                n_obj = int(np.load(scores).shape[0])
            except Exception:
                n_obj = None
        items.append(
            {
                "prompt": prompt,
                "tag": tag,
                "base": base,
                "mask_path": mp,
                "ann_path": ann if ann.is_file() else None,
                "scores_path": scores if scores.is_file() else None,
                "n_objects": n_obj,
            }
        )
    return items


def items_from_summary(summary: dict[str, Any], result_dir: Path) -> list[dict[str, Any]]:
    items: list[dict[str, Any]] = []
    for r in summary.get("results", []):
        n = int(r.get("n_objects") or 0)
        if n <= 0 and not r.get("mask"):
            continue
        mask_s = r.get("mask")
        mask_path = Path(mask_s) if mask_s else None
        if mask_path is None or not mask_path.is_file():
            # try local name
            prompt = r.get("prompt", "")
            tag = "".join(c if c.isalnum() or c in "-_" else "_" for c in prompt)
            cands = list(result_dir.glob(f"*_{tag}_mask.png"))
            mask_path = cands[0] if cands else None
        if mask_path is None or not mask_path.is_file():
            continue
        ann_s = r.get("ann")
        ann_path = Path(ann_s) if ann_s and Path(ann_s).is_file() else None
        items.append(
            {
                "prompt": r.get("prompt", mask_path.stem),
                "tag": mask_path.stem.replace("_mask", "").split("_")[-1],
                "base": mask_path.name[: -len("_mask.png")],
                "mask_path": mask_path,
                "ann_path": ann_path,
                "scores_path": Path(r["scores_npy"])
                if r.get("scores_npy") and Path(r["scores_npy"]).is_file()
                else None,
                "n_objects": n,
                "scores_min": r.get("scores_min"),
                "scores_max": r.get("scores_max"),
            }
        )
    return items


def resolve_base_image(
    args_image: Path | None,
    summary: dict[str, Any] | None,
    result_dir: Path,
) -> Path:
    cands: list[Path] = []
    if args_image:
        cands.append(args_image)
    if summary and summary.get("image"):
        cands.append(Path(summary["image"]))
    # common lab defaults
    cands.append(ROOT / "data" / "DJI_20260306100802_0016.JPG")
    cands.append(
        Path(r"D:\MYCLAUDE_PROJECT\segment-geospatial\sample\DJI_20260306100802_0016.JPG")
    )
    for c in cands:
        if c and Path(c).is_file():
            return Path(c).resolve()
    raise FileNotFoundError(
        "Base image not found. Pass --image. Tried:\n  "
        + "\n  ".join(str(c) for c in cands)
    )


def resize_max(im: Image.Image, max_side: int) -> Image.Image:
    w, h = im.size
    m = max(w, h)
    if m <= max_side:
        return im
    scale = max_side / m
    nw, nh = int(w * scale), int(h * scale)
    return im.resize((nw, nh), Image.Resampling.BILINEAR)


def load_mask_bool(path: Path, size: tuple[int, int]) -> np.ndarray:
    m = Image.open(path)
    if m.mode not in ("L", "I", "I;16", "P"):
        m = m.convert("L")
    if m.size != size:
        m = m.resize(size, Image.Resampling.NEAREST)
    arr = np.array(m)
    return arr > 0


def try_font(size: int) -> ImageFont.ImageFont:
    for name in (
        "C:/Windows/Fonts/malgun.ttf",
        "C:/Windows/Fonts/segoeui.ttf",
        "C:/Windows/Fonts/arial.ttf",
    ):
        if Path(name).is_file():
            try:
                return ImageFont.truetype(name, size)
            except Exception:
                pass
    return ImageFont.load_default()


def build_overlay(
    base_rgb: Image.Image,
    items: list[dict[str, Any]],
    alpha: float,
) -> tuple[Image.Image, list[dict[str, Any]]]:
    base = np.asarray(base_rgb.convert("RGB"), dtype=np.float32)
    h, w = base.shape[:2]
    out = base.copy()
    legend: list[dict[str, Any]] = []

    for i, it in enumerate(items):
        color = PALETTE[i % len(PALETTE)]
        try:
            mask = load_mask_bool(it["mask_path"], (w, h))
        except Exception as e:
            print(f"  [WARN] skip mask {it['mask_path'].name}: {e}")
            continue
        pix = int(mask.sum())
        if pix == 0:
            continue
        c = np.array(color, dtype=np.float32)
        out[mask] = out[mask] * (1.0 - alpha) + c * alpha
        legend.append(
            {
                "prompt": it["prompt"],
                "color": color,
                "n_objects": it.get("n_objects"),
                "pixels": pix,
                "scores_min": it.get("scores_min"),
                "scores_max": it.get("scores_max"),
            }
        )

    blended = Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))
    return blended, legend


def draw_legend_panel(
    legend: list[dict[str, Any]],
    width: int = 420,
    row_h: int = 36,
) -> Image.Image:
    font = try_font(16)
    title_font = try_font(20)
    n = max(len(legend), 1)
    height = 56 + n * row_h + 16
    panel = Image.new("RGB", (width, height), (24, 28, 36))
    draw = ImageDraw.Draw(panel)
    draw.text((16, 14), "Merged classes (prompt)", fill=(240, 244, 248), font=title_font)
    y = 52
    for ent in legend:
        r, g, b = ent["color"]
        draw.rectangle([16, y + 4, 40, y + 28], fill=(r, g, b))
        n_obj = ent.get("n_objects")
        n_txt = f"n={n_obj}" if n_obj is not None else f"px={ent['pixels']}"
        smin, smax = ent.get("scores_min"), ent.get("scores_max")
        if smin is not None and smax is not None:
            score_txt = f"  conf {smin:.2f}~{smax:.2f}"
        else:
            score_txt = ""
        label = f"{ent['prompt']}  ({n_txt}{score_txt})"
        draw.text((52, y + 6), label, fill=(220, 226, 234), font=font)
        y += row_h
    return panel


def compose_with_legend(overlay: Image.Image, legend_panel: Image.Image) -> Image.Image:
    gap = 16
    w = overlay.width + gap + legend_panel.width
    h = max(overlay.height, legend_panel.height)
    canvas = Image.new("RGB", (w, h), (18, 20, 26))
    canvas.paste(overlay, (0, 0))
    canvas.paste(legend_panel, (overlay.width + gap, 0))
    return canvas


def build_grid(
    base: Image.Image,
    items: list[dict[str, Any]],
    colors: list[tuple[int, int, int]],
    cols: int,
    thumb: int,
    alpha: float,
) -> Image.Image:
    font = try_font(18)
    cells: list[Image.Image] = []
    for i, it in enumerate(items):
        color = colors[i % len(colors)]
        # small overlay per class
        b = resize_max(base.copy(), thumb)
        arr = np.asarray(b.convert("RGB"), dtype=np.float32)
        h, w = arr.shape[:2]
        try:
            mask = load_mask_bool(it["mask_path"], (w, h))
        except Exception:
            mask = np.zeros((h, w), dtype=bool)
        c = np.array(color, dtype=np.float32)
        if mask.any():
            arr[mask] = arr[mask] * (1.0 - alpha) + c * alpha
        cell = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
        # title bar
        bar_h = 32
        framed = Image.new("RGB", (cell.width, cell.height + bar_h), (30, 34, 42))
        framed.paste(cell, (0, bar_h))
        draw = ImageDraw.Draw(framed)
        draw.rectangle([0, 0, framed.width, bar_h], fill=(30, 34, 42))
        draw.rectangle([8, 8, 24, 24], fill=color)
        n = it.get("n_objects")
        title = f"{it['prompt']}" + (f"  (n={n})" if n is not None else "")
        draw.text((32, 6), title[:40], fill=(235, 240, 245), font=font)
        cells.append(framed)

    if not cells:
        return Image.new("RGB", (thumb, thumb), (40, 40, 40))

    cols = max(1, cols)
    rows = (len(cells) + cols - 1) // cols
    cw = max(c.width for c in cells)
    ch = max(c.height for c in cells)
    pad = 8
    grid_w = cols * cw + (cols + 1) * pad
    grid_h = rows * ch + (rows + 1) * pad
    grid = Image.new("RGB", (grid_w, grid_h), (18, 20, 26))
    for i, cell in enumerate(cells):
        r, c = divmod(i, cols)
        x = pad + c * (cw + pad)
        y = pad + r * (ch + pad)
        grid.paste(cell, (x, y))
    return grid


def enrich_item_scores(items: list[dict[str, Any]]) -> None:
    """Fill scores_min/max/mean/n_objects from *_scores.npy when present."""
    for it in items:
        sp = it.get("scores_path")
        if sp is None or not Path(sp).is_file():
            continue
        try:
            arr = np.load(sp).astype(np.float64).ravel()
            if arr.size == 0:
                continue
            it["n_objects"] = int(arr.size)
            it["scores_min"] = float(arr.min())
            it["scores_max"] = float(arr.max())
            it["scores_mean"] = float(arr.mean())
        except Exception:
            continue


def mask_to_rle(mask: np.ndarray) -> list[int]:
    """Run-length encode binary mask (row-major). [start, length, start, length, ...].

    Used so the browser can hit-test without canvas getImageData (file:// safe).
    """
    flat = np.ascontiguousarray(mask.astype(bool).ravel())
    n = int(flat.size)
    if n == 0:
        return []
    # transitions where value changes
    # find runs of True
    rle: list[int] = []
    i = 0
    while i < n:
        if flat[i]:
            start = i
            while i < n and flat[i]:
                i += 1
            rle.append(int(start))
            rle.append(int(i - start))
        else:
            i += 1
    return rle


def export_hitmasks(
    out_dir: Path,
    ordered_items: list[dict[str, Any]],
    legend: list[dict[str, Any]],
    size: tuple[int, int],
) -> list[dict[str, Any]]:
    """Build interactive class meta with RLE hit data (file:// / no CORS needed)."""
    w, h = size
    classes_meta: list[dict[str, Any]] = []
    for i, ent in enumerate(legend):
        it = None
        for cand in ordered_items:
            if cand["prompt"] == ent["prompt"] or cand["prompt"].replace(
                "_", " "
            ) == ent["prompt"]:
                it = cand
                break
        if it is None and i < len(ordered_items):
            it = ordered_items[i]
        if it is None:
            continue

        mask = load_mask_bool(it["mask_path"], (w, h))
        rle = mask_to_rle(mask)

        r, g, b = ent["color"]
        smin = ent.get("scores_min")
        if smin is None:
            smin = it.get("scores_min")
        smax = ent.get("scores_max")
        if smax is None:
            smax = it.get("scores_max")
        smean = it.get("scores_mean")
        if smean is None and smin is not None and smax is not None:
            smean = (float(smin) + float(smax)) / 2.0

        classes_meta.append(
            {
                "id": i,
                "prompt": ent["prompt"],
                "color": [int(r), int(g), int(b)],
                "n_objects": ent.get("n_objects")
                if ent.get("n_objects") is not None
                else it.get("n_objects"),
                "pixels": int(ent.get("pixels") or int(mask.sum())),
                "scores_min": smin,
                "scores_max": smax,
                "scores_mean": smean,
                "rle": rle,
            }
        )

    meta = {
        "width": w,
        "height": h,
        "overlay": "combined_overlay.png",
        "hit_mode": "rle",
        "classes": classes_meta,
    }
    meta_path = out_dir / "interactive_meta.json"
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)
    return classes_meta


def write_html(
    out_path: Path,
    overlay_name: str,
    grid_name: str | None,
    legend: list[dict[str, Any]],
    meta: dict[str, Any],
    interactive_meta: dict[str, Any] | None = None,
) -> None:
    """Write interactive index.html (click image → class list + scores)."""
    # Embed meta so file:// works without fetch CORS issues
    imeta = interactive_meta
    if imeta is None:
        imeta_path = out_path.parent / "interactive_meta.json"
        if imeta_path.is_file():
            with open(imeta_path, encoding="utf-8") as f:
                imeta = json.load(f)
        else:
            imeta = {
                "width": 0,
                "height": 0,
                "overlay": overlay_name,
                "classes": [],
            }
    imeta_js = json.dumps(imeta, ensure_ascii=False)

    # Static fallback rows still useful if JS fails
    rows = []
    for ent in legend:
        r, g, b = ent["color"]
        n = ent.get("n_objects")
        smin, smax = ent.get("scores_min"), ent.get("scores_max")
        score_cell = (
            f"{smin:.3f}~{smax:.3f}" if smin is not None and smax is not None else "-"
        )
        rows.append(
            f"<tr data-prompt=\"{ent['prompt']}\">"
            f"<td><span class='swatch' style='background:rgb({r},{g},{b})'></span></td>"
            f"<td class='prompt-cell'>{ent['prompt']}</td>"
            f"<td>{n if n is not None else '-'}</td>"
            f"<td>{score_cell}</td>"
            f"<td>{ent['pixels']:,}</td></tr>"
        )
    grid_block = (
        f'<details class="grid-details"><summary>Per-class grid</summary>'
        f'<img src="{grid_name}" alt="grid" class="grid-img"/></details>'
        if grid_name
        else ""
    )
    n_cls = len(legend)
    html = f"""<!DOCTYPE html>
<html lang="ko">
<head>
<meta charset="utf-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1"/>
<title>Merged multi-prompt — interactive</title>
<style>
:root {{
  --bg:#0f1419; --panel:#1a2332; --panel2:#243044; --text:#e7ecf3; --muted:#9aa8bc;
  --accent:#3d9cf0; --accent2:#5eead4; --border:#2d3a4f; --hit:#fbbf24;
}}
* {{ box-sizing:border-box; }}
body {{
  margin:0; font-family:Segoe UI, Malgun Gothic, sans-serif;
  background:var(--bg); color:var(--text); min-height:100vh;
}}
header {{
  padding:.85rem 1.25rem; border-bottom:1px solid var(--border);
  background:linear-gradient(90deg,#1e3a5f,#0f766e);
}}
header h1 {{ margin:0; font-size:1.15rem; font-weight:600; }}
header p {{ margin:.25rem 0 0; font-size:.85rem; color:rgba(255,255,255,.88); }}
.layout {{
  display:grid; grid-template-columns:1fr minmax(300px,380px);
  gap:0; min-height:calc(100vh - 70px);
}}
@media (max-width:900px) {{
  .layout {{ grid-template-columns:1fr; }}
}}
.stage {{
  padding:1rem; overflow:auto; border-right:1px solid var(--border);
  background:#0b1020;
}}
.canvas-wrap {{
  position:relative; display:inline-block; max-width:100%;
  border:1px solid var(--border); border-radius:8px; overflow:hidden;
  cursor:crosshair; background:#000;
}}
.canvas-wrap canvas {{ display:block; max-width:100%; height:auto; vertical-align:top; }}
.hint {{
  margin:.6rem 0 0; font-size:.82rem; color:var(--muted);
}}
.side {{
  padding:1rem; background:var(--panel); overflow:auto;
  display:flex; flex-direction:column; gap:.85rem;
}}
.side h2 {{
  margin:0; font-size:.95rem; color:var(--accent2);
  text-transform:uppercase; letter-spacing:.04em;
}}
.pick-meta {{
  font-size:.82rem; color:var(--muted); font-family:Consolas, monospace;
}}
.hit-list {{ list-style:none; margin:0; padding:0; }}
.hit-list li {{
  display:grid; grid-template-columns:18px 1fr auto;
  gap:.5rem; align-items:center;
  padding:.55rem .6rem; margin-bottom:.4rem;
  background:var(--panel2); border:1px solid var(--border); border-radius:8px;
  border-left:4px solid #666;
}}
.hit-list li.top {{
  box-shadow:0 0 0 1px var(--hit);
}}
.hit-list .sw {{
  width:14px; height:14px; border-radius:3px; display:inline-block;
}}
.hit-list .name {{ font-weight:600; font-size:.9rem; }}
.hit-list .score {{
  font-family:Consolas, monospace; font-size:.8rem; color:var(--accent2);
  text-align:right; white-space:nowrap;
}}
.hit-list .sub {{
  grid-column:2 / -1; font-size:.75rem; color:var(--muted);
}}
.empty-hit {{ color:var(--muted); font-size:.9rem; padding:.5rem 0; }}
.rank-badge {{
  font-size:.7rem; background:#0b1220; border:1px solid var(--border);
  border-radius:999px; padding:.1rem .45rem; color:var(--hit); margin-left:.35rem;
}}
.all-table {{ width:100%; border-collapse:collapse; font-size:.78rem; }}
.all-table th, .all-table td {{
  border:1px solid var(--border); padding:.35rem .4rem; text-align:left;
}}
.all-table th {{ background:#121a28; color:var(--accent2); position:sticky; top:0; }}
.all-table tr.hit-row {{ background:rgba(251,191,36,.12); }}
.all-table tr.hit-row.top {{ background:rgba(251,191,36,.22); }}
.swatch {{
  display:inline-block; width:12px; height:12px; border-radius:2px; vertical-align:middle;
}}
.status {{ font-size:.8rem; color:var(--muted); }}
.status.ok {{ color:var(--accent2); }}
.status.err {{ color:#f87171; }}
.grid-details {{ margin-top:1rem; color:var(--muted); font-size:.85rem; }}
.grid-img {{ max-width:100%; border:1px solid var(--border); border-radius:8px; margin-top:.5rem; }}
.priority-note {{
  font-size:.8rem; color:var(--muted); line-height:1.45;
  border-left:3px solid var(--accent); padding-left:.6rem;
}}
</style>
</head>
<body>
<header>
  <h1>Merged multi-prompt — click to inspect</h1>
  <p>이미지를 클릭하면 해당 픽셀에 겹친 클래스(프롬프트)와 confidence 점수가 오른쪽에 표시됩니다. 겹침 시 점수 높은 순으로 정렬됩니다.</p>
</header>
<div class="layout">
  <div class="stage">
    <div class="canvas-wrap" id="wrap">
      <canvas id="view"></canvas>
    </div>
    <p class="hint">클릭: 픽셀 조회 · Shift+클릭: 선택 유지(누적) · Esc: 초기화 · 클래스 {n_cls}개</p>
    <p class="status" id="status">loading…</p>
    {grid_block}
  </div>
  <aside class="side">
    <div>
      <h2>Selection</h2>
      <div class="pick-meta" id="pickMeta">클릭하여 조회</div>
      <ul class="hit-list" id="hitList"></ul>
      <div class="empty-hit" id="emptyHit" hidden>이 픽셀에 마스크 없음</div>
      <p class="priority-note">
        겹침 우선순위 참고: 목록은 <strong>scores_max</strong> 내림차순입니다.
        점수가 비슷하면 픽셀 면적(n_objects)과 장면 문맥으로 판단하세요.
      </p>
    </div>
    <div>
      <h2>All classes</h2>
      <div style="max-height:45vh; overflow:auto;">
        <table class="all-table" id="allTable">
          <thead>
            <tr><th></th><th>prompt</th><th>n</th><th>score max</th><th>px</th></tr>
          </thead>
          <tbody>
            {''.join(rows)}
          </tbody>
        </table>
      </div>
    </div>
  </aside>
</div>
<script>
(async function () {{
  /* file:// safe: NO getImageData (avoids "canvas tainted by cross-origin data")
     Hit test uses RLE from merge step, not mask PNGs. */
  const statusEl = document.getElementById('status');
  const canvas = document.getElementById('view');
  const ctx = canvas.getContext('2d');
  const hitList = document.getElementById('hitList');
  const emptyHit = document.getElementById('emptyHit');
  const pickMeta = document.getElementById('pickMeta');

  let meta = null;
  let overlayImg = null;
  let classes = [];
  let lastHits = [];
  let marker = null;

  function setStatus(msg, ok) {{
    statusEl.textContent = msg;
    statusEl.className = 'status ' + (ok === true ? 'ok' : ok === false ? 'err' : '');
  }}

  function loadImage(src) {{
    return new Promise((resolve, reject) => {{
      const im = new Image();
      im.onload = () => resolve(im);
      im.onerror = () => reject(new Error('fail load ' + src));
      im.src = src;
    }});
  }}

  /** RLE: [start, length, ...] of True runs on flat row-major index */
  function rleHas(rle, idx) {{
    if (!rle || !rle.length) return false;
    for (let j = 0; j < rle.length; j += 2) {{
      const s = rle[j];
      const len = rle[j + 1];
      if (idx < s) return false;
      if (idx < s + len) return true;
    }}
    return false;
  }}

  function draw() {{
    if (!overlayImg || !meta) return;
    const w = meta.width, h = meta.height;
    if (canvas.width !== w || canvas.height !== h) {{
      canvas.width = w;
      canvas.height = h;
    }}
    canvas.style.width = '100%';
    canvas.style.height = 'auto';
    // drawImage only — never getImageData (taint-safe display)
    ctx.drawImage(overlayImg, 0, 0, w, h);

    if (marker) {{
      const {{x, y}} = marker;
      ctx.save();
      ctx.strokeStyle = '#fbbf24';
      ctx.fillStyle = 'rgba(251,191,36,0.35)';
      ctx.lineWidth = 2;
      ctx.beginPath();
      ctx.arc(x + 0.5, y + 0.5, 8, 0, Math.PI * 2);
      ctx.fill();
      ctx.stroke();
      ctx.beginPath();
      ctx.moveTo(x - 16, y + 0.5); ctx.lineTo(x + 16, y + 0.5);
      ctx.moveTo(x + 0.5, y - 16); ctx.lineTo(x + 0.5, y + 16);
      ctx.stroke();
      // stacked color dots for each hit
      lastHits.forEach((h, i) => {{
        ctx.beginPath();
        ctx.fillStyle = `rgb(${{h.color[0]}},${{h.color[1]}},${{h.color[2]}})`;
        ctx.strokeStyle = '#fff';
        ctx.lineWidth = 1;
        ctx.arc(x + 14 + i * 10, y - 14, 5, 0, Math.PI * 2);
        ctx.fill();
        ctx.stroke();
      }});
      ctx.restore();
    }}
  }}

  function eventToNatural(ev) {{
    const rect = canvas.getBoundingClientRect();
    const sx = canvas.width / rect.width;
    const sy = canvas.height / rect.height;
    const x = Math.floor((ev.clientX - rect.left) * sx);
    const y = Math.floor((ev.clientY - rect.top) * sy);
    return {{
      x: Math.max(0, Math.min(canvas.width - 1, x)),
      y: Math.max(0, Math.min(canvas.height - 1, y)),
    }};
  }}

  function queryPixel(x, y, accumulate) {{
    const idx = y * meta.width + x;
    const hits = [];
    for (const cls of classes) {{
      if (rleHas(cls.rle, idx)) {{
        hits.push({{
          id: cls.id,
          prompt: cls.prompt,
          color: cls.color,
          scores_min: cls.scores_min,
          scores_max: cls.scores_max,
          scores_mean: cls.scores_mean,
          n_objects: cls.n_objects,
          pixels: cls.pixels,
          rank_score: (cls.scores_max != null ? Number(cls.scores_max) : -1),
        }});
      }}
    }}
    hits.sort((a, b) => b.rank_score - a.rank_score || (b.pixels || 0) - (a.pixels || 0));

    if (accumulate && lastHits.length) {{
      const map = new Map(lastHits.map(h => [h.id, h]));
      for (const h of hits) map.set(h.id, h);
      lastHits = Array.from(map.values()).sort(
        (a, b) => b.rank_score - a.rank_score || (b.pixels || 0) - (a.pixels || 0)
      );
    }} else {{
      lastHits = hits;
    }}
    marker = {{ x, y }};
    renderHits();
    draw();
  }}

  function fmtScore(h) {{
    if (h.scores_max == null && h.scores_min == null) return 'score n/a';
    const a = h.scores_min != null ? Number(h.scores_min).toFixed(3) : '?';
    const b = h.scores_max != null ? Number(h.scores_max).toFixed(3) : '?';
    const m = h.scores_mean != null ? Number(h.scores_mean).toFixed(3) : null;
    return m != null ? `max ${{b}} · mean ${{m}}` : `${{a}} ~ ${{b}}`;
  }}

  function renderHits() {{
    hitList.innerHTML = '';
    document.querySelectorAll('#allTable tr').forEach(tr => tr.classList.remove('hit-row', 'top'));

    if (!lastHits.length) {{
      emptyHit.hidden = false;
      pickMeta.textContent = marker
        ? `pixel (${{marker.x}}, ${{marker.y}}) — 0 classes`
        : '클릭하여 조회';
      return;
    }}
    emptyHit.hidden = true;
    pickMeta.textContent = `pixel (${{marker.x}}, ${{marker.y}}) — ${{lastHits.length}} class(es) overlapping`;

    lastHits.forEach((h, i) => {{
      const li = document.createElement('li');
      if (i === 0) li.classList.add('top');
      li.style.borderLeftColor = `rgb(${{h.color[0]}},${{h.color[1]}},${{h.color[2]}})`;
      const rank = i === 0 ? '<span class="rank-badge">priority #1</span>' : `<span class="rank-badge">#${{i+1}}</span>`;
      li.innerHTML = `
        <span class="sw" style="background:rgb(${{h.color[0]}},${{h.color[1]}},${{h.color[2]}})"></span>
        <div><div class="name">${{h.prompt}}${{rank}}</div></div>
        <div class="score">${{fmtScore(h)}}</div>
        <div class="sub">n_objects=${{h.n_objects ?? '-'}} · pixels=${{(h.pixels != null ? Number(h.pixels).toLocaleString() : '-')}}</div>
      `;
      hitList.appendChild(li);

      const tr = document.querySelector('#allTable tr[data-prompt="' + h.prompt.replace(/"/g, '') + '"]');
      if (tr) {{
        tr.classList.add('hit-row');
        if (i === 0) tr.classList.add('top');
      }}
    }});
  }}

  try {{
    meta = {imeta_js};
    if (!meta || !meta.classes || !meta.classes.length) {{
      throw new Error('interactive_meta empty — re-run merge_multi_results.py');
    }}
    classes = meta.classes.map(c => ({{
      id: c.id,
      prompt: c.prompt,
      color: c.color,
      rle: c.rle || [],
      scores_min: c.scores_min,
      scores_max: c.scores_max,
      scores_mean: c.scores_mean,
      n_objects: c.n_objects,
      pixels: c.pixels,
    }}));
    setStatus('loading overlay…');
    overlayImg = await loadImage(meta.overlay || '{overlay_name}');
    draw();
    setStatus('ready — ' + classes.length + ' classes · click image (RLE hit-test)', true);
  }} catch (e) {{
    setStatus(String(e.message || e), false);
    try {{
      overlayImg = await loadImage('{overlay_name}');
      meta = meta || {{ width: overlayImg.naturalWidth, height: overlayImg.naturalHeight }};
      canvas.width = overlayImg.naturalWidth;
      canvas.height = overlayImg.naturalHeight;
      ctx.drawImage(overlayImg, 0, 0);
    }} catch (_) {{}}
  }}

  canvas.addEventListener('click', (ev) => {{
    if (!meta || !classes.length) return;
    const pt = eventToNatural(ev);
    queryPixel(pt.x, pt.y, ev.shiftKey);
  }});

  window.addEventListener('keydown', (ev) => {{
    if (ev.key === 'Escape') {{
      lastHits = [];
      marker = null;
      renderHits();
      draw();
    }}
  }});
}})();
</script>
</body>
</html>
"""
    out_path.write_text(html, encoding="utf-8")


def main() -> int:
    args = parse_args()
    result_dir = args.result_dir.resolve()
    if not result_dir.is_dir():
        print(f"ERROR: result-dir not found: {result_dir}", file=sys.stderr)
        return 2

    summary = load_summary(result_dir)
    # Prefer all masks on disk (summary.json is only the *last* multi run)
    disk_items = discover_masks(result_dir)
    if summary:
        meta_by_prompt = {
            r.get("prompt"): r for r in summary.get("results", []) if r.get("prompt")
        }
        for it in disk_items:
            # exact or underscore-normalized match
            meta = meta_by_prompt.get(it["prompt"])
            if meta is None:
                meta = meta_by_prompt.get(it["prompt"].replace(" ", "_"))
            if meta is None:
                for k, v in meta_by_prompt.items():
                    if k.replace(" ", "_") == it["tag"] or k == it["tag"].replace(
                        "_", " "
                    ):
                        meta = v
                        break
            if meta:
                if it.get("n_objects") is None and meta.get("n_objects") is not None:
                    it["n_objects"] = meta.get("n_objects")
                it["scores_min"] = meta.get("scores_min")
                it["scores_max"] = meta.get("scores_max")
        items = disk_items
        # also include any summary-only masks not found by discover
        if not items:
            items = items_from_summary(summary, result_dir)
    else:
        items = disk_items

    # only keep masks that have content
    filtered: list[dict[str, Any]] = []
    for it in items:
        try:
            m = Image.open(it["mask_path"])
            arr = np.array(m.convert("L") if m.mode != "L" else m)
            if (arr > 0).any():
                filtered.append(it)
        except Exception:
            continue
    items = filtered
    enrich_item_scores(items)

    if not items:
        print(f"ERROR: no non-empty mask files in {result_dir}", file=sys.stderr)
        return 1

    try:
        image_path = resolve_base_image(args.image, summary, result_dir)
    except FileNotFoundError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 2

    out_dir = (args.output_dir or (result_dir / "merged")).resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    print("=== merge multi-prompt results ===")
    print(f"result_dir: {result_dir}")
    print(f"image:      {image_path}")
    print(f"classes:    {len(items)}")
    print(f"output:     {out_dir}")

    base_full = Image.open(image_path).convert("RGB")
    base = resize_max(base_full, args.max_side)
    print(f"base size:  {base_full.size} -> overview {base.size}")

    # resize masks via load_mask_bool to base.size
    overlay, legend = build_overlay(base, items, alpha=args.alpha)
    if not legend:
        print("ERROR: all masks empty after resize", file=sys.stderr)
        return 1

    # propagate scores from items into legend
    by_prompt = {it["prompt"]: it for it in items}
    for ent in legend:
        it = by_prompt.get(ent["prompt"])
        if not it:
            continue
        for k in ("scores_min", "scores_max", "scores_mean", "n_objects"):
            if ent.get(k) is None and it.get(k) is not None:
                ent[k] = it[k]

    legend_panel = draw_legend_panel(legend)
    combined = compose_with_legend(overlay, legend_panel)

    overlay_path = out_dir / "combined_overlay.png"
    with_legend_path = out_dir / "combined_with_legend.png"
    overlay.save(overlay_path, optimize=True)
    combined.save(with_legend_path, optimize=True)
    print(f"saved: {overlay_path}")
    print(f"saved: {with_legend_path}")

    # ordered items matching legend (for hitmasks + grid)
    ordered = []
    for ent in legend:
        for it in items:
            if it["prompt"] == ent["prompt"] or it["prompt"].replace(
                "_", " "
            ) == ent["prompt"]:
                ordered.append(it)
                break
    if not ordered:
        ordered = items

    grid_name = None
    if not args.no_grid:
        colors = [PALETTE[i % len(PALETTE)] for i in range(len(ordered))]
        grid = build_grid(
            base_full,
            ordered,
            colors,
            cols=args.grid_cols,
            thumb=args.thumb_size,
            alpha=args.alpha,
        )
        grid_path = out_dir / "per_class_grid.png"
        grid.save(grid_path, optimize=True)
        grid_name = grid_path.name
        print(f"saved: {grid_path}")

    # RLE hit data for interactive click query (file:// safe, no getImageData)
    print("exporting RLE hit data for interactive HTML...")
    classes_meta = export_hitmasks(out_dir, ordered, legend, base.size)
    print(f"saved: {out_dir / 'interactive_meta.json'} ({len(classes_meta)} classes, RLE)")

    merge_summary = {
        "result_dir": str(result_dir),
        "image": str(image_path),
        "n_classes": len(legend),
        "max_side": args.max_side,
        "alpha": args.alpha,
        "overview_size": list(base.size),
        "classes": legend,
        "outputs": {
            "combined_overlay": str(overlay_path),
            "combined_with_legend": str(with_legend_path),
            "per_class_grid": str(out_dir / "per_class_grid.png")
            if grid_name
            else None,
            "interactive_meta": str(out_dir / "interactive_meta.json"),
            "index_html": str(out_dir / "index.html"),
        },
    }
    sum_path = out_dir / "merge_summary.json"
    with open(sum_path, "w", encoding="utf-8") as f:
        json.dump(merge_summary, f, ensure_ascii=False, indent=2)
    print(f"saved: {sum_path}")

    if not args.no_html:
        html_path = out_dir / "index.html"
        interactive_meta = {
            "width": base.size[0],
            "height": base.size[1],
            "overlay": overlay_path.name,
            "classes": classes_meta,
        }
        # Use overlay without side legend so click coords map 1:1 to hitmasks
        write_html(
            html_path,
            overlay_name=overlay_path.name,
            grid_name=grid_name,
            legend=legend,
            meta=merge_summary,
            interactive_meta=interactive_meta,
        )
        print(f"saved: {html_path} (interactive click inspect)")

    print("\nClasses merged:")
    for ent in legend:
        n = ent.get("n_objects")
        print(f"  - {ent['prompt']}: n={n} px={ent['pixels']:,}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())