Aislo/B04_PreProcess/B04_PreProcess_Engine_ModelBuild.py

"""B04 지표면 표현 빌더 (DTM/implicit).

TerrainContext에서 파생 격자·샘플을 받아 각 표현의 모델(npz)과 프리뷰
(GLB)를 저장하고 메타데이터를 반환한다.
"""

from contextlib import suppress
from pathlib import Path
from typing import Any

import numpy as np
from scipy.interpolate import RBFInterpolator

from B04_PreProcess.B04_PreProcess_Engine_ModelContext import (
    ProgressCallback,
    TerrainContext,
    artifact_size,
    atomic_npz,
    clip_and_compact_mesh,
    grid_faces,
    grid_vertices,
    with_footprint,
    write_glb,
)
from common_util.common_util_surface_sampler import implicit_grid_npz


def build_dtm(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    progress(10)
    resolution = float(context.config["dtm_grid_resolution_meters"])
    x_coords, y_coords, z_grid = context.grid(resolution)
    progress(55)
    preview_x, preview_y, preview_z = context.preview_grid(resolution)
    vertices = grid_vertices(preview_x, preview_y, preview_z)
    faces = grid_faces(len(preview_y), len(preview_x))
    valid_grid = context.contains_xy(*np.meshgrid(x_coords, y_coords)).reshape(
        len(y_coords), len(x_coords)
    )
    preview_valid = context.contains_xy(vertices[:, 0], vertices[:, 1])
    vertices, faces = clip_and_compact_mesh(vertices, faces, preview_valid)
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.glb"
    atomic_npz(
        model_path,
        x=x_coords,
        y=y_coords,
        z=z_grid,
        valid_mask=valid_grid,
        filled_mask=context.filled(resolution) & valid_grid,
        resolution=np.array([resolution], np.float32),
    )
    progress(75)
    write_glb(preview_path, vertices, faces, context.bounds)
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "regular_grid",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "model/gltf-binary",
            "grid_rows": int(len(y_coords)),
            "grid_columns": int(len(x_coords)),
            "grid_resolution_meters": resolution,
            "vertex_count": int(len(vertices)),
            "face_count": int(len(faces)),
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


def implicit_point_count(ground_count: int, config: dict[str, Any]) -> int:
    """implicit 중심점 수 = clip(지면점 × 비율, 하한, 상한) · 지면점이 하한보다 적으면 전부."""
    minimum = int(config["implicit_min_points"])
    if ground_count <= minimum:
        return int(ground_count)
    ratio = round(ground_count * float(config["implicit_point_ratio"]))
    return int(min(max(ratio, minimum), int(config["implicit_max_points"])))


def build_implicit(
    context: TerrainContext, output_dir: Path, stem: str, progress: ProgressCallback
) -> dict[str, Any]:
    count = implicit_point_count(context.source_count, context.config)
    # 구멍(766-2)은 메운 dtm 칸을 중심점으로 더해 dtm 과 같은 면으로 이음
    points = np.concatenate([context.grid_sample(count), context.hole_guides()])
    unique_xy, unique_indices = np.unique(points[:, :2], axis=0, return_index=True)
    points = points[unique_indices]
    if len(points) < 4:
        raise ValueError("Implicit 생성에 필요한 포인트가 부족합니다.")
    progress(20)
    interpolator = RBFInterpolator(
        unique_xy.astype(np.float64),
        points[:, 2].astype(np.float64),
        neighbors=min(64, len(points)),
        smoothing=float(context.config["implicit_smoothing"]),
        kernel="thin_plate_spline",
    )
    x_preview, y_preview, _ = context.preview_grid(
        float(context.config["dtm_grid_resolution_meters"])
    )
    xx, yy = np.meshgrid(x_preview, y_preview)
    query = np.column_stack([xx.ravel(), yy.ravel()])
    z_values = np.empty(len(query), dtype=np.float32)
    for start in range(0, len(query), 50_000):
        end = min(start + 50_000, len(query))
        z_values[start:end] = interpolator(query[start:end]).astype(np.float32)
        progress(25 + int(45 * end / len(query)))
    z_grid = z_values.reshape(len(y_preview), len(x_preview))
    vertices = grid_vertices(x_preview, y_preview, z_grid)
    faces = grid_faces(len(y_preview), len(x_preview))
    valid_preview = context.contains_xy(vertices[:, 0], vertices[:, 1])
    vertices, faces = clip_and_compact_mesh(vertices, faces, valid_preview)
    model_path = output_dir / f"{stem}.npz"
    preview_path = output_dir / f"{stem}_preview.glb"
    atomic_npz(
        model_path,
        centers_xy=unique_xy.astype(np.float32),
        center_z=points[:, 2].astype(np.float32),
        smoothing=np.array([float(context.config["implicit_smoothing"])], np.float32),
    )
    write_glb(preview_path, vertices, faces, context.bounds)
    with suppress(OSError, ValueError):  # 641-1 굳힌 격자 — 못 지으면 처음 읽을 때 지음
        implicit_grid_npz(output_dir, stem.removeprefix("implicit_"))
    progress(100)
    return with_footprint(
        context,
        {
            "representation": "local_rbf_height_field",
            "model_file": model_path.name,
            "preview_file": preview_path.name,
            "preview_media_type": "model/gltf-binary",
            "center_count": int(len(points)),
            "vertex_count": int(len(vertices)),
            "face_count": int(len(faces)),
            "artifact_bytes": artifact_size(model_path, preview_path),
        },
    )


BUILDERS = {
    "dtm": build_dtm,
    "implicit": build_implicit,
}