Files
railway-client/tools/detect_hollow_section.py
minsung ccba1266b5 프로젝트 분리 이동
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-05-20 14:28:27 +09:00

153 lines
5.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
얇은 중공단면(도너츠 형태) 검출 및 표시 스크립트
드론 영상에서 전철주 상단 원형 단면을 찾아 표시
사용법:
python detect_hollow_section.py <image> [--radius <px>] [--tol <px>] [--topk <n>]
--radius : 3D 프로그램 줌 기준 단면 반지름 (픽셀). 미지정 시 자동 탐색.
--tol : radius ± 허용 오차 (기본 30px)
--topk : 표시할 최대 후보 수 (기본 3)
"""
import argparse
import cv2
import numpy as np
from pathlib import Path
def ring_score(edges: np.ndarray, cx: int, cy: int, r: int,
H: int, W: int, inner_ratio: float = 0.60) -> float:
"""
링(도넛) 특성 점수.
- 링 영역(외곽~내곽) 엣지 강도 / 내부 엣지 강도 비율 (최대 5 클리핑)
- 이미지 중심 근접도 보정 (중심에 가까울수록 가산점)
Canny edges는 외부에서 한 번만 계산해 전달받음.
"""
inner_r = max(int(r * inner_ratio), 3)
mask_full = np.zeros((H, W), np.uint8)
mask_inner = np.zeros((H, W), np.uint8)
cv2.circle(mask_full, (cx, cy), r, 255, -1)
cv2.circle(mask_inner, (cx, cy), inner_r, 255, -1)
mask_ring = cv2.subtract(mask_full, mask_inner)
if mask_ring.any() == 0 or mask_inner.any() == 0:
return 0.0
ring_edge = float(edges[mask_ring > 0].mean())
inner_edge = float(edges[mask_inner > 0].mean()) + 1e-6
edge_ratio = min(ring_edge / inner_edge, 5.0)
dist_from_center = np.hypot(cx - W / 2, cy - H / 2)
max_dist = np.hypot(W / 2, H / 2)
center_bonus = 1.0 - dist_from_center / max_dist # 0~1
return edge_ratio * (0.7 + 0.3 * center_bonus)
def detect_hollow_section(
image_path: str,
output_path: str | None = None,
radius_px: int | None = None, # 3D 프로그램 줌 기준 알려진 반지름
tol_px: int = 30, # radius_px ± 허용 오차
top_k: int = 3,
) -> str:
img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
if img is None:
raise FileNotFoundError(f"이미지 로드 실패: {image_path}")
H, W = img.shape[:2]
print(f"이미지 크기: {W}×{H}")
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
blurred = cv2.GaussianBlur(clahe.apply(gray), (7, 7), 2)
# Canny 엣지를 미리 한 번만 계산 (ring_score에서 재사용)
edges = cv2.Canny(blurred, 30, 90)
# 반경 범위 결정
if radius_px is not None:
min_r = max(5, radius_px - tol_px)
max_r = radius_px + tol_px
else:
short = min(W, H)
min_r = max(10, short // 20)
max_r = short // 3
print(f"탐색 반경: {min_r} ~ {max_r} px"
+ (f" (기준={radius_px}px ±{tol_px})" if radius_px else " (자동)"))
circles = cv2.HoughCircles(
blurred,
cv2.HOUGH_GRADIENT,
dp=1.0,
minDist=min_r * 3,
param1=80,
param2=35,
minRadius=min_r,
maxRadius=max_r,
)
result = img.copy()
if circles is None:
print("HoughCircles 미검출")
else:
circles = np.round(circles[0]).astype(int)
print(f"HoughCircles 후보: {len(circles)}개 → 링 스코어 계산 중...")
scored = sorted(
[(ring_score(edges, cx, cy, r, H, W), cx, cy, r)
for cx, cy, r in circles],
reverse=True,
)
# 1위: 녹색(굵게), 2위: 노랑, 3위: 파랑
colors = [(0, 255, 0), (0, 220, 255), (255, 120, 0)]
print(f"\n상위 {min(top_k, len(scored))}개 중공단면 후보:")
for rank, (score, cx, cy, r) in enumerate(scored[:top_k]):
color = colors[rank % len(colors)]
inner_r = max(int(r * 0.60), 5)
lw = 3 if rank == 0 else 2
cv2.circle(result, (cx, cy), r, color, lw)
cv2.circle(result, (cx, cy), inner_r, color, 1)
cv2.circle(result, (cx, cy), 6, (0, 0, 255), -1)
arm = r + 15
cv2.line(result, (cx - arm, cy), (cx + arm, cy), color, 1)
cv2.line(result, (cx, cy - arm), (cx, cy + arm), color, 1)
cv2.putText(result, f"#{rank+1} r={r}px s={score:.2f}",
(cx - r, cy - r - 8),
cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 255, 255), 2)
print(f" #{rank+1}: 중심=({cx},{cy}), 반지름={r}px, 링스코어={score:.3f}")
if output_path is None:
p = Path(image_path)
output_path = str(p.parent / (p.stem + "_hollow_detected" + p.suffix))
cv2.imencode(Path(output_path).suffix, result)[1].tofile(output_path)
print(f"\n결과 저장: {output_path}")
return output_path
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="얇은 중공단면 검출")
parser.add_argument("image", nargs="?",
default=r"d:\MYCLAUDE_PROJECT\x-anylabeling01\data\Message_2026-04-23T11_13_29+09_00.png")
parser.add_argument("--radius", type=int, default=None,
help="3D 프로그램 줌 기준 단면 반지름(px). 미지정 시 자동 탐색.")
parser.add_argument("--tol", type=int, default=30,
help="radius ± 허용 오차 (기본 30px)")
parser.add_argument("--topk", type=int, default=3,
help="표시할 최대 후보 수 (기본 3)")
args = parser.parse_args()
detect_hollow_section(
args.image,
radius_px=args.radius,
tol_px=args.tol,
top_k=args.topk,
)