프로젝트 분리 이동

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

tools/rail_centerline_dxf.py

@@ -0,0 +1,175 @@
+"""
+2cm GSD 드론 TIF 전체에서 레일 중심선 추출 → DXF 저장
+
+중심선 추출 방법:
+  masks.xy 폴리곤 → PCA 주축(길이 방향) 찾기
+  → 주축 방향으로 슬라이싱 → 폭 방향 중앙값
+  → 계단 현상 없는 매끄러운 중심선
+"""
+import numpy as np
+import rasterio
+from rasterio.windows import Window
+from ultralytics import YOLO
+from PIL import Image
+import ezdxf
+
+
+def polygon_to_centerline(xy_tile, spacing_px=10):
+    """
+    타일 픽셀 좌표 폴리곤 → 중심선 점 목록
+
+    Parameters
+    ----------
+    xy_tile    : list of (x, y)  타일 픽셀 좌표 (masks.xy × sx/sy)
+    spacing_px : int  슬라이싱 간격 [픽셀]  (10px = 0.2m at 2cm GSD)
+
+    Returns
+    -------
+    list of (x, y)  중심선 타일 픽셀 좌표
+    """
+    pts = np.asarray(xy_tile, dtype=float)
+    if len(pts) < 4:
+        return []
+
+    # PCA: 주축(레일 길이 방향) 탐색
+    center = pts.mean(axis=0)
+    _, _, Vt = np.linalg.svd(pts - center, full_matrices=False)
+    v_long = Vt[0]   # 길이 방향 단위벡터
+    v_perp = Vt[1]   # 폭 방향 단위벡터
+
+    local = pts - center
+    t = local @ v_long   # 길이 방향 좌표
+
+    t_min, t_max = t.min(), t.max()
+    if t_max - t_min < spacing_px:
+        # 너무 짧으면 무게중심 1점
+        return [center.tolist()]
+
+    n_bins = max(2, int((t_max - t_min) / spacing_px) + 1)
+    bins = np.linspace(t_min, t_max, n_bins + 1)
+
+    centerline = []
+    s_all = local @ v_perp   # 폭 방향 좌표 (전체)
+
+    for i in range(n_bins):
+        mask = (t >= bins[i]) & (t <= bins[i + 1])
+        if mask.sum() < 2:
+            continue
+        t_c = (bins[i] + bins[i + 1]) / 2
+        s_vals = s_all[mask]
+        # 폭 방향: 폴리곤 양 가장자리의 기하학적 중앙
+        s_c = (s_vals.max() + s_vals.min()) / 2
+        pt = center + t_c * v_long + s_c * v_perp
+        centerline.append(pt.tolist())
+
+    return centerline
+
+
+# ── 설정 ──────────────────────────────────────────
+SRC_TIF   = "drone_2cm/22)조치원(STA.127+570~131+300).tif"
+MODEL_PT  = "runs/segment/output/yolo_train_2cm/rail_seg_v2/weights/best.pt"
+OUT_DXF   = "output/rail_centerline_2cm_pca.dxf"
+TILE      = 1024
+OVERLAP   = 128
+STEP      = TILE - OVERLAP
+CONF      = 0.3
+BLACK_THR = 0.4    # 검은 픽셀 비율 임계값
+SPACING   = 10     # 중심선 샘플 간격 [픽셀]  10px ≈ 0.2m
+
+# ── 모델 로드 ──────────────────────────────────────
+print("모델 로드 중...")
+model = YOLO(MODEL_PT)
+
+# ── DXF 초기화 ────────────────────────────────────
+doc = ezdxf.new()
+msp = doc.modelspace()
+doc.layers.add("RAIL_CENTERLINE", color=3)
+
+# ── TIF 열기 ──────────────────────────────────────
+src = rasterio.open(SRC_TIF)
+W, H = src.width, src.height
+transform = src.transform
+print(f"이미지 크기: {W} x {H}")
+
+
+def pixel_to_world(px, py):
+    """픽셀 좌표 → 세계 좌표 (EPSG:5186)"""
+    wx = transform.c + px * transform.a
+    wy = transform.f + py * transform.e
+    return wx, wy
+
+
+# ── 타일 순회 ─────────────────────────────────────
+total_polylines = 0
+xs_range = range(0, W - TILE // 2, STEP)
+ys_range = range(0, H - TILE // 2, STEP)
+total_tiles = len(xs_range) * len(ys_range)
+processed = 0
+
+print(f"총 타일 수(예상): {total_tiles}, 처리 시작...")
+
+for ty in ys_range:
+    for tx in xs_range:
+        tw = min(TILE, W - tx)
+        th = min(TILE, H - ty)
+        if tw < 64 or th < 64:
+            continue
+
+        win = Window(tx, ty, tw, th)
+        data = src.read([1, 2, 3], window=win)
+        img_arr = np.transpose(data, (1, 2, 0)).astype(np.uint8)
+
+        black_ratio = np.mean(np.all(img_arr < 10, axis=2))
+        if black_ratio > BLACK_THR:
+            processed += 1
+            continue
+
+        pil_img = Image.fromarray(img_arr)
+        results = model.predict(pil_img, imgsz=TILE, conf=CONF,
+                                device='cuda:0', verbose=False)
+
+        if not results or results[0].masks is None:
+            processed += 1
+            continue
+
+        masks_data = results[0].masks.data.cpu().numpy()
+        h_r, w_r = masks_data.shape[1], masks_data.shape[2]
+        sx = tw / w_r
+        sy = th / h_r
+
+        for xy in results[0].masks.xy:
+            if len(xy) < 4:
+                continue
+
+            # YOLO 좌표 → 타일 픽셀 좌표
+            xy_tile = [(float(lx) * sx, float(ly) * sy) for lx, ly in xy]
+
+            # PCA 슬라이싱으로 중심선 추출 (타일 픽셀 좌표)
+            cl_tile = polygon_to_centerline(xy_tile, spacing_px=SPACING)
+            if len(cl_tile) < 2:
+                continue
+
+            # overlap 가장자리 제거 + 세계 좌표 변환
+            pts_world = []
+            for lx, ly in cl_tile:
+                if tx + tw < W and lx > (tw - OVERLAP // 2):
+                    continue
+                if ty + th < H and ly > (th - OVERLAP // 2):
+                    continue
+                pts_world.append(pixel_to_world(tx + lx, ty + ly))
+
+            if len(pts_world) >= 2:
+                msp.add_lwpolyline(pts_world,
+                                   dxfattribs={"layer": "RAIL_CENTERLINE"})
+                total_polylines += 1
+
+        processed += 1
+        if processed % 500 == 0:
+            pct = processed / total_tiles * 100
+            print(f"  진행: {processed}/{total_tiles} ({pct:.1f}%)"
+                  f" | 폴리라인: {total_polylines}")
+
+src.close()
+
+doc.saveas(OUT_DXF)
+print(f"\n완료: {total_polylines}개 폴리라인 → {OUT_DXF}")