302 lines
10 KiB
Python
302 lines
10 KiB
Python
"""
|
|
auto_rail_detect.py
|
|
===================
|
|
항공 이미지에서 레일 라인을 자동 검출하여 Rhino용 DXF로 저장.
|
|
수동 라벨링 없이 이미지 1장에서 바로 실행.
|
|
|
|
사용법:
|
|
python tools/auto_rail_detect.py <image.png> [output.dxf]
|
|
|
|
예시:
|
|
python tools/auto_rail_detect.py "경부선_2-2구간_미션33-34(5cm).png"
|
|
python tools/auto_rail_detect.py "경부선.png" output/rail.dxf
|
|
|
|
원리:
|
|
이미지 → 엣지 검출 → Hough 직선 검출 → 방향별 클러스터링 → DXF 폴리라인
|
|
"""
|
|
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
|
|
# ─── 설정 (조정 가능) ─────────────────────────────────────────────────────────
|
|
CANNY_LOW = 30 # Canny 엣지 낮은 임계값 (낮을수록 엣지 많이 검출)
|
|
CANNY_HIGH = 80 # Canny 엣지 높은 임계값
|
|
HOUGH_THRESHOLD = 100 # Hough 투표 임계값 (높을수록 긴 선만 검출)
|
|
HOUGH_MIN_LEN = 200 # 최소 선 길이 (픽셀)
|
|
HOUGH_MAX_GAP = 30 # 선 간격 허용 (픽셀, 클수록 끊긴 선도 연결)
|
|
ANGLE_TOLERANCE = 5 # 같은 방향으로 볼 각도 범위 (도)
|
|
CLUSTER_DIST = 20 # 같은 레일로 볼 선 간격 (픽셀)
|
|
MIN_TOTAL_LEN = 500 # 최종 레일 최소 총 길이 (픽셀, 짧은 잡선 제거)
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def detect_edges(img_gray):
|
|
"""CLAHE 대비 강화 → Gaussian 블러 → Canny 엣지"""
|
|
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
|
enhanced = clahe.apply(img_gray)
|
|
blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
|
|
edges = cv2.Canny(blurred, CANNY_LOW, CANNY_HIGH, apertureSize=3)
|
|
return edges
|
|
|
|
|
|
def detect_lines(edges):
|
|
"""확률적 Hough 변환으로 직선 검출"""
|
|
lines = cv2.HoughLinesP(
|
|
edges,
|
|
rho=1,
|
|
theta=np.pi / 180,
|
|
threshold=HOUGH_THRESHOLD,
|
|
minLineLength=HOUGH_MIN_LEN,
|
|
maxLineGap=HOUGH_MAX_GAP,
|
|
)
|
|
if lines is None:
|
|
return []
|
|
return [tuple(l[0]) for l in lines] # [(x1,y1,x2,y2), ...]
|
|
|
|
|
|
def line_angle(x1, y1, x2, y2):
|
|
"""선의 각도 (0~180도)"""
|
|
angle = np.degrees(np.arctan2(y2 - y1, x2 - x1)) % 180
|
|
return angle
|
|
|
|
|
|
def line_length(x1, y1, x2, y2):
|
|
return np.hypot(x2 - x1, y2 - y1)
|
|
|
|
|
|
def line_midpoint(x1, y1, x2, y2):
|
|
return ((x1 + x2) / 2, (y1 + y2) / 2)
|
|
|
|
|
|
def perpendicular_dist(x1, y1, x2, y2, px, py):
|
|
"""점 (px,py)에서 선 (x1,y1)-(x2,y2)까지 수직거리"""
|
|
dx, dy = x2 - x1, y2 - y1
|
|
length = np.hypot(dx, dy)
|
|
if length == 0:
|
|
return np.hypot(px - x1, py - y1)
|
|
return abs(dy * px - dx * py + x2 * y1 - y2 * x1) / length
|
|
|
|
|
|
def cluster_lines(lines):
|
|
"""비슷한 방향 + 가까운 위치의 선들을 같은 레일로 묶기"""
|
|
if not lines:
|
|
return []
|
|
|
|
# 각도별 그룹 먼저 분리
|
|
angle_groups = {}
|
|
for line in lines:
|
|
x1, y1, x2, y2 = line
|
|
ang = round(line_angle(x1, y1, x2, y2) / ANGLE_TOLERANCE) * ANGLE_TOLERANCE
|
|
angle_groups.setdefault(ang, []).append(line)
|
|
|
|
clusters = []
|
|
for ang, grp in angle_groups.items():
|
|
# 같은 방향 그룹 내에서 거리 기반 클러스터링
|
|
used = [False] * len(grp)
|
|
for i, line_i in enumerate(grp):
|
|
if used[i]:
|
|
continue
|
|
cluster = [line_i]
|
|
used[i] = True
|
|
mx_i, my_i = line_midpoint(*line_i)
|
|
for j, line_j in enumerate(grp):
|
|
if used[j]:
|
|
continue
|
|
mx_j, my_j = line_midpoint(*line_j)
|
|
dist = perpendicular_dist(*line_i, mx_j, my_j)
|
|
if dist < CLUSTER_DIST:
|
|
cluster.append(line_j)
|
|
used[j] = True
|
|
clusters.append(cluster)
|
|
|
|
return clusters
|
|
|
|
|
|
def merge_cluster_to_polyline(cluster):
|
|
"""클러스터의 선들을 하나의 정렬된 폴리라인으로 합치기"""
|
|
# 모든 끝점 수집
|
|
all_pts = []
|
|
for x1, y1, x2, y2 in cluster:
|
|
all_pts.append((x1, y1))
|
|
all_pts.append((x2, y2))
|
|
|
|
if not all_pts:
|
|
return []
|
|
|
|
# 주성분 방향으로 정렬
|
|
pts_arr = np.array(all_pts, dtype=float)
|
|
mean = pts_arr.mean(axis=0)
|
|
centered = pts_arr - mean
|
|
cov = np.cov(centered.T)
|
|
if cov.ndim < 2:
|
|
direction = np.array([1.0, 0.0])
|
|
else:
|
|
eigenvalues, eigenvectors = np.linalg.eig(cov)
|
|
direction = eigenvectors[:, np.argmax(eigenvalues)]
|
|
|
|
# 방향으로 투영하여 정렬
|
|
projections = centered.dot(direction)
|
|
sorted_idx = np.argsort(projections)
|
|
sorted_pts = pts_arr[sorted_idx]
|
|
|
|
# 중복 제거 (가까운 점 합치기)
|
|
merged = [sorted_pts[0]]
|
|
for pt in sorted_pts[1:]:
|
|
if np.hypot(pt[0] - merged[-1][0], pt[1] - merged[-1][1]) > 5:
|
|
merged.append(pt)
|
|
|
|
return merged
|
|
|
|
|
|
def total_polyline_length(polyline):
|
|
if len(polyline) < 2:
|
|
return 0
|
|
total = 0
|
|
for i in range(len(polyline) - 1):
|
|
total += np.hypot(
|
|
polyline[i+1][0] - polyline[i][0],
|
|
polyline[i+1][1] - polyline[i][1]
|
|
)
|
|
return total
|
|
|
|
|
|
def save_debug_image(img, all_polylines, output_path):
|
|
"""검출 결과를 이미지에 시각화 (확인용)"""
|
|
vis = img.copy() if len(img.shape) == 3 else cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
|
colors = [(0, 0, 255), (0, 255, 0), (255, 0, 0), (0, 255, 255),
|
|
(255, 0, 255), (255, 255, 0)]
|
|
for i, poly in enumerate(all_polylines):
|
|
color = colors[i % len(colors)]
|
|
pts = np.array([[int(p[0]), int(p[1])] for p in poly])
|
|
cv2.polylines(vis, [pts], False, color, 2)
|
|
# 번호 표시
|
|
cx, cy = int(poly[len(poly)//2][0]), int(poly[len(poly)//2][1])
|
|
cv2.putText(vis, str(i+1), (cx, cy), cv2.FONT_HERSHEY_SIMPLEX,
|
|
1.5, color, 3)
|
|
cv2.imwrite(str(output_path), vis)
|
|
print(f" 시각화 저장: {output_path}")
|
|
|
|
|
|
def process(image_path: str, dxf_path: str):
|
|
import ezdxf
|
|
|
|
print(f"[입력] {image_path}")
|
|
# 한글 경로 대응: numpy로 읽어서 cv2 디코딩
|
|
import numpy as np_io
|
|
with open(image_path, "rb") as f:
|
|
data = f.read()
|
|
img_arr = np_io.frombuffer(data, dtype=np_io.uint8)
|
|
img = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
print(f"오류: 이미지를 열 수 없습니다: {image_path}")
|
|
sys.exit(1)
|
|
|
|
H, W = img.shape[:2]
|
|
print(f"[이미지] {W} x {H} px")
|
|
|
|
# 전처리: 작은 이미지로 축소 (속도)
|
|
scale = 1.0
|
|
if W > 4000:
|
|
scale = 4000 / W
|
|
small = cv2.resize(img, (int(W * scale), int(H * scale)))
|
|
print(f" 축소: {int(W*scale)} x {int(H*scale)} (scale={scale:.2f})")
|
|
else:
|
|
small = img.copy()
|
|
|
|
gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)
|
|
|
|
# 엣지 검출
|
|
print("[1단계] 엣지 검출...")
|
|
edges = detect_edges(gray)
|
|
edge_count = int(edges.sum() / 255)
|
|
print(f" 엣지 픽셀: {edge_count:,}")
|
|
|
|
# Hough 직선 검출
|
|
print("[2단계] Hough 직선 검출...")
|
|
lines = detect_lines(edges)
|
|
print(f" 검출된 선: {len(lines)}개")
|
|
|
|
if not lines:
|
|
print("선을 검출하지 못했습니다. HOUGH_THRESHOLD를 낮춰보세요.")
|
|
sys.exit(1)
|
|
|
|
# 클러스터링
|
|
print("[3단계] 레일 클러스터링...")
|
|
clusters = cluster_lines(lines)
|
|
print(f" 클러스터: {len(clusters)}개")
|
|
|
|
# 폴리라인 변환 + 길이 필터
|
|
polylines = []
|
|
for cluster in clusters:
|
|
poly = merge_cluster_to_polyline(cluster)
|
|
length = total_polyline_length(poly)
|
|
if length >= MIN_TOTAL_LEN / scale:
|
|
# 원본 해상도로 역변환
|
|
poly_orig = [(p[0] / scale, p[1] / scale) for p in poly]
|
|
polylines.append((poly_orig, length / scale))
|
|
|
|
# 길이순 정렬
|
|
polylines.sort(key=lambda x: -x[1])
|
|
print(f" 유효 레일 라인: {len(polylines)}개")
|
|
|
|
if not polylines:
|
|
print("유효한 레일을 찾지 못했습니다. MIN_TOTAL_LEN을 낮춰보세요.")
|
|
sys.exit(1)
|
|
|
|
for i, (poly, length) in enumerate(polylines):
|
|
print(f" 레일 {i+1}: {length:.0f}px ({length*0.05:.1f}m)")
|
|
|
|
# DXF 저장
|
|
print("[4단계] DXF 저장...")
|
|
doc = ezdxf.new("R2010")
|
|
msp = doc.modelspace()
|
|
doc.layers.add("RAIL_AUTO", color=1) # 빨강 — 자동검출 레일
|
|
doc.layers.add("RAIL_AUTO_LABEL", color=2) # 노랑 — 번호
|
|
|
|
for i, (poly, length) in enumerate(polylines):
|
|
# Y축 반전 (이미지→DXF)
|
|
dxf_pts = [(float(p[0]), float(-p[1])) for p in poly]
|
|
msp.add_lwpolyline(dxf_pts, dxfattribs={"layer": "RAIL_AUTO"})
|
|
|
|
# 중간점에 번호 텍스트
|
|
mid = poly[len(poly) // 2]
|
|
msp.add_text(
|
|
f"Rail_{i+1}",
|
|
dxfattribs={
|
|
"layer": "RAIL_AUTO_LABEL",
|
|
"height": 20,
|
|
"insert": (float(mid[0]), float(-mid[1])),
|
|
}
|
|
)
|
|
|
|
Path(dxf_path).parent.mkdir(parents=True, exist_ok=True)
|
|
doc.saveas(dxf_path)
|
|
print(f"[완료] DXF 저장: {dxf_path}")
|
|
print(f" 레이어: RAIL_AUTO(빨강) = {len(polylines)}개 레일 중심선")
|
|
|
|
# 디버그 이미지 저장
|
|
debug_path = Path(dxf_path).parent / (Path(dxf_path).stem + "_debug.jpg")
|
|
all_polys = [p for p, _ in polylines]
|
|
save_debug_image(small, [([(x*scale, y*scale) for x, y in p]) for p in all_polys],
|
|
debug_path)
|
|
|
|
print(f"\nRhino 사용법:")
|
|
print(f" 1. File -> Import -> {Path(dxf_path).name}")
|
|
print(f" 2. RAIL_AUTO 레이어 선택")
|
|
print(f" 3. Sweep1 명령 -> 레일 단면 선택 -> 3D 레일 완성")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("사용법: python tools/auto_rail_detect.py <image.png> [output.dxf]")
|
|
sys.exit(1)
|
|
|
|
img_path = sys.argv[1]
|
|
dxf_path = sys.argv[2] if len(sys.argv) >= 3 else str(
|
|
Path(img_path).with_suffix(".dxf")
|
|
)
|
|
process(img_path, dxf_path)
|