railway-client/tools/detect_raamen.py

"""
드론 사선 촬영 이미지에서 라멘형(門자형) 전철주 검출.
기존 픽셀 위상수학(Skeleton) 방식을 공간 기하학 방식으로 교체.

파이프라인:
  Phase 1: 폴리곤 단순화(approxPolyDP) + 소실점(Vanishing Point) 계산
  Phase 2: 동적 C/B 분류 (소실점 기반 기대 각도) — C=기둥(Column), B=빔(Beam)
  Phase 3: 근접성 기반 그룹핑 (B 앵커 → 아래 C 탐색)
  Phase 4: 라멘 구조 판정 + 예외(가림) 처리

사용:
    python tools/detect_raamen.py \
        --image <path> --label <path> --output <path> \
        [--class-ids 1] [--epsilon 4.0] [--c-thresh 20.0]
"""
import argparse
import numpy as np
import cv2
from pathlib import Path


# ── Phase 1: 파싱 + 단순화 + 소실점 ─────────────────────────────────────

def load_polygons(label_path, W, H, class_ids=None, class_names=None):
    """
    AnyLabeling JSON 또는 YOLO .txt 자동 감지 후 파싱.
    Returns: (픽셀 좌표 폴리곤 리스트, 절대 shapes[] 인덱스 리스트)
    절대 인덱스는 AnyLabeling JSON shapes[] 배열 인덱스와 동일.
    """
    import json as _json
    if not label_path.exists():
        raise FileNotFoundError(f"라벨 파일을 찾을 수 없습니다: {label_path}")

    if label_path.suffix.lower() == ".json":
        data = _json.loads(label_path.read_text(encoding="utf-8"))
        shapes = data.get("shapes", [])
        polys, abs_indices = [], []
        for abs_idx, s in enumerate(shapes):
            if class_names is not None and s.get("label", "") not in class_names:
                continue
            pts = np.array([[float(p[0]), float(p[1])] for p in s.get("points", [])],
                           dtype=np.float32)
            if len(pts) >= 3:
                polys.append(pts)
                abs_indices.append(abs_idx)
        return polys, abs_indices

    # YOLO .txt
    polys, abs_indices = [], []
    for abs_idx, line in enumerate(label_path.read_text(encoding="utf-8").splitlines()):
        parts = line.split()
        if not parts:
            continue
        if class_ids is not None and int(parts[0]) not in class_ids:
            continue
        coords = list(map(float, parts[1:]))
        pts = np.array(
            [[coords[i] * W, coords[i + 1] * H] for i in range(0, len(coords), 2)],
            dtype=np.float32,
        )
        if len(pts) >= 3:
            polys.append(pts)
            abs_indices.append(abs_idx)
    return polys, abs_indices


def smooth_polygon(pts, epsilon):
    """approxPolyDP로 노이즈 경계 → 직선 위주 단순화."""
    approx = cv2.approxPolyDP(pts.astype(np.int32), epsilon, closed=True)
    result = approx.reshape(-1, 2).astype(np.float32)
    return result if len(result) >= 3 else pts.astype(np.float32)


def _minrect(pts):
    """cv2.minAreaRect 래퍼 (int32 변환 포함)."""
    return cv2.minAreaRect(pts.astype(np.int32))


def _long_axis_angle(pts):
    """minAreaRect 장축 각도 (degrees) 반환."""
    (cx, cy), (rw, rh), angle = _minrect(pts)
    return angle if rw >= rh else angle + 90, cx, cy, max(rw, rh), min(rw, rh)


def _radial_cos_sim(pts, img_cx, img_cy):
    """VP 후보 bootstrap용 기존 radial cos_sim."""
    long_angle_deg, cx, cy, long_side, short_side = _long_axis_angle(pts)
    if short_side < 1 or long_side / short_side < 1.3:
        return 0.0
    lx = np.cos(np.radians(long_angle_deg))
    ly = np.sin(np.radians(long_angle_deg))
    rdx, rdy = cx - img_cx, cy - img_cy
    n = (rdx ** 2 + rdy ** 2) ** 0.5
    return abs(lx * rdx / n + ly * rdy / n) if n > 1 else 0.0


def compute_vanishing_point(polys):
    """
    후보 폴리곤 장축 선분들의 최소자승 교점 (소실점) 계산.
    각 장축 선분: 법선 n=(-dy, dx), 방정식 -dy*x + dx*y = -dy*cx + dx*cy
    """
    A_rows, b_vals = [], []
    for pts in polys:
        long_angle_deg, cx, cy, _, _ = _long_axis_angle(pts)
        dx = np.cos(np.radians(long_angle_deg))
        dy = np.sin(np.radians(long_angle_deg))
        A_rows.append([-dy, dx])
        b_vals.append(-dy * cx + dx * cy)
    vp, *_ = np.linalg.lstsq(np.array(A_rows), np.array(b_vals), rcond=None)
    return float(vp[0]), float(vp[1])


def _estimate_vp_iterative(polys, seed_indices, c_thresh, b_max_diff, vp_min_len, max_iter=6):
    """초기 후보에서 반복 정제 VP 추정. Returns (vp_x, vp_y, n_c, orients, adiffs)."""
    n = len(polys)
    orients = ['?'] * n
    adiffs  = [90.0] * n
    if len(seed_indices) < 2:
        return None, None, 0, orients, adiffs
    vp_x, vp_y = compute_vanishing_point([polys[i] for i in seed_indices])
    for _ in range(max_iter):
        for i, pts in enumerate(polys):
            orients[i], adiffs[i] = classify_cb(pts, vp_x, vp_y, c_thresh, b_max_diff)
        c_cands = [i for i in range(n)
                   if orients[i] == 'C' and _long_axis_angle(polys[i])[3] > vp_min_len]
        if len(c_cands) < 3:
            break
        nx, ny = compute_vanishing_point([polys[i] for i in c_cands])
        shift = ((nx - vp_x) ** 2 + (ny - vp_y) ** 2) ** 0.5
        vp_x, vp_y = nx, ny
        if shift < 5.0:
            for i, pts in enumerate(polys):
                orients[i], adiffs[i] = classify_cb(pts, vp_x, vp_y, c_thresh, b_max_diff)
            break
    return vp_x, vp_y, orients.count('C'), orients, adiffs


# ── Phase 2: 동적 C/B 분류 ──────────────────────────────────────────────

def classify_cb(pts, vp_x, vp_y, c_thresh, b_max_diff=75.0):
    """
    소실점 기준 C/B 분류. C=기둥(Column), B=빔(Beam).
    - diff < c_thresh               → C (기둥)
    - c_thresh ≤ diff < b_max_diff  → B (빔)
    - diff ≥ b_max_diff             → '?' (레일/전선 등 비라멘 수평 구조물)
    Returns: (orient, angle_diff_deg)
      orient = 'C' | 'B' | '?'
    """
    long_angle_deg, cx, cy, long_side, short_side = _long_axis_angle(pts)
    if short_side < 1 or long_side / short_side < 1.3:
        return '?', 90.0
    # 기대 수직 각도: 폴리곤 중심 → 소실점 방향
    exp_angle = np.degrees(np.arctan2(vp_y - cy, vp_x - cx))
    # 장축은 방향 무관 → 0~90° 범위로 정규화
    diff = abs(long_angle_deg - exp_angle) % 180.0
    if diff > 90.0:
        diff = 180.0 - diff
    if diff < c_thresh:
        return 'C', diff
    if diff < b_max_diff:
        return 'B', diff
    return '?', diff


# ── Phase 3: 폴리곤 접촉/교차 기반 그룹핑 ───────────────────────────────

def _poly_union(polys_list):
    """여러 폴리곤의 rasterization union 외곽 contour 반환 (실제 union shape)."""
    all_pts = np.vstack(polys_list)
    x0 = int(all_pts[:, 0].min()) - 1
    y0 = int(all_pts[:, 1].min()) - 1
    x1 = int(all_pts[:, 0].max()) + 2
    y1 = int(all_pts[:, 1].max()) + 2
    w, h = x1 - x0, y1 - y0
    mask = np.zeros((h, w), dtype=np.uint8)
    off = np.array([[x0, y0]], dtype=np.float32)
    for pts in polys_list:
        cv2.fillPoly(mask, [(pts - off).astype(np.int32)], 255)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return all_pts.astype(np.float32)
    cnt = max(contours, key=cv2.contourArea).reshape(-1, 2).astype(np.float32)
    return cnt + [x0, y0]


def _polys_intersect(pts_i, pts_j, bbox_i, bbox_j):
    """두 폴리곤의 실제 픽셀 교차 여부 확인 (rasterization 기반)."""
    ax0, ay0, ax1, ay1 = bbox_i
    bx0, by0, bx1, by1 = bbox_j
    ix0 = max(int(ax0), int(bx0))
    iy0 = max(int(ay0), int(by0))
    ix1 = min(int(ax1), int(bx1))
    iy1 = min(int(ay1), int(by1))
    if ix0 >= ix1 or iy0 >= iy1:
        return False
    w, h = ix1 - ix0 + 1, iy1 - iy0 + 1
    off = np.array([[ix0, iy0]], dtype=np.float32)
    mi = np.zeros((h, w), dtype=np.uint8)
    mj = np.zeros((h, w), dtype=np.uint8)
    cv2.fillPoly(mi, [(pts_i - off).astype(np.int32)], 1)
    cv2.fillPoly(mj, [(pts_j - off).astype(np.int32)], 1)
    return bool(np.any(mi & mj))


def remove_beam_center_small(polys, orients, poly_abs_idx,
                             center_ratio=0.30, area_ratio=0.15, min_beam_len=100):
    """
    B(빔) 폴리곤 장축의 중앙 center_ratio 구간에 위치한 소형 폴리곤 제거.
    소형 기준: 해당 B 면적의 area_ratio 미만.
    """
    n = len(polys)
    remove_set = set()

    for i in range(n):
        if orients[i] != 'B':
            continue
        la, bcx, bcy, long_side, short_side = _long_axis_angle(polys[i])
        if long_side < min_beam_len:
            continue
        b_area = cv2.contourArea(polys[i].astype(np.int32))
        if b_area <= 0:
            continue

        angle_rad = np.radians(la)
        ux, uy = np.cos(angle_rad), np.sin(angle_rad)

        pts_b = polys[i]
        proj_b = pts_b[:, 0] * ux + pts_b[:, 1] * uy
        proj_center = (proj_b.min() + proj_b.max()) / 2.0
        half = (proj_b.max() - proj_b.min()) * center_ratio / 2.0

        for j in range(n):
            if j == i or j in remove_set:
                continue
            if orients[j] not in ('B', '?'):  # C(기둥)은 절대 제거 안 함
                continue
            j_area = cv2.contourArea(polys[j].astype(np.int32))
            if j_area >= b_area * area_ratio:
                continue
            jcx = float(polys[j][:, 0].mean())
            jcy = float(polys[j][:, 1].mean())
            proj_j = jcx * ux + jcy * uy
            if proj_center - half <= proj_j <= proj_center + half:
                remove_set.add(j)

    keep = [i for i in range(n) if i not in remove_set]
    if remove_set:
        print(f"  [B 중앙부 소형 제거] {len(remove_set)}개 idx={sorted(remove_set)}")
        polys        = [polys[i]        for i in keep]
        orients      = [orients[i]      for i in keep]
        poly_abs_idx = [poly_abs_idx[i] for i in keep]

    return polys, orients, poly_abs_idx, keep


def connectivity_groups(polys, orients, margin=30):
    """
    폴리곤 bbox가 margin px 이내로 닿거나 교차하면 같은 그룹 (B/C 구분 없음).
    Union-Find로 연결된 폴리곤들을 묶은 뒤, 각 그룹 내에서 B/C 목록 분리.
    Returns: list of {'id': int, 'B': [idx,...], 'C': [idx,...]}
    """
    n = len(polys)
    parent = list(range(n))

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

    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    # 폴리곤 bbox 미리 계산
    bboxes = []
    for pts in polys:
        bboxes.append((pts[:, 0].min(), pts[:, 1].min(),
                       pts[:, 0].max(), pts[:, 1].max()))

    # bbox가 margin 이내로 닿거나 교차하면 union
    for i in range(n):
        ax0, ay0, ax1, ay1 = bboxes[i]
        for j in range(i + 1, n):
            bx0, by0, bx1, by1 = bboxes[j]
            if (ax1 + margin >= bx0 and bx1 + margin >= ax0 and
                    ay1 + margin >= by0 and by1 + margin >= ay0):
                union(i, j)

    # 연결 컴포넌트 수집
    from collections import defaultdict
    comp = defaultdict(list)
    for i in range(n):
        comp[find(i)].append(i)

    groups = []
    for gid, members in enumerate(comp.values(), 1):
        b_list = sorted(i for i in members if orients[i] == 'B')
        c_list = sorted(i for i in members if orients[i] == 'C')
        groups.append({'id': gid, 'B': b_list, 'C': c_list})

    # 면적 내림차순으로 ID 재부여 (큰 그룹이 G1)
    for g in groups:
        g['area'] = sum(cv2.contourArea(polys[i].astype(np.int32))
                        for i in g['B'] + g['C'])
    groups.sort(key=lambda x: x['area'], reverse=True)
    for gid, g in enumerate(groups, 1):
        g['id'] = gid

    return groups


def reclassify_center_groups(groups, polys, W, center_ratio=0.2):
    """
    RAAMEN_CENTER 후보 그룹 (B 없음, C 2개 이상, 이미지 중앙):
    가장 아래(+Y 최대) C 폴리곤만 C(기둥)로 유지, 나머지 C → B(빔)로 재분류.
    """
    for g in groups:
        if g['B'] or len(g['C']) < 2:
            continue
        all_c_pts = np.vstack([polys[i] for i in g['C']])
        gcx = float(all_c_pts[:, 0].mean())
        if abs(gcx - W / 2) >= W * center_ratio:
            continue
        bottom_c = max(g['C'], key=lambda i: float(polys[i][:, 1].mean()))
        g['B'] = sorted(i for i in g['C'] if i != bottom_c)
        g['C'] = [bottom_c]
    return groups


def merge_intersecting_same_type(groups, polys, poly_abs_idx):
    """
    각 그룹 내에서 교차하는 동일 타입(B↔B, C↔C) 폴리곤들을 convex hull로 병합.
    교차하는 폴리곤 클러스터 → 하나의 합성 폴리곤으로 대체.
    Returns: (groups, extended_polys, extended_abs_idx)
      - extended_polys: 원본 + 새 병합 폴리곤 (append 방식, 원본 인덱스 유지)
      - 병합된 폴리곤의 abs_idx는 원본 인덱스 리스트를 저장 (추적용)
    """
    from collections import defaultdict

    ext_polys = list(polys)
    ext_abs_idx = list(poly_abs_idx)

    for g in groups:
        for key in ('B', 'C'):
            members = g[key]
            if len(members) <= 1:
                continue

            n = len(members)
            parent = list(range(n))

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

            def union(a, b):
                ra, rb = find(a), find(b)
                if ra != rb:
                    parent[ra] = rb

            bboxes_m = [(ext_polys[members[k]][:, 0].min(), ext_polys[members[k]][:, 1].min(),
                         ext_polys[members[k]][:, 0].max(), ext_polys[members[k]][:, 1].max())
                        for k in range(n)]

            for a in range(n):
                for b in range(a + 1, n):
                    if _polys_intersect(ext_polys[members[a]], ext_polys[members[b]],
                                        bboxes_m[a], bboxes_m[b]):
                        union(a, b)

            clusters = defaultdict(list)
            for k in range(n):
                clusters[find(k)].append(members[k])

            new_members = []
            for cluster in clusters.values():
                if len(cluster) == 1:
                    new_members.append(cluster[0])
                else:
                    cluster_polys = [ext_polys[i] for i in cluster]
                    union_pts = _poly_union(cluster_polys)
                    new_idx = len(ext_polys)
                    ext_polys.append(union_pts)
                    flat = []
                    for i in cluster:
                        v = ext_abs_idx[i]
                        flat.extend(v) if isinstance(v, list) else flat.append(v)
                    ext_abs_idx.append(sorted(flat))
                    new_members.append(new_idx)

            g[key] = sorted(new_members)

    return groups, ext_polys, ext_abs_idx


# ── Phase 4: 라멘 구조 판정 ─────────────────────────────────────────────

def _cluster_polys(indices, polys, margin=60):
    """
    인접 폴리곤을 클러스터링 → 물리적 객체(기둥) 단위 반환.
    Returns: list of [idx, ...] clusters
    """
    if not indices:
        return []
    n = len(indices)
    parent = list(range(n))

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

    bboxes = [(polys[i][:, 0].min(), polys[i][:, 1].min(),
               polys[i][:, 0].max(), polys[i][:, 1].max()) for i in indices]
    for a in range(n):
        ax0, ay0, ax1, ay1 = bboxes[a]
        for b in range(a + 1, n):
            bx0, by0, bx1, by1 = bboxes[b]
            if (ax1 + margin >= bx0 and bx1 + margin >= ax0 and
                    ay1 + margin >= by0 and by1 + margin >= ay0):
                ra, rb = find(a), find(b)
                if ra != rb:
                    parent[ra] = rb

    from collections import defaultdict
    clusters = defaultdict(list)
    for i in range(n):
        clusters[find(i)].append(indices[i])
    return list(clusters.values())


def judge_raamen(group, polys, W, center_ratio=0.2, v_cluster_margin=60):
    """
    라멘 구조 판정.
    - 빔 기준: 그룹 내 가장 큰 B 폴리곤
    - 기둥 수: C 폴리곤을 근접 클러스터링한 클러스터 수
    - 빔 x 범위(±50%) 밖의 C 클러스터는 잡폴리곤으로 무시
    Returns: ('RAAMEN' | 'RAAMEN_OCCLUDED' | 'PARTIAL' | '', n_poles)
    """
    bs, cs = group['B'], group['C']

    # B 없음: 중앙 영역이면 C/B 분류 자체가 신뢰 불가 (기둥·빔 각도 수렴)
    # → 2개 이상의 C 폴리곤이 중앙에 모여 있으면 RAAMEN_CENTER 처리
    if not bs:
        if len(cs) >= 2:
            all_c_pts = np.vstack([polys[i] for i in cs])
            gcx = float(all_c_pts[:, 0].mean())
            if abs(gcx - W / 2) < W * center_ratio:
                return 'RAAMEN_CENTER', len(cs)
        return '', 0

    # 큰 B 폴리곤 최대 2개를 빔 기준으로 사용 (2nd가 1st 면적의 50% 이상이면 포함)
    b_by_area = sorted(bs, key=lambda i: cv2.contourArea(polys[i].astype(np.int32)), reverse=True)
    main_bs = [b_by_area[0]]
    if len(b_by_area) > 1:
        a0 = cv2.contourArea(polys[b_by_area[0]].astype(np.int32))
        a1 = cv2.contourArea(polys[b_by_area[1]].astype(np.int32))
        if a1 >= a0 * 0.5:
            main_bs.append(b_by_area[1])

    # 이미지 중앙 여부 판정 (X축 기반 유지)
    hx0 = int(min(polys[i][:, 0].min() for i in main_bs))
    hx1 = int(max(polys[i][:, 0].max() for i in main_bs))
    hcx = (hx0 + hx1) / 2.0
    is_center = abs(hcx - W / 2) < W * center_ratio

    # 빔 장축(major axis) 방향으로 투영 — 대각 빔도 정확하게 span 계산
    la, _, _, _, _ = _long_axis_angle(polys[main_bs[0]])
    angle_rad = np.radians(la)
    ux, uy = np.cos(angle_rad), np.sin(angle_rad)
    beam_pts = np.vstack([polys[i] for i in main_bs])
    proj_beam = beam_pts[:, 0] * ux + beam_pts[:, 1] * uy
    proj_min, proj_max = float(proj_beam.min()), float(proj_beam.max())
    span = proj_max - proj_min

    # C 폴리곤 클러스터링 → 기둥 단위
    pole_clusters = _cluster_polys(cs, polys, margin=v_cluster_margin)

    # 기둥은 빔 양 끝단(좌 35% / 우 35%)에만 존재 (장축 투영 기준)
    p_tol      = max(span * 0.5, 50)
    left_zone  = proj_min + span * 0.35
    right_zone = proj_max - span * 0.35
    valid_projs = []
    for cluster in pole_clusters:
        # 멤버별 투영값 계산
        member_projs = [polys[i][:, 0].mean() * ux + polys[i][:, 1].mean() * uy
                        for i in cluster]
        proj_c = float(np.mean(member_projs))   # 중심값 (range 체크)
        proj_lo = min(member_projs)              # 가장 왼쪽 끝
        proj_hi = max(member_projs)              # 가장 오른쪽 끝
        in_range    = proj_min - p_tol <= proj_c <= proj_max + p_tol
        # 클러스터 안 멤버 중 하나라도 끝단에 있으면 유효 (43C 같은 중간 노이즈가 흡수돼도 기둥 인식)
        in_end_zone = proj_lo <= left_zone or proj_hi >= right_zone
        if in_range and in_end_zone:
            # 끝단 대표값: 왼쪽 끝단이면 proj_lo, 오른쪽 끝단이면 proj_hi
            rep = proj_lo if proj_lo <= left_zone else proj_hi
            valid_projs.append(rep)

    n_poles = len(valid_projs)

    if n_poles >= 2:
        lp, rp = min(valid_projs), max(valid_projs)
        if lp <= proj_min + span * 0.4 and rp >= proj_max - span * 0.4:
            return 'RAAMEN', n_poles
        return 'PARTIAL', n_poles

    if n_poles == 1:
        return ('RAAMEN_OCCLUDED', 1) if not is_center else ('', 1)

    return '', 0


def _merge_poly_hull(indices, polys):
    """여러 폴리곤 점들을 합쳐 Convex Hull 좌표 반환 (AnyLabeling points 형식)."""
    all_pts = np.vstack([polys[i] for i in indices])
    hull = cv2.convexHull(all_pts.astype(np.int32))
    return [[float(p[0][0]), float(p[0][1])] for p in hull]


def group_detail(group, polys, W, center_ratio=0.2, v_cluster_margin=60):
    """
    라멘 그룹의 세부 구성 분석.
    Returns dict: main_b, junk_b, valid_pole_clusters, attach_clusters
    """
    bs, cs = group['B'], group['C']
    if not bs:
        return {'main_b': None, 'junk_b': [], 'valid_pole_clusters': [], 'attach_clusters': []}

    b_by_area = sorted(bs, key=lambda i: cv2.contourArea(polys[i].astype(np.int32)), reverse=True)
    main_bs = [b_by_area[0]]
    if len(b_by_area) > 1:
        a0 = cv2.contourArea(polys[b_by_area[0]].astype(np.int32))
        a1 = cv2.contourArea(polys[b_by_area[1]].astype(np.int32))
        if a1 >= a0 * 0.5:
            main_bs.append(b_by_area[1])
    main_b = main_bs[0]  # JSON 출력용 대표 빔
    junk_b = [i for i in bs if i not in main_bs]

    # 빔 장축(major axis) 투영 기반 span 계산
    la, _, _, _, _ = _long_axis_angle(polys[main_bs[0]])
    angle_rad = np.radians(la)
    ux, uy = np.cos(angle_rad), np.sin(angle_rad)
    beam_pts = np.vstack([polys[i] for i in main_bs])
    proj_beam = beam_pts[:, 0] * ux + beam_pts[:, 1] * uy
    proj_min, proj_max = float(proj_beam.min()), float(proj_beam.max())
    span = proj_max - proj_min

    pole_clusters = _cluster_polys(cs, polys, margin=v_cluster_margin)
    p_tol      = max(span * 0.5, 50)
    left_zone  = proj_min + span * 0.35
    right_zone = proj_max - span * 0.35

    valid_pole_clusters, attach_clusters = [], []
    for cluster in pole_clusters:
        member_projs = [polys[i][:, 0].mean() * ux + polys[i][:, 1].mean() * uy
                        for i in cluster]
        proj_c  = float(np.mean(member_projs))
        proj_lo = min(member_projs)
        proj_hi = max(member_projs)
        in_range    = proj_min - p_tol <= proj_c <= proj_max + p_tol
        in_end_zone = proj_lo <= left_zone or proj_hi >= right_zone
        if in_range and in_end_zone:
            valid_pole_clusters.append(cluster)
        elif in_range:
            attach_clusters.append(cluster)

    return {
        'main_b': main_b,
        'main_bs': main_bs,
        'junk_b': junk_b,
        'valid_pole_clusters': valid_pole_clusters,
        'attach_clusters': attach_clusters,
    }


# ── 시각화 상수 ─────────────────────────────────────────────────────────

_CB_COLOR = {
    'C': (255,  80,   0),   # 주황 (기둥 Column)
    'B': (  0,  80, 255),   # 파란 (빔 Beam)
    '?': (140, 140, 140),   # 회색 (미분류)
}
_RAAMEN_COLOR = {
    'RAAMEN':          (  0, 255,   0),   # 초록
    'RAAMEN_CENTER':   (  0, 255, 255),   # 노랑 (중앙 영역, C/B 분류 불신뢰)
    'RAAMEN_OCCLUDED': (  0, 165, 255),   # 주황 (가림/부분 검출)
    'PARTIAL':         (128, 128, 128),   # 회색
}


# ── 메인 렌더링 ─────────────────────────────────────────────────────────

def render(image_path, label_path, output_path, args,
           class_ids=None, class_names=None, epsilon=4.0, c_thresh=20.0,
           b_max_diff=75.0, vp_min_ar=2.5, vp_min_len=80.0, vp_outer_ratio=0.2):

    buf = np.fromfile(str(image_path), dtype=np.uint8)
    img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    H, W = img.shape[:2]

    # ── Phase 1 ──────────────────────────────────────────────────────────
    raw_polys, poly_abs_idx = load_polygons(label_path, W, H, class_ids, class_names)
    polys_all = [smooth_polygon(p, epsilon) for p in raw_polys]

    # 가로:세로 > 4:1 → 전철주 아님 (레일·전선 등), 제거
    keep = []
    skipped = []
    for i, pts in enumerate(polys_all):
        bw = pts[:, 0].max() - pts[:, 0].min()
        bh = pts[:, 1].max() - pts[:, 1].min()
        if bh > 0 and bw / bh > 4.0:
            skipped.append(poly_abs_idx[i])
        else:
            keep.append(i)
    polys = [polys_all[i] for i in keep]
    poly_abs_idx = [poly_abs_idx[i] for i in keep]
    print(f"  {len(polys_all)}개 파싱 → {len(skipped)}개 가로형 필터 제거 → {len(polys)}개 처리")
    if skipped:
        print(f"    제거 shape idx: {skipped}")

    img_cx, img_cy = W / 2.0, H / 2.0

    # ── Phase 1+2: 두 가지 VP 시드 방식으로 시도 → C 폴리곤이 더 많은 VP 채택 ──
    elong_idx = []
    for i, pts in enumerate(polys):
        la, cx, cy, long_side, short_side = _long_axis_angle(pts)
        if short_side > 0 and long_side / short_side > vp_min_ar and long_side > vp_min_len:
            elong_idx.append(i)

    # 시드 A: 지배적 장축 각도 클러스터 (이미지 방향 비의존)
    seeds_A = []
    if len(elong_idx) >= 2:
        avals_all = [(i, _long_axis_angle(polys[i])[0] % 180.0) for i in elong_idx]
        hist, edges = np.histogram([a for _, a in avals_all], bins=12, range=(0.0, 180.0))
        pk = int(np.argmax(hist))
        peak_c = (edges[pk] + edges[pk + 1]) / 2.0
        seeds_A = [i for i, a in avals_all
                   if min(abs(a - peak_c), 180.0 - abs(a - peak_c)) < 20.0]

    # 시드 B: radial cos_sim (이미지 중심 방사 방향 정렬)
    seeds_B = []
    for i in elong_idx:
        _, cx, cy, _, _ = _long_axis_angle(polys[i])
        dist = ((cx - img_cx) ** 2 + (cy - img_cy) ** 2) ** 0.5
        if dist > min(W, H) * vp_outer_ratio and _radial_cos_sim(polys[i], img_cx, img_cy) > 0.5:
            seeds_B.append(i)

    # 두 시드 모두 시도 → C가 더 많이 나오는 VP 채택
    best_vp_x, best_vp_y = img_cx, -H * 3.0
    best_n_c = 0
    orients, adiffs = ['?'] * len(polys), [90.0] * len(polys)
    for label, seeds in [('지배각', seeds_A), ('radial', seeds_B)]:
        if len(seeds) < 2:
            continue
        vx, vy, nc, ors, ads = _estimate_vp_iterative(
            polys, seeds, c_thresh, b_max_diff, vp_min_len)
        print(f"  VP [{label}]: ({vx:.0f}, {vy:.0f})  C={nc}")
        if nc > best_n_c:
            best_vp_x, best_vp_y, best_n_c = vx, vy, nc
            orients, adiffs = ors, ads
    vp_x, vp_y = best_vp_x, best_vp_y
    print(f"  → 채택 VP: ({vp_x:.1f}, {vp_y:.1f})")

    print(f"\n  [C/B 분류] threshold=±{c_thresh}°  b_max=±{b_max_diff}°")
    for i in range(len(polys)):
        print(f"    poly {i:>2d}: {orients[i]}  diff={adiffs[i]:.1f}°")
    print(f"  C:{orients.count('C')}  B:{orients.count('B')}  ?:{orients.count('?')}")

    # B 장축 중앙부 소형 폴리곤 제거 (Phase 3 전)
    polys, orients, poly_abs_idx, beam_keep = remove_beam_center_small(polys, orients, poly_abs_idx)
    old_to_new = {old: new for new, old in enumerate(beam_keep)}
    seeds_A = [old_to_new[i] for i in seeds_A if i in old_to_new]
    seeds_B = [old_to_new[i] for i in seeds_B if i in old_to_new]

    # ── Phase 3 ──────────────────────────────────────────────────────────
    groups = connectivity_groups(polys, orients, margin=args.margin)
    print(f"\n  [그룹핑] 연결 컴포넌트 {len(groups)}개 (margin={args.margin}px)")
    for g in groups:
        print(f"    G{g['id']}: B={g['B']}  C={g['C']}")

    # ── Phase 4 ──────────────────────────────────────────────────────────
    print(f"\n  [라멘 판정]")
    for g in groups:
        verdict, n_poles = judge_raamen(g, polys, W)
        g['verdict'] = verdict
        g['n_poles'] = n_poles
        pole_str = f"{n_poles}poles" if n_poles else "-"
        print(f"    G{g['id']}: B={g['B']} C={g['C']} → {verdict or '-':18s} ({pole_str})")

    # RAAMEN_CENTER: 최하단(+Y 최대) C만 기둥으로 유지, 나머지 C → B (판정 유지)
    for g in groups:
        if g['verdict'] != 'RAAMEN_CENTER' or len(g['C']) < 2:
            continue
        bottom_c = max(g['C'], key=lambda i: float(polys[i][:, 1].mean()))
        for i in g['C']:
            if i != bottom_c:
                orients[i] = 'B'
        g['B'] = sorted(g['B'] + [i for i in g['C'] if i != bottom_c])
        g['C'] = [bottom_c]
        print(f"    G{g['id']} [CENTER 재분류]: B={g['B']} C={g['C']}")

    # 그룹 내 동일 타입 교차 폴리곤 병합 (B↔B, C↔C, 폴리곤 교차 기준)
    in_group_before = set()
    for g in groups:
        in_group_before.update(g['B'] + g['C'])

    orig_poly_n = len(polys)
    groups, polys, poly_abs_idx = merge_intersecting_same_type(groups, polys, poly_abs_idx)

    # 새 병합 폴리곤의 orient 확장
    new_orient_map = {}
    for g in groups:
        for key in ('B', 'C'):
            for idx in g[key]:
                if idx >= orig_poly_n:
                    new_orient_map[idx] = key
    for i in range(orig_poly_n, len(polys)):
        orients.append(new_orient_map.get(i, '?'))

    # 병합으로 흡수된 원본 폴리곤 인덱스 (시각화에서 제외)
    in_group_after = set()
    for g in groups:
        in_group_after.update(g['B'] + g['C'])
    merged_away = {i for i in in_group_before if i not in in_group_after}

    valid_items = [g for g in groups if g['verdict']]
    valid_items.sort(key=lambda x: x['area'], reverse=True)

    print(f"\n  [최종 라멘 객체] {len(valid_items)}개 (면적순)")
    for g in valid_items:
        print(f"    G{g['id']}: B={g['B']} C={g['C']} → {g['verdict']:18s} ({g['n_poles']}poles)  Area={g['area']:,.0f}")

    # 최소 면적 필터링
    if args.min_group_area > 0:
        before = len(valid_items)
        valid_items = [g for g in valid_items if g['area'] >= args.min_group_area]
        print(f"  [필터] 최소 면적 {args.min_group_area} 미만 제거: {before} → {len(valid_items)}개")

    # ── 시각화 ───────────────────────────────────────────────────────────

    # 1. 폴리곤 C/B 색상 반투명 오버레이 (fill + outline) — 병합 흡수된 원본은 제외
    for i, (pts, orient) in enumerate(zip(polys, orients)):
        if i in merged_away:
            continue
        color = _CB_COLOR[orient]
        pts_i = pts.astype(np.int32)
        ov = img.copy()
        cv2.fillPoly(ov, [pts_i], color)
        cv2.addWeighted(ov, 0.25, img, 0.75, 0, img)
        cv2.polylines(img, [pts_i], True, color, 2)

    # 2. VP 시드 폴리곤에 청록 테두리 (채택된 시드셋 재구성)
    for i in seeds_A + [i for i in seeds_B if i not in seeds_A]:
        if i not in merged_away:
            cv2.polylines(img, [polys[i].astype(np.int32)], True, (0, 220, 220), 3)

    # 3. 소실점 표시 (이미지 내부: 원, 외부: 방향 화살표)
    vp_ix, vp_iy = int(vp_x), int(vp_y)
    if 0 <= vp_ix < W and 0 <= vp_iy < H:
        cv2.circle(img, (vp_ix, vp_iy), 20, (0, 255, 255), 3)
        cv2.putText(img, "VP", (vp_ix + 25, vp_iy),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2)
    else:
        ax_s, ay_s = W // 2, H // 5
        dx, dy = vp_x - img_cx, vp_y - img_cy
        n = (dx ** 2 + dy ** 2) ** 0.5
        if n > 0:
            ax_e = max(0, min(W - 1, int(ax_s + dx / n * 80)))
            ay_e = max(0, min(H - 1, int(ay_s + dy / n * 80)))
            cv2.arrowedLine(img, (ax_s, ay_s), (ax_e, ay_e), (0, 255, 255), 3)
        cv2.putText(img, f"VP({vp_ix},{vp_iy})", (ax_s + 5, ay_s - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2)

    # 4. 라멘 그룹 강조: 바운딩 박스 + 그룹 ID + 판정 라벨
    for g in valid_items:
        verdict = g['verdict']
        color = _RAAMEN_COLOR[verdict]
        all_idx = g['B'] + g['C']
        all_pts = np.vstack([polys[i] for i in all_idx]).astype(np.int32)
        x0 = all_pts[:, 0].min() - 15; y0 = all_pts[:, 1].min() - 15
        x1 = all_pts[:, 0].max() + 15; y1 = all_pts[:, 1].max() + 15
        cv2.rectangle(img, (x0, y0), (x1, y1), color, 4)
        lbl = f"G{g['id']} {verdict} ({g['n_poles']}poles)"
        cv2.putText(img, lbl, (x0, y0 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 0), 6)
        cv2.putText(img, lbl, (x0, y0 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 1.2, color, 2)

    # 5. 폴리곤 라벨 — 모든 그래픽 최상단에 그리기 (병합 흡수된 원본 제외)
    for i, (pts, orient) in enumerate(zip(polys, orients)):
        if i in merged_away:
            continue
        color = _CB_COLOR[orient]
        cx, cy = int(pts[:, 0].mean()), int(pts[:, 1].mean())
        lbl = f"{i}{orient}"
        cv2.putText(img, lbl, (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 4)
        cv2.putText(img, lbl, (cx, cy), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)

    # 이미지 저장
    scale = min(1.0, 4096 / max(H, W))
    if scale < 1.0:
        img = cv2.resize(img, (int(W * scale), int(H * scale)))
    output_path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imencode(output_path.suffix, img)[1].tofile(str(output_path))
    print(f"\n  → {output_path}")

    # AnyLabeling JSON 저장
    import json

    def _shape(label, points, gid, desc):
        return {"label": label, "score": None, "points": points,
                "group_id": gid, "description": desc,
                "shape_type": "polygon", "flags": None}

    shapes = []
    for g in valid_items:
        gid     = g['id']
        verdict = g['verdict']
        desc_base = f"G{gid} {verdict}"

        def abs_ids(rel_list):
            """상대 폴리곤 인덱스 → 절대 JSON shapes[] 인덱스 변환. 병합 폴리곤은 list."""
            result = []
            for i in rel_list:
                v = poly_abs_idx[i]
                if isinstance(v, list):
                    result.extend(v)
                else:
                    result.append(v)
            return sorted(result)

        if not g['B']:
            # RAAMEN_CENTER: C 폴리곤만 있는 중앙 영역 그룹
            for ci in g['C']:
                shapes.append(_shape("raamen_center",
                                     [[float(p[0]), float(p[1])] for p in polys[ci]],
                                     gid, f"{desc_base} shape#{poly_abs_idx[ci]}"))
            continue

        det = group_detail(g, polys, W)

        # 주 빔 (1~2개: 면적 기준 상위, 2nd ≥ 50% 조건 충족 시 포함)
        for bi, mb in enumerate(det['main_bs'], 1):
            label = "raamen_beam" if bi == 1 else "raamen_beam2"
            shapes.append(_shape(label,
                                 _merge_poly_hull([mb], polys),
                                 gid, f"{desc_base} beam{bi} shape#{poly_abs_idx[mb]}"))

        # 잡 B 폴리곤 클러스터 (브라켓 등 부속)
        if det['junk_b']:
            junk_clusters = _cluster_polys(det['junk_b'], polys)
            for jc in junk_clusters:
                shapes.append(_shape("raamen_beam_sub",
                                     _merge_poly_hull(jc, polys),
                                     gid, f"{desc_base} beam_sub{abs_ids(jc)}"))

        # 유효 기둥 클러스터 (끝단)
        for pi, cluster in enumerate(det['valid_pole_clusters'], 1):
            shapes.append(_shape("raamen_pole",
                                 _merge_poly_hull(cluster, polys),
                                 gid, f"{desc_base} pole{pi}{abs_ids(cluster)}"))

        # 중앙 부속물 클러스터 (기둥 아님)
        for cluster in det['attach_clusters']:
            shapes.append(_shape("raamen_pole_attach",
                                 _merge_poly_hull(cluster, polys),
                                 gid, f"{desc_base} attach{abs_ids(cluster)}"))

    anylabel_json = {
        "version": "3.3.9",
        "flags": {},
        "shapes": shapes,
        "imagePath": image_path.name,
        "imageData": None,
        "imageHeight": H,
        "imageWidth": W,
    }
    json_path = output_path.with_suffix(".json")
    json_path.write_text(json.dumps(anylabel_json, ensure_ascii=False, indent=2),
                         encoding="utf-8")
    print(f"  → {json_path}")

    # ── 원본 폴리곤 분류 JSON 저장 ───────────────────────────────────────
    # 입력 catenary_pole 폴리곤마다 그룹/타입 레이블 부여
    # 그룹 소속: catenary_pole_B / catenary_pole_C  (group_id=G번호)
    # 그룹 미소속: catenary_pole  (원본 그대로)

    abs_to_info = {}  # abs_idx → (group_id, verdict, 'B'|'C')
    for g in groups:
        v = g.get('verdict', '')
        for type_char, members in (('B', g['B']), ('C', g['C'])):
            for idx in members:
                av = poly_abs_idx[idx]
                for a in (av if isinstance(av, list) else [av]):
                    abs_to_info[a] = (g['id'], v, type_char)

    input_data = json.loads(label_path.read_text(encoding="utf-8"))
    input_shapes = input_data.get("shapes", [])

    cls_shapes = []
    for abs_idx, s in enumerate(input_shapes):
        if class_names and s.get("label", "") not in class_names:
            continue
        if abs_idx in abs_to_info:
            gid, verdict, type_char = abs_to_info[abs_idx]
            lbl  = f"catenary_pole_{type_char}"
            desc = f"G{gid} {verdict}"
            gid_out = gid
        else:
            lbl     = s.get("label", "catenary_pole")
            desc    = s.get("description", "")
            gid_out = None
        cls_shapes.append({
            "label":      lbl,
            "score":      s.get("score"),
            "points":     s.get("points", []),
            "group_id":   gid_out,
            "description": desc,
            "shape_type": s.get("shape_type", "polygon"),
            "flags":      s.get("flags"),
        })

    cls_json = {
        "version":     input_data.get("version", "3.3.9"),
        "flags":       input_data.get("flags", {}),
        "shapes":      cls_shapes,
        "imagePath":   input_data.get("imagePath", image_path.name),
        "imageData":   None,
        "imageHeight": H,
        "imageWidth":  W,
    }
    cls_path = output_path.parent / (label_path.stem + "_classified.json")
    cls_path.write_text(json.dumps(cls_json, ensure_ascii=False, indent=2),
                        encoding="utf-8")
    print(f"  → {cls_path}  ({len(cls_shapes)}개 폴리곤)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image",      required=True)
    ap.add_argument("--label",      required=True)
    ap.add_argument("--output",     default=None,
                    help="출력 jpg 경로 (기본: output/raamen/{이미지명}/{이미지명}_raamen_NNN.jpg)")
    ap.add_argument("--class-ids",  default="",
                    help="포함할 클래스 ID, 콤마 구분 (.txt 전용)")
    ap.add_argument("--class-names", default="catenary_pole",
                    help="포함할 클래스 이름, 콤마 구분 (.json 전용, 기본: 'catenary_pole')")
    ap.add_argument("--epsilon",    type=float, default=4.0,
                    help="approxPolyDP epsilon (기본 4.0px)")
    ap.add_argument("--c-thresh",   type=float, default=20.0,
                    help="C/B 분류 각도 임계값 degrees (기본 20°)")
    ap.add_argument("--b-max-diff", type=float, default=75.0,
                    help="B(빔) 최대 각도 diff; 이 이상은 레일/전선 등으로 제외 (기본 75°)")
    ap.add_argument("--margin",         type=int,   default=30,
                    help="폴리곤 접촉 판정 margin px (기본 30)")
    ap.add_argument("--min-group-area", type=float, default=0,
                    help="라멘 그룹의 최소 면적 합계 (기본 0)")
    args = ap.parse_args()

    class_ids = ({int(x) for x in args.class_ids.split(',') if x.strip()}
                 if args.class_ids else None)
    class_names = ({x.strip() for x in args.class_names.split(',') if x.strip()}
                   if args.class_names else None)

    if args.output:
        out_path = Path(args.output)
    else:
        img_stem = Path(args.image).stem
        out_dir = Path("output") / "raamen" / img_stem
        out_dir.mkdir(parents=True, exist_ok=True)
        n = 1
        while True:
            out_path = out_dir / f"{img_stem}_raamen_{n:03d}.jpg"
            if not out_path.exists():
                break
            n += 1

    render(Path(args.image), Path(args.label), out_path, args,
           class_ids=class_ids, class_names=class_names,
           epsilon=args.epsilon, c_thresh=args.c_thresh,
           b_max_diff=args.b_max_diff)


if __name__ == "__main__":
    main()