railway-client/tools/merge_labels.py

"""
검출 결과 후처리 — 같은 병합 그룹에 속한 라벨끼리 외곽선이 gap px 이내로
인접하면 하나로 합치고 그룹 대표 라벨을 붙인다.

예) "building" + "building rooftop" 이 맞닿아 있으면 → "building" 하나로.

사용법:
    python tools/merge_labels.py \
        --json   "data/everyimage/output/0006_multi.json" \
        --groups configs/merge_groups.txt \
        --gap 2
출력:
    <입력>_merged.json  (뷰어에 그대로 넣을 수 있음)
"""
import argparse
import json
import sys
from collections import Counter, OrderedDict
from pathlib import Path

import cv2
import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from sam3_everything_explore import _bbox, _polys_touch  # noqa: E402


def load_merge_groups(path: Path) -> "OrderedDict[str, str]":
    """라벨 → 대표 라벨 매핑. '# [대표]' 헤더 아래 라벨들이 그 그룹."""
    mapping, current = OrderedDict(), None
    for line in path.read_text(encoding="utf-8").splitlines():
        s = line.strip()
        if not s:
            continue
        if s.startswith("#"):
            body = s.lstrip("#").strip()
            current = body[1:-1].strip() if body.startswith("[") and body.endswith("]") else None
            continue
        if current:
            mapping[s] = current
    return mapping


def load_claim_rules(path: Path):
    """'<그룹> : <비율> : <대상 라벨들>' → [(그룹, 비율, {대상라벨})]"""
    rules = []
    for line in path.read_text(encoding="utf-8").splitlines():
        s = line.strip()
        if not s or s.startswith("#"):
            continue
        parts = [p.strip() for p in s.split(":")]
        if len(parts) != 3:
            raise SystemExit(f"claim 규칙 형식 오류: {line}")
        rules.append((parts[0], float(parts[1]),
                      {t.strip() for t in parts[2].split(",") if t.strip()}))
    return rules


def _overlap_ratio(pa, pb):
    """pa 면적 대비 pa∩pb 비율."""
    xs = [p[0] for p in pa] + [p[0] for p in pb]
    ys = [p[1] for p in pa] + [p[1] for p in pb]
    x0, y0 = int(min(xs)) - 1, int(min(ys)) - 1
    x1, y1 = int(max(xs)) + 1, int(max(ys)) + 1
    ca = np.zeros((y1 - y0, x1 - x0), np.uint8)
    cb = np.zeros_like(ca)
    cv2.fillPoly(ca, [np.array(pa, np.int32) - (x0, y0)], 255)
    cv2.fillPoly(cb, [np.array(pb, np.int32) - (x0, y0)], 255)
    area = int(np.count_nonzero(ca))
    return 0.0 if area == 0 else np.count_nonzero(cv2.bitwise_and(ca, cb)) / area


def claim_labels(shapes, mapping, rules):
    """겹침 기준으로 라벨 소유권을 재배정한다. 반환: 바뀐 개수."""
    changed = 0
    boxes = [_bbox(s["points"]) for s in shapes]
    for group, ratio, targets in rules:
        owners = [i for i, s in enumerate(shapes)
                  if mapping.get(s.get("label", "")) == group]
        if not owners:
            print(f"  [claim] 그룹 '{group}' 폴리곤 없음 — 규칙 무시")
            continue
        for i, s in enumerate(shapes):
            if s.get("label", "") not in targets:
                continue
            ax0, ay0, ax1, ay1 = boxes[i]
            best, best_r = None, 0.0
            for j in owners:
                bx0, by0, bx1, by1 = boxes[j]
                if ax1 < bx0 or bx1 < ax0 or ay1 < by0 or by1 < ay0:
                    continue
                r = _overlap_ratio(s["points"], shapes[j]["points"])
                if r > best_r:
                    best, best_r = shapes[j], r
            if best is not None and best_r >= ratio:
                s["claimed_from"] = s["label"]
                s["label"] = best["label"]
                changed += 1
    return changed


def merge_by_group(shapes, mapping, gap=2, epsilon=1.5):
    """대표 라벨이 같은 것끼리 외곽선 인접 시 병합. 그룹 밖 라벨은 그대로 통과."""
    grouped, passthrough = {}, []
    for s in shapes:
        rep = mapping.get(s.get("label", ""))
        if rep is None:
            passthrough.append(s)
        else:
            grouped.setdefault(rep, []).append(s)

    merged = list(passthrough)
    for rep, items in grouped.items():
        parent = list(range(len(items)))

        def find(i):
            while parent[i] != i:
                parent[i] = parent[parent[i]]
                i = parent[i]
            return i

        boxes = [_bbox(s["points"]) for s in items]
        for i in range(len(items)):
            for j in range(i + 1, len(items)):
                if find(i) == find(j):
                    continue
                ax0, ay0, ax1, ay1 = boxes[i]
                bx0, by0, bx1, by1 = boxes[j]
                if ax1 + gap < bx0 or bx1 + gap < ax0 or ay1 + gap < by0 or by1 + gap < ay0:
                    continue                       # 조기 탈락 (판정은 아래 픽셀 단위)
                if _polys_touch(items[i]["points"], items[j]["points"], gap):
                    parent[find(j)] = find(i)

        clusters = {}
        for i in range(len(items)):
            clusters.setdefault(find(i), []).append(i)

        for members in clusters.values():
            src = [items[i] for i in members]
            best = max(src, key=lambda s: float(s.get("score", 0)))
            if len(src) == 1:
                merged.append({**src[0], "label": rep,
                               "merged_from": [src[0].get("label", "")]})
                continue
            pts = [p for s in src for p in s["points"]]
            x0 = int(min(p[0] for p in pts)) - gap - 1
            y0 = int(min(p[1] for p in pts)) - gap - 1
            x1 = int(max(p[0] for p in pts)) + gap + 1
            y1 = int(max(p[1] for p in pts)) + gap + 1
            canvas = np.zeros((y1 - y0, x1 - x0), np.uint8)
            for s in src:
                cv2.fillPoly(canvas, [np.array(s["points"], np.int32) - (x0, y0)], 255)
            if gap > 0:                            # 틈 메우기
                k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (gap * 2 + 1, gap * 2 + 1))
                canvas = cv2.morphologyEx(canvas, cv2.MORPH_CLOSE, k)
            contours, _ = cv2.findContours(canvas, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            names = sorted({s.get("label", "") for s in src})
            for cnt in contours:
                approx = cv2.approxPolyDP(cnt, epsilon, True)
                if len(approx) < 3:
                    continue
                merged.append({
                    "label": rep,
                    "score": float(best.get("score", 0)),
                    "shape_type": "polygon",
                    "merged_from": names,
                    "points": [[float(p[0][0] + x0), float(p[0][1] + y0)] for p in approx],
                })
    return merged


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--json",   required=True)
    ap.add_argument("--groups", default="configs/merge_groups.txt")
    ap.add_argument("--claims", default="configs/claim_rules.txt",
                    help="소유권 재배정 규칙 (없으면 빈 문자열로 끄기)")
    ap.add_argument("--gap",    type=int, default=2, help="외곽선 인접 판정 px (기본 2)")
    ap.add_argument("--output", default=None, help="기본: <입력>_merged.json")
    args = ap.parse_args()

    jpath = Path(args.json)
    data = json.loads(jpath.read_text(encoding="utf-8"))
    shapes = data.get("segments", [])
    mapping = load_merge_groups(Path(args.groups))

    reps = sorted(set(mapping.values()))
    print(f"입력 {len(shapes)}개 · 병합 그룹 {len(reps)}개: {', '.join(reps)}")

    if args.claims:
        rules = load_claim_rules(Path(args.claims))
        n = claim_labels(shapes, mapping, rules)
        print(f"소유권 재배정: {n}개 라벨 변경")

    out = merge_by_group(shapes, mapping, gap=args.gap)
    counts = Counter(s.get("label", "") for s in out)
    print(f"병합(gap={args.gap}px) {len(shapes)} → {len(out)}개\n")
    for lb, c in counts.most_common():
        print(f"  {lb:34s} {c:4d}")

    out_path = Path(args.output) if args.output else jpath.with_name(jpath.stem + "_merged.json")
    out_path.write_text(json.dumps({
        "total_segments": len(out),
        "label_counts": dict(counts),
        "segments": [{"label": s.get("label", ""), "score": s.get("score", 0),
                      "merged_from": s.get("merged_from", []),
                      "bbox": list(_bbox(s["points"])), "points": s["points"]}
                     for s in out],
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\n저장: {out_path}")


if __name__ == "__main__":
    main()