176 lines
5.8 KiB
Python
176 lines
5.8 KiB
Python
"""
|
||
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}")
|