Files
Aislo/utils/terrain_skeleton.py
T
2026-07-05 14:05:22 +09:00

333 lines
13 KiB
Python

"""지형 스켈레톤(주/지 능선·계곡) 추출 모듈.
"물이 모이는 곳 = 계곡, 나머지 = 능선" 이라는 수문학적 정의를 따른다:
- D8 흐름누적(flow accumulation)이 임계값 이상인 셀 = 계곡
- DEM 을 반전(z' = -z)해 같은 계산을 하면 분수계(능선) 셀을 얻는다
- 누적값 크기의 2차 임계값으로 주(main)/지(minor) 를 나눈다
흐름누적은 requirements 에 이미 있는 whitebox(WhiteboxTools) 를 우선 사용하고,
실행 실패 시 numpy 기반 자체 D8 구현으로 폴백한다.
산출물은 route_solver 의 비용면(2m 격자)과 동일한 좌표계의 polyline 목록이며,
route_design/terrain_skeleton_{filter}_{method}{_smooth}.json 에 캐시된다.
"""
import json
import math
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import config
from utils.route_solver import _load_or_build_cost_surface, _cost_surface_signature
# 스켈레톤 클래스 키 (산출 dict 의 키와 동일)
SKELETON_CLASSES = ("main_ridge", "minor_ridge", "main_valley", "minor_valley")
# D8 이웃: (dr, dc), 대각선 여부
_D8_OFFSETS = [
(-1, -1), (-1, 0), (-1, 1),
(0, -1), (0, 1),
(1, -1), (1, 0), (1, 1),
]
def _default_thresholds() -> Dict[str, float]:
return {
"valley_acc": float(config.SKELETON_VALLEY_ACC_THRESHOLD_CELLS),
"main_valley_acc": float(config.SKELETON_MAIN_VALLEY_ACC_THRESHOLD_CELLS),
"ridge_acc": float(config.SKELETON_RIDGE_ACC_THRESHOLD_CELLS),
"main_ridge_acc": float(config.SKELETON_MAIN_RIDGE_ACC_THRESHOLD_CELLS),
}
def d8_flow_accumulation_numpy(z_grid: np.ndarray, valid_mask: np.ndarray) -> np.ndarray:
"""numpy 기반 D8 흐름누적 (whitebox 폴백).
각 셀은 8방향 이웃 중 하강경사가 가장 급한 이웃 하나로 전량 흘러가고,
누적값은 자기 자신 1 + 상류 기여 합이다. 표고 내림차순 처리라서 별도의
위상정렬 없이 한 번의 순회로 계산된다. 함몰지(depression) 채움은 하지
않는다 — 스켈레톤 분류 용도로는 국지 함몰에서 누적이 끊겨도 충분하다.
"""
rows, cols = z_grid.shape
n = rows * cols
z_flat = z_grid.ravel()
valid_flat = valid_mask.ravel()
# 각 셀의 수용(receiver) 셀 인덱스 (-1 = 유출 없음/무효)
receiver = np.full(n, -1, dtype=np.int64)
for idx in range(n):
if not valid_flat[idx]:
continue
r, c = divmod(idx, cols)
zc = z_flat[idx]
best_slope = 0.0
best = -1
for dr, dc in _D8_OFFSETS:
nr, nc = r + dr, c + dc
if not (0 <= nr < rows and 0 <= nc < cols):
continue
nidx = nr * cols + nc
if not valid_flat[nidx]:
continue
drop = zc - z_flat[nidx]
if drop <= 0:
continue
dist = math.sqrt(2.0) if (dr != 0 and dc != 0) else 1.0
slope = drop / dist
if slope > best_slope:
best_slope = slope
best = nidx
receiver[idx] = best
acc = np.where(valid_flat, 1.0, 0.0)
order = np.argsort(-z_flat, kind="stable") # 높은 곳부터 아래로 흘려보냄
for idx in order:
recv = receiver[idx]
if recv >= 0 and valid_flat[idx]:
acc[recv] += acc[idx]
return acc.reshape(rows, cols)
def _d8_flow_accumulation_whitebox(
z_grid: np.ndarray, x_coords: np.ndarray, y_coords: np.ndarray,
valid_mask: np.ndarray, grid_res: float,
) -> Optional[np.ndarray]:
"""WhiteboxTools 로 D8 흐름누적을 계산한다. 실패 시 None (폴백 유도).
rasterio 로 임시 GeoTIFF 를 쓰고(북상향 정렬을 위해 행 뒤집기),
fill_depressions → d8_flow_accumulation(out_type=cells) 순으로 실행한다.
"""
try:
import rasterio
from rasterio.transform import from_origin
from whitebox import WhiteboxTools
except Exception as exc:
print(f"[terrain_skeleton] whitebox/rasterio 사용 불가 → numpy 폴백: {exc}")
return None
nodata = -9999.0
rows, cols = z_grid.shape
# GeoTIFF 는 북상향(위쪽 행이 큰 y)이므로 y 오름차순 격자를 뒤집어 저장한다.
z_out = np.where(valid_mask, z_grid, nodata).astype(np.float32)[::-1, :]
transform = from_origin(
float(x_coords[0]) - grid_res / 2.0,
float(y_coords[-1]) + grid_res / 2.0,
grid_res, grid_res,
)
try:
with tempfile.TemporaryDirectory(prefix="wbt_skel_") as tmp:
tmp_path = Path(tmp)
dem_tif = tmp_path / "dem.tif"
filled_tif = tmp_path / "filled.tif"
acc_tif = tmp_path / "acc.tif"
with rasterio.open(
dem_tif, "w", driver="GTiff", height=rows, width=cols, count=1,
dtype="float32", nodata=nodata, transform=transform,
) as dst:
dst.write(z_out, 1)
wbt = WhiteboxTools()
wbt.set_verbose_mode(False)
wbt.set_working_dir(str(tmp_path))
if wbt.fill_depressions("dem.tif", "filled.tif") != 0:
raise RuntimeError("fill_depressions 실패")
if wbt.d8_flow_accumulation("filled.tif", "acc.tif", out_type="cells") != 0:
raise RuntimeError("d8_flow_accumulation 실패")
with rasterio.open(acc_tif) as src:
acc = src.read(1).astype(np.float64)[::-1, :] # 다시 y 오름차순으로
acc = np.where(np.isfinite(acc) & (acc > 0) & valid_mask, acc, 0.0)
return acc
except Exception as exc:
print(f"[terrain_skeleton] whitebox 흐름누적 실패 → numpy 폴백: {exc}")
return None
def _flow_accumulation(
z_grid: np.ndarray, x_coords: np.ndarray, y_coords: np.ndarray,
valid_mask: np.ndarray, grid_res: float,
) -> np.ndarray:
acc = _d8_flow_accumulation_whitebox(z_grid, x_coords, y_coords, valid_mask, grid_res)
if acc is None:
acc = d8_flow_accumulation_numpy(z_grid, valid_mask)
return acc
def _trace_polylines(mask: np.ndarray) -> List[List[Tuple[int, int]]]:
"""1픽셀 폭 스켈레톤 마스크를 (r, c) polyline 목록으로 변환한다.
끝점(이웃 1개)과 분기점(이웃 3개 이상)에서 경로를 시작해 다음 분기/끝점까지
걷는 표준 추적 방식. 순수 루프(끝점/분기점 없는 폐곡선)는 임의 지점에서
한 바퀴 추적한다.
"""
pixels = set(zip(*np.nonzero(mask)))
if not pixels:
return []
def neighbors(p):
r, c = p
return [(r + dr, c + dc) for dr, dc in _D8_OFFSETS if (r + dr, c + dc) in pixels]
degree = {p: len(neighbors(p)) for p in pixels}
# 경로 시작점: 끝점과 분기점
seeds = [p for p in pixels if degree[p] != 2]
visited_edges = set() # 정규화된 (p, q) 픽셀 간선
polylines: List[List[Tuple[int, int]]] = []
def edge_key(a, b):
return (a, b) if a <= b else (b, a)
def walk(start, nxt):
"""start → nxt 방향으로 다음 끝점/분기점까지 걷는다."""
path = [start, nxt]
visited_edges.add(edge_key(start, nxt))
prev, curr = start, nxt
while degree[curr] == 2:
candidates = [q for q in neighbors(curr) if q != prev]
if not candidates:
break
q = candidates[0]
if edge_key(curr, q) in visited_edges:
break
visited_edges.add(edge_key(curr, q))
path.append(q)
prev, curr = curr, q
return path
for seed in seeds:
for nb in neighbors(seed):
if edge_key(seed, nb) not in visited_edges:
polylines.append(walk(seed, nb))
# 남은 순수 루프 처리
for p in pixels:
if degree[p] == 2:
for nb in neighbors(p):
if edge_key(p, nb) not in visited_edges:
polylines.append(walk(p, nb))
# 픽셀 2개 미만 조각 제거
return [pl for pl in polylines if len(pl) >= 2]
def _mask_to_polylines(
mask: np.ndarray, x_coords: np.ndarray, y_coords: np.ndarray, z_grid: np.ndarray,
) -> List[Dict[str, Any]]:
"""셀 마스크를 세선화한 뒤 모델좌표 polyline 목록으로 변환한다."""
if not mask.any():
return []
try:
from skimage.morphology import skeletonize
skel = skeletonize(mask)
except Exception as exc:
print(f"[terrain_skeleton] skeletonize 실패(마스크 그대로 사용): {exc}")
skel = mask
out = []
for pixel_path in _trace_polylines(skel):
poly = [
[float(x_coords[c]), float(y_coords[r]), float(z_grid[r, c])]
for r, c in pixel_path
]
out.append({"polyline": poly})
return out
def extract_skeleton_from_grid(
x_coords: np.ndarray,
y_coords: np.ndarray,
z_grid: np.ndarray,
valid_mask: np.ndarray,
grid_res: float,
thresholds: Optional[Dict[str, float]] = None,
use_whitebox: bool = True,
) -> Dict[str, List[Dict[str, Any]]]:
"""격자에서 주/지 능선·계곡 polyline 을 추출한다 (캐시 없음, 테스트용 공개 API)."""
th = thresholds or _default_thresholds()
if use_whitebox:
acc_valley = _flow_accumulation(z_grid, x_coords, y_coords, valid_mask, grid_res)
acc_ridge = _flow_accumulation(-z_grid, x_coords, y_coords, valid_mask, grid_res)
else:
acc_valley = d8_flow_accumulation_numpy(z_grid, valid_mask)
acc_ridge = d8_flow_accumulation_numpy(-z_grid, valid_mask)
# 계곡이 우선: 계곡으로 분류된 셀은 능선에서 제외한다 (정의상 배타).
valley_mask = valid_mask & (acc_valley >= th["valley_acc"])
main_valley_mask = valley_mask & (acc_valley >= th["main_valley_acc"])
minor_valley_mask = valley_mask & ~main_valley_mask
ridge_mask = valid_mask & ~valley_mask & (acc_ridge >= th["ridge_acc"])
main_ridge_mask = ridge_mask & (acc_ridge >= th["main_ridge_acc"])
minor_ridge_mask = ridge_mask & ~main_ridge_mask
return {
"main_ridge": _mask_to_polylines(main_ridge_mask, x_coords, y_coords, z_grid),
"minor_ridge": _mask_to_polylines(minor_ridge_mask, x_coords, y_coords, z_grid),
"main_valley": _mask_to_polylines(main_valley_mask, x_coords, y_coords, z_grid),
"minor_valley": _mask_to_polylines(minor_valley_mask, x_coords, y_coords, z_grid),
}
def _skeleton_signature(terrain_dir: Path, filter_key: str, method: str, smooth: bool) -> str:
"""소스 모델 + 격자 해상도 + 분류 임계값이 바뀌면 캐시를 재생성하도록 하는 서명."""
th = _default_thresholds()
th_part = "|".join(f"{k}={v}" for k, v in sorted(th.items()))
return _cost_surface_signature(terrain_dir, filter_key, method, smooth) + "|" + th_part
def load_or_build_skeleton(
terrain_dir: Path,
filter_key: str,
method: str,
smooth: bool,
) -> Dict[str, Any]:
"""스켈레톤을 캐시에서 로드하거나 새로 계산해 저장한다.
반환 dict: {"main_ridge": [...], "minor_ridge": [...], "main_valley": [...],
"minor_valley": [...], "grid_res": float}
"""
terrain_dir = Path(terrain_dir)
cache_dir = terrain_dir.parent / "route_design"
suffix = "_smooth" if smooth else ""
cache_path = cache_dir / f"terrain_skeleton_{filter_key}_{method}{suffix}.json"
signature = _skeleton_signature(terrain_dir, filter_key, method, smooth)
if cache_path.exists():
try:
with open(cache_path, "r", encoding="utf-8") as f:
cached = json.load(f)
if cached.get("signature") == signature:
return cached
except Exception as exc:
print(f"[terrain_skeleton] 캐시 무시(재생성): {exc}")
import time
t0 = time.time()
(x_coords, y_coords, z_grid, valid_mask, _dz_dx, _dz_dy, grid_res) = (
_load_or_build_cost_surface(terrain_dir, filter_key, method, smooth)
)
z_grid = np.asarray(z_grid, dtype=np.float64)
valid_mask = np.asarray(valid_mask, dtype=bool)
skeleton = extract_skeleton_from_grid(
np.asarray(x_coords), np.asarray(y_coords), z_grid, valid_mask, float(grid_res)
)
result: Dict[str, Any] = dict(skeleton)
result["grid_res"] = float(grid_res)
result["signature"] = signature
counts = {k: len(result[k]) for k in SKELETON_CLASSES}
print(f"[terrain_skeleton] built {counts} time={time.time()-t0:.2f}s "
f"[{method}/{filter_key}{suffix}]")
try:
cache_dir.mkdir(parents=True, exist_ok=True)
with open(cache_path, "w", encoding="utf-8") as f:
json.dump(result, f, ensure_ascii=False)
except Exception as exc: # 캐시 저장 실패는 치명적이지 않다
print(f"[terrain_skeleton] 캐시 저장 실패(무시): {exc}")
return result