Aislo/B04_PreProcess/B04_PreProcess_Engine_Structurize.py

"""B04 LAS/LAZ 고속 구조화 엔진."""

import os
import tempfile
from collections.abc import Callable
from pathlib import Path

import laspy
import numpy as np

from common_util.common_util_json import replace_with_retry
from config.config_system import SURFACE_DEFAULT_RGB_VALUE, SURFACE_LAS_CHUNK_SIZE


def structurize_las(
    las_path: str | Path,
    output_dir: str | Path,
    progress_callback: Callable[[int], None] | None = None,
) -> Path:
    """LAS/LAZ 속성을 청크로 읽어 B04 structured.npz로 원자적 저장한다."""
    source = Path(las_path)
    target_dir = Path(output_dir)
    target_dir.mkdir(parents=True, exist_ok=True)
    target = target_dir / "structured.npz"

    with laspy.open(source) as las_file:
        header = las_file.header
        total_points = int(header.point_count)
        point_format = header.point_format
        dimensions = set(point_format.dimension_names)
        has_rgb = {"red", "green", "blue"}.issubset(dimensions)
        has_intensity = "intensity" in dimensions
        has_returns = {"return_number", "number_of_returns"}.issubset(dimensions)
        has_classification = "classification" in dimensions
        bounds = np.array(
            [
                [float(header.mins[0]), float(header.maxs[0])],
                [float(header.mins[1]), float(header.maxs[1])],
                [float(header.mins[2]), float(header.maxs[2])],
            ],
            dtype=np.float64,
        )

        xyz = np.empty((total_points, 3), dtype=np.float64)
        intensity = np.zeros(total_points, dtype=np.uint16)
        rgb = np.full((total_points, 3), SURFACE_DEFAULT_RGB_VALUE, dtype=np.uint8)
        return_number = np.ones(total_points, dtype=np.uint8)
        number_of_returns = np.ones(total_points, dtype=np.uint8)
        classification = np.zeros(total_points, dtype=np.uint8)

        offset = 0
        for chunk in las_file.chunk_iterator(SURFACE_LAS_CHUNK_SIZE):
            chunk_size = len(chunk)
            section = slice(offset, offset + chunk_size)
            xyz[section, 0] = np.asarray(chunk.x, dtype=np.float64)
            xyz[section, 1] = np.asarray(chunk.y, dtype=np.float64)
            xyz[section, 2] = np.asarray(chunk.z, dtype=np.float64)
            if has_intensity:
                intensity[section] = np.asarray(chunk.intensity, dtype=np.uint16)
            if has_rgb:
                colors = np.stack(
                    [
                        np.asarray(chunk.red, dtype=np.float64),
                        np.asarray(chunk.green, dtype=np.float64),
                        np.asarray(chunk.blue, dtype=np.float64),
                    ],
                    axis=1,
                )
                if colors.size and float(colors.max()) > 255.0:
                    colors /= 256.0
                rgb[section] = colors.clip(0, 255).astype(np.uint8)
            if has_returns:
                return_number[section] = np.asarray(chunk.return_number, dtype=np.uint8)
                number_of_returns[section] = np.asarray(chunk.number_of_returns, dtype=np.uint8)
            if has_classification:
                classification[section] = np.asarray(chunk.classification, dtype=np.uint8)
            offset += chunk_size
            if progress_callback:
                progress_callback(int(offset / total_points * 100) if total_points else 100)

    temporary_path: Path | None = None
    try:
        with tempfile.NamedTemporaryFile(
            mode="wb",
            dir=target_dir,
            prefix=".structured.",
            suffix=".npz.tmp",
            delete=False,
        ) as temporary:
            temporary_path = Path(temporary.name)
            np.savez_compressed(
                temporary,
                xyz=xyz,
                intensity=intensity,
                rgb=rgb,
                return_number=return_number,
                number_of_returns=number_of_returns,
                classification=classification,
                bounds=bounds,
                total_points=np.array([total_points], dtype=np.int64),
                has_rgb=np.array([int(has_rgb)], dtype=np.int8),
            )
            temporary.flush()
            os.fsync(temporary.fileno())
        # 이 교체가 윈도우 공유 위반으로 죽으면 WF1 분석 전체가 실패한다 — 재시도한다.
        replace_with_retry(temporary_path, target)
        temporary_path = None
    finally:
        if temporary_path is not None:
            temporary_path.unlink(missing_ok=True)

    if progress_callback and total_points == 0:
        progress_callback(100)
    return target