"""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: /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"" f"" f"{ent['prompt']}" f"{n if n is not None else '-'}" f"{score_cell}" f"{ent['pixels']:,}" ) grid_block = ( f'
Per-class grid' f'grid
' if grid_name else "" ) n_cls = len(legend) html = f""" Merged multi-prompt — interactive

Merged multi-prompt — click to inspect

이미지를 클릭하면 해당 픽셀에 겹친 클래스(프롬프트)와 confidence 점수가 오른쪽에 표시됩니다. 겹침 시 점수 높은 순으로 정렬됩니다.

클릭: 픽셀 조회 · Shift+클릭: 선택 유지(누적) · Esc: 초기화 · 클래스 {n_cls}개

loading…

{grid_block}
""" 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())