Miroslav Simko <ms@iolabs.ch> 2026-09-02T07:50:17+02:00
Commit #115 ยท 66 snippets
.../cluster_io.py | 22 ++-- src/iolabs_point_cloud_mask_clustering/input_io.py | 55 ++++---- src/iolabs_point_cloud_mask_clustering/pipeline.py | 140 ++++++++++----------- tests/test_cluster_io.py | 59 +++++---- 4 files changed, 133 insertions(+), 143 deletions(-)
| 43 | np.rint(np.asarray(values, dtype=np.float64)), info.min, info.max | 41 | np.rint(np.asarray(values, dtype=np.float64)), info.min, info.max |
| 44 | ).astype(dtype) | 42 | ).astype(dtype) |
| 45 | 43 | ||
| 46 | 44 | ||
| 47 | def _extra_channel_arrays(data: ColorIntensityData) -> dict[str, np.ndarray]: | 45 | def _extra_channel_arrays(data: color_intensity_data.ColorIntensityData) -> dict[str, np.ndarray]: |
| 48 | """Return the channels to forward verbatim, keyed by NPZ member name. | 46 | """Return the channels to forward verbatim, keyed by NPZ member name. |
| 49 | 47 | ||
| 50 | The installed dataclass is inspected with :func:`dataclasses.fields` rather | 48 | The installed dataclass is inspected with :func:`dataclasses.fields` rather |
| 51 | than matched against a version, so this works on any ``iolabs-common``: | 49 | than matched against a version, so this works on any ``iolabs-common``: |
| 1 | """Run the mask-clustering step end to end for one segment.""" | 1 | """Run the mask-clustering step end to end for one segment.""" |
| 2 | 2 | ||
| 3 | import datetime | ||
| 3 | import os | 4 | import os |
| 5 | import pathlib | ||
| 4 | import tempfile | 6 | import tempfile |
| 5 | from datetime import UTC, datetime | 7 | from importlib import metadata |
| 6 | from importlib.metadata import PackageNotFoundError, version | ||
| 7 | from pathlib import Path | ||
| 8 | from typing import Any | 8 | from typing import Any |
| 9 | 9 | ||
| 10 | import numpy as np | 10 | import numpy as np |
| 11 | from iolabs.common.color_intensity_data import ColorIntensityData | 11 | from iolabs import logstash |
| 12 | from iolabs.common.run_stats import read_stats, write_stats | 12 | from iolabs.common import color_intensity_data, run_stats, segments, version_info |
| 13 | from iolabs.common.segments import segment_record_files | ||
| 14 | from iolabs.common.version_info import save_version_json | ||
| 15 | from iolabs.logstash import get_props_logger | ||
| 16 | 13 | ||
| 17 | from ._config import MaskClusteringConfig | 14 | from . import ( |
| 18 | from ._log_props import LOG_PROPS | 15 | _config, |
| 19 | from .cluster_io import write_cluster_npz | 16 | _log_props, |
| 20 | from .geometry import component_geojson_geometry, mask_polygons_xy, write_prism_ply | 17 | cluster_io, |
| 21 | from .input_io import channel_dtypes_of, load_separation_arrays, load_step3_file | 18 | geometry, |
| 22 | from .intensity_separation import ( | 19 | input_io, |
| 23 | ClusterSeparation, | 20 | intensity_separation, |
| 24 | dilate_mask, | 21 | mask_components, |
| 25 | separate_cluster, | 22 | point_projection, |
| 26 | write_separation_pdf, | 23 | raster_frame, |
| 24 | types, | ||
| 27 | ) | 25 | ) |
| 28 | from .mask_components import label_components, load_classified_mask | ||
| 29 | from .point_projection import assign_point_chunks | ||
| 30 | from .raster_frame import reconstruct_raster_frame | ||
| 31 | from .types import ComponentResult, PointChannels, SegmentResult, SegmentType | ||
| 32 | 26 | ||
| 33 | logger = get_props_logger(__name__, LOG_PROPS) | 27 | logger = logstash.get_props_logger(__name__, _log_props.LOG_PROPS) |
| 34 | 28 | ||
| 35 | 29 | ||
| 36 | def apply_sparse_policy( | 30 | def apply_sparse_policy( |
| 37 | *, | 31 | *, |
| 38 | segment_name: str, | 32 | segment_name: str, |
| 39 | component_id: int, | 33 | component_id: int, |
| 40 | segment_type: SegmentType, | 34 | segment_type: types.SegmentType, |
| 41 | point_count: int, | 35 | point_count: int, |
| 42 | warn_below_points: int, | 36 | warn_below_points: int, |
| 43 | min_points_per_cluster: int, | 37 | min_points_per_cluster: int, |
| 44 | logger: Any, | 38 | logger: Any, |
| 387 | 381 | ||
| 388 | separation_pdf_relative: str | None = None | 382 | separation_pdf_relative: str | None = None |
| 389 | if separation_enabled: | 383 | if separation_enabled: |
| 390 | separation_pdf_path = output_dir / str(separation_cfg["pdf_filename"]) | 384 | separation_pdf_path = output_dir / str(separation_cfg["pdf_filename"]) |
| 391 | write_separation_pdf(separations, separation_pdf_path, separation_cfg) | 385 | intensity_separation.write_separation_pdf(separations, separation_pdf_path, separation_cfg) |
| 392 | separation_pdf_relative = separation_pdf_path.relative_to(output_dir).as_posix() | 386 | separation_pdf_relative = separation_pdf_path.relative_to(output_dir).as_posix() |
| 393 | 387 | ||
| 394 | versions_path = output_dir / "run6c_versions.json" | 388 | versions_path = output_dir / "run6c_versions.json" |
| 395 | save_version_json(versions_path, "step6c_mask_clustering") | 389 | version_info.save_version_json(versions_path, "step6c_mask_clustering") |
| 396 | dropped_count = len(components) - retained_index | 390 | dropped_count = len(components) - retained_index |
| 397 | package_version = _package_version() | 391 | package_version = _package_version() |
| 398 | completed_at = datetime.now(UTC) | 392 | completed_at = datetime.datetime.now(datetime.UTC) |
| 399 | manifest = { | 393 | manifest = { |
| 400 | "schema": "iolabs.mask_clustering/1", | 394 | "schema": "iolabs.mask_clustering/1", |
| 401 | "package_version": package_version, | 395 | "package_version": package_version, |
| 402 | "timestamps": { | 396 | "timestamps": { |
| 497 | 491 | ||
| 498 | 492 | ||
| 499 | def _package_version() -> str: | 493 | def _package_version() -> str: |
| 500 | try: | 494 | try: |
| 501 | return version("iolabs-point-cloud-mask-clustering") | 495 | return metadata.version("iolabs-point-cloud-mask-clustering") |
| 502 | except PackageNotFoundError: | 496 | except metadata.PackageNotFoundError: |
| 503 | return "0.1.0" | 497 | return "0.1.0" |
| 10 | archive is byte-identical to the pre-AI3D-382 output. | 10 | archive is byte-identical to the pre-AI3D-382 output. |
| 11 | """ | 11 | """ |
| 12 | 12 | ||
| 13 | import dataclasses | 13 | import dataclasses |
| 14 | from pathlib import Path | 14 | import pathlib |
| 15 | 15 | ||
| 16 | import numpy as np | 16 | import numpy as np |
| 17 | from iolabs.common.atomic_io import atomic_savez | 17 | from iolabs import logstash |
| 18 | from iolabs.common.color_intensity_data import ColorIntensityData | 18 | from iolabs.common import atomic_io, color_intensity_data |
| 19 | from iolabs.logstash import get_props_logger | ||
| 20 | 19 | ||
| 21 | from ._log_props import LOG_PROPS | 20 | from . import _log_props, types |
| 22 | from .types import PointChannels, SegmentType | ||
| 23 | 21 | ||
| 24 | logger = get_props_logger(__name__, LOG_PROPS) | 22 | logger = logstash.get_props_logger(__name__, _log_props.LOG_PROPS) |
| 25 | 23 | ||
| 26 | #: ``ColorIntensityData`` fields :func:`write_cluster_npz` writes itself, with | 24 | #: ``ColorIntensityData`` fields :func:`write_cluster_npz` writes itself, with |
| 27 | #: their own dtype handling; they must not be forwarded a second time. | 25 | #: their own dtype handling; they must not be forwarded a second time. |
| 28 | EXPLICIT_CHANNEL_FIELDS: frozenset[str] = frozenset( | 26 | EXPLICIT_CHANNEL_FIELDS: frozenset[str] = frozenset( |
| 80 | return extras | 78 | return extras |
| 81 | 79 | ||
| 82 | 80 | ||
| 83 | def write_cluster_npz( | 81 | def write_cluster_npz( |
| 84 | output_path: Path, | 82 | output_path: pathlib.Path, |
| 85 | channels: PointChannels, | 83 | channels: types.PointChannels, |
| 86 | segment_type: SegmentType, | 84 | segment_type: types.SegmentType, |
| 87 | ) -> None: | 85 | ) -> None: |
| 88 | """Write one cluster to *output_path* as an uncompressed, atomically-replaced NPZ. | 86 | """Write one cluster to *output_path* as an uncompressed, atomically-replaced NPZ. |
| 89 | 87 | ||
| 90 | Members are STORED, not deflated: tablecloth's reader chunk-streams cluster | 88 | Members are STORED, not deflated: tablecloth's reader chunk-streams cluster |
| 97 | channel beyond the five explicit ones is forwarded under its field | 95 | channel beyond the five explicit ones is forwarded under its field |
| 98 | name, with its own dtype. | 96 | name, with its own dtype. |
| 99 | segment_type: Road-marking class recorded alongside the points. | 97 | segment_type: Road-marking class recorded alongside the points. |
| 100 | """ | 98 | """ |
| 101 | atomic_savez( | 99 | atomic_io.atomic_savez( |
| 102 | output_path, | 100 | output_path, |
| 103 | compress=False, | 101 | compress=False, |
| 104 | points=np.asarray(channels.points, dtype=np.float64), | 102 | points=np.asarray(channels.points, dtype=np.float64), |
| 105 | scan_angle=_round_clip(channels.data.scan_angle_rank, np.dtype(np.int8)), | 103 | scan_angle=_round_clip(channels.data.scan_angle_rank, np.dtype(np.int8)), |
| 6 | ``ColorIntensityData`` flows through this module with no code change here. | 6 | ``ColorIntensityData`` flows through this module with no code change here. |
| 7 | """ | 7 | """ |
| 8 | 8 | ||
| 9 | import dataclasses | 9 | import dataclasses |
| 10 | from collections.abc import Iterable, Mapping | 10 | import pathlib |
| 11 | from dataclasses import dataclass | 11 | from collections import abc |
| 12 | from pathlib import Path | ||
| 13 | 12 | ||
| 14 | import numpy as np | 13 | import numpy as np |
| 15 | from iolabs.common import segment_points_io | 14 | from iolabs import logstash |
| 16 | from iolabs.common.color_intensity_data import ColorIntensityData | 15 | from iolabs.common import color_intensity_data, segment_points_io |
| 17 | from iolabs.logstash import get_props_logger | ||
| 18 | 16 | ||
| 19 | from ._log_props import LOG_PROPS | 17 | from . import _log_props, types |
| 20 | from .types import PointChannels, RasterFrame | ||
| 21 | 18 | ||
| 22 | logger = get_props_logger(__name__, LOG_PROPS) | 19 | logger = logstash.get_props_logger(__name__, _log_props.LOG_PROPS) |
| 23 | 20 | ||
| 24 | #: Deprecated alias for the shared run3 point-record schema; use | 21 | #: Deprecated alias for the shared run3 point-record schema; use |
| 25 | #: :data:`iolabs.common.segment_points_io.POINT_RECORD_KEYS` instead. | 22 | #: :data:`iolabs.common.segment_points_io.POINT_RECORD_KEYS` instead. |
| 26 | REQUIRED_ARRAYS: tuple[str, ...] = segment_points_io.POINT_RECORD_KEYS | 23 | REQUIRED_ARRAYS: tuple[str, ...] = segment_points_io.POINT_RECORD_KEYS |
| 27 | 24 | ||
| 28 | #: Record keys whose name differs from the matching ``ColorIntensityData`` field. | 25 | #: Record keys whose name differs from the matching ``ColorIntensityData`` field. |
| 29 | _RECORD_KEY_TO_CHANNEL_FIELD: Mapping[str, str] = {"scan_angle": "scan_angle_rank"} | 26 | _RECORD_KEY_TO_CHANNEL_FIELD: abc.Mapping[str, str] = {"scan_angle": "scan_angle_rank"} |
| 30 | 27 | ||
| 31 | #: Keys the lean separation path reads directly from the npz. | 28 | #: Keys the lean separation path reads directly from the npz. |
| 32 | _LEAN_KEYS: tuple[str, ...] = ("points", "intensity") | 29 | _LEAN_KEYS: tuple[str, ...] = ("points", "intensity") |
| 33 | 30 |
| 43 | Derived from the installed ``iolabs.common`` contract: a key added to both | 40 | Derived from the installed ``iolabs.common`` contract: a key added to both |
| 44 | the point-record schema and ``ColorIntensityData`` appears here | 41 | the point-record schema and ``ColorIntensityData`` appears here |
| 45 | automatically, and keys the installed dataclass cannot hold are skipped. | 42 | automatically, and keys the installed dataclass cannot hold are skipped. |
| 46 | """ | 43 | """ |
| 47 | fields = {field.name for field in dataclasses.fields(ColorIntensityData)} | 44 | fields = {field.name for field in dataclasses.fields(color_intensity_data.ColorIntensityData)} |
| 48 | return tuple( | 45 | return tuple( |
| 49 | key | 46 | key |
| 50 | for key in segment_points_io.POINT_RECORD_KEYS | 47 | for key in segment_points_io.POINT_RECORD_KEYS |
| 51 | if key != "points" and _channel_field(key) in fields | 48 | if key != "points" and _channel_field(key) in fields |
| 52 | ) | 49 | ) |
| 53 | 50 | ||
| 54 | 51 | ||
| 55 | def channel_dtypes_of(data: ColorIntensityData) -> dict[str, np.dtype]: | 52 | def channel_dtypes_of(data: color_intensity_data.ColorIntensityData) -> dict[str, np.dtype]: |
| 56 | """Return *data*'s per-channel storage dtypes, keyed by run3 record key.""" | 53 | """Return *data*'s per-channel storage dtypes, keyed by run3 record key.""" |
| 57 | return { | 54 | return { |
| 58 | key: getattr(data, _channel_field(key)).dtype for key in channel_record_keys() | 55 | key: getattr(data, _channel_field(key)).dtype for key in channel_record_keys() |
| 59 | } | 56 | } |
| 60 | 57 | ||
| 61 | 58 | ||
| 62 | def _channels_from_record(record: Mapping[str, np.ndarray]) -> ColorIntensityData: | 59 | def _channels_from_record( |
| 60 | record: abc.Mapping[str, np.ndarray], | ||
| 61 | ) -> color_intensity_data.ColorIntensityData: | ||
| 63 | """Build channels from a run3 record without naming each field.""" | 62 | """Build channels from a run3 record without naming each field.""" |
| 64 | return ColorIntensityData( | 63 | return color_intensity_data.ColorIntensityData( |
| 65 | **{ | 64 | **{ |
| 66 | _channel_field(key): record[key] | 65 | _channel_field(key): record[key] |
| 67 | for key in channel_record_keys() | 66 | for key in channel_record_keys() |
| 68 | if key in record | 67 | if key in record |
| 69 | } | 68 | } |
| 70 | ) | 69 | ) |
| 71 | 70 | ||
| 72 | 71 | ||
| 73 | def _as_points(path: Path, raw: np.ndarray) -> np.ndarray: | 72 | def _as_points(path: pathlib.Path, raw: np.ndarray) -> np.ndarray: |
| 74 | points = np.asarray(raw, dtype=np.float64) | 73 | points = np.asarray(raw, dtype=np.float64) |
| 75 | if points.ndim != 2 or points.shape[1] != 3: | 74 | if points.ndim != 2 or points.shape[1] != 3: |
| 76 | raise ValueError(f"{path}: points must have shape (N, 3), got {points.shape}") | 75 | raise ValueError(f"{path}: points must have shape (N, 3), got {points.shape}") |
| 77 | return points | 76 | return points |
| 78 | 77 | ||
| 79 | 78 | ||
| 80 | def _as_channel(path: Path, name: str, raw: np.ndarray, count: int) -> np.ndarray: | 79 | def _as_channel(path: pathlib.Path, name: str, raw: np.ndarray, count: int) -> np.ndarray: |
| 81 | array = np.asarray(raw) | 80 | array = np.asarray(raw) |
| 82 | if array.ndim != 1 or len(array) != count: | 81 | if array.ndim != 1 or len(array) != count: |
| 83 | raise ValueError( | 82 | raise ValueError( |
| 84 | f"{path}: {name} must have shape ({count},), got {array.shape}" | 83 | f"{path}: {name} must have shape ({count},), got {array.shape}" |
| 85 | ) | 84 | ) |
| 86 | return array | 85 | return array |
| 87 | 86 | ||
| 88 | 87 | ||
| 89 | @dataclass(frozen=True) | 88 | @dataclasses.dataclass(frozen=True) |
| 90 | class SeparationArrays: | 89 | class SeparationArrays: |
| 91 | """Preallocated full-segment arrays for intensity separation.""" | 90 | """Preallocated full-segment arrays for intensity separation.""" |
| 92 | 91 | ||
| 93 | xyz: np.ndarray | 92 | xyz: np.ndarray |
| 94 | rows: np.ndarray | 93 | rows: np.ndarray |
| 95 | cols: np.ndarray | 94 | cols: np.ndarray |
| 96 | intensity: np.ndarray | 95 | intensity: np.ndarray |
| 97 | channels: ColorIntensityData | None | 96 | channels: color_intensity_data.ColorIntensityData | None |
| 98 | 97 | ||
| 99 | 98 | ||
| 100 | def load_step3_file(path: Path) -> PointChannels: | 99 | def load_step3_file(path: pathlib.Path) -> types.PointChannels: |
| 101 | """Load one Step 3 chunk with all its channels. | 100 | """Load one Step 3 chunk with all its channels. |
| 102 | 101 | ||
| 103 | The record is read and schema-validated by | 102 | The record is read and schema-validated by |
| 104 | :func:`iolabs.common.segment_points_io.load_points_npz`; XYZ is then cast to | 103 | :func:`iolabs.common.segment_points_io.load_points_npz`; XYZ is then cast to |
| 113 | Raises: | 112 | Raises: |
| 114 | ValueError: A required array is missing or has the wrong shape. | 113 | ValueError: A required array is missing or has the wrong shape. |
| 115 | """ | 114 | """ |
| 116 | record = segment_points_io.load_points_npz(path) | 115 | record = segment_points_io.load_points_npz(path) |
| 117 | return PointChannels( | 116 | return types.PointChannels( |
| 118 | points=_as_points(path, record["points"]), | 117 | points=_as_points(path, record["points"]), |
| 119 | data=_channels_from_record(record), | 118 | data=_channels_from_record(record), |
| 120 | ) | 119 | ) |
| 121 | 120 | ||
| 122 | 121 | ||
| 123 | def _load_separation_chunk( | 122 | def _load_separation_chunk( |
| 124 | path: Path, *, with_channels: bool | 123 | path: pathlib.Path, *, with_channels: bool |
| 125 | ) -> tuple[np.ndarray, np.ndarray, ColorIntensityData | None]: | 124 | ) -> tuple[np.ndarray, np.ndarray, color_intensity_data.ColorIntensityData | None]: |
| 126 | """Load one Step 3 chunk for separation. | 125 | """Load one Step 3 chunk for separation. |
| 127 | 126 | ||
| 128 | When ``with_channels`` is false, only ``points`` and ``intensity`` are read | 127 | When ``with_channels`` is false, only ``points`` and ``intensity`` are read |
| 129 | from the npz; the other record keys are neither materialized nor validated | 128 | from the npz; the other record keys are neither materialized nor validated |
| 146 | intensity = _as_channel(path, "intensity", payload["intensity"], len(points)) | 145 | intensity = _as_channel(path, "intensity", payload["intensity"], len(points)) |
| 147 | return points, intensity, None | 146 | return points, intensity, None |
| 148 | 147 | ||
| 149 | 148 | ||
| 150 | def load_step3_points(paths: Iterable[Path]) -> PointChannels: | 149 | def load_step3_points(paths: abc.Iterable[pathlib.Path]) -> types.PointChannels: |
| 151 | """Load and concatenate several Step 3 chunks. | 150 | """Load and concatenate several Step 3 chunks. |
| 152 | 151 | ||
| 153 | Args: | 152 | Args: |
| 154 | paths: Step 3 ``.npz`` chunks; loaded in sorted order. | 153 | paths: Step 3 ``.npz`` chunks; loaded in sorted order. |
| 165 | points = np.concatenate([chunk.points for chunk in chunks]) | 164 | points = np.concatenate([chunk.points for chunk in chunks]) |
| 166 | data = chunks[0].data | 165 | data = chunks[0].data |
| 167 | for chunk in chunks[1:]: | 166 | for chunk in chunks[1:]: |
| 168 | data = data.append(chunk.data) | 167 | data = data.append(chunk.data) |
| 169 | return PointChannels(points=points, data=data) | 168 | return types.PointChannels(points=points, data=data) |
| 170 | 169 | ||
| 171 | 170 | ||
| 172 | def load_separation_arrays( | 171 | def load_separation_arrays( |
| 173 | paths: Iterable[Path], | 172 | paths: abc.Iterable[pathlib.Path], |
| 174 | frame: RasterFrame, | 173 | frame: types.RasterFrame, |
| 175 | total_count: int, | 174 | total_count: int, |
| 176 | *, | 175 | *, |
| 177 | channel_dtypes: Mapping[str, np.dtype], | 176 | channel_dtypes: abc.Mapping[str, np.dtype], |
| 178 | with_channels: bool, | 177 | with_channels: bool, |
| 179 | ) -> SeparationArrays: | 178 | ) -> SeparationArrays: |
| 180 | """Stream Step 3 chunks into preallocated separation arrays. | 179 | """Stream Step 3 chunks into preallocated separation arrays. |
| 181 | 180 |
| 250 | ) | 249 | ) |
| 251 | 250 | ||
| 252 | channels = None | 251 | channels = None |
| 253 | if with_channels: | 252 | if with_channels: |
| 254 | channels = ColorIntensityData( | 253 | channels = color_intensity_data.ColorIntensityData( |
| 255 | intensity=intensity, | 254 | intensity=intensity, |
| 256 | **{_channel_field(key): buffers[key] for key in extra_keys}, | 255 | **{_channel_field(key): buffers[key] for key in extra_keys}, |
| 257 | ) | 256 | ) |
| 258 | return SeparationArrays( | 257 | return SeparationArrays( |
| 76 | 70 | ||
| 77 | 71 | ||
| 78 | def process_segment( | 72 | def process_segment( |
| 79 | *, | 73 | *, |
| 80 | segment_dir: Path, | 74 | segment_dir: pathlib.Path, |
| 81 | mask_or_vectors_path: Path, | 75 | mask_or_vectors_path: pathlib.Path, |
| 82 | raster_metadata_path: Path, | 76 | raster_metadata_path: pathlib.Path, |
| 83 | output_dir: Path, | 77 | output_dir: pathlib.Path, |
| 84 | config: MaskClusteringConfig | dict[str, Any], | 78 | config: _config.MaskClusteringConfig | dict[str, Any], |
| 85 | overwrite: bool = False, | 79 | overwrite: bool = False, |
| 86 | ) -> SegmentResult: | 80 | ) -> types.SegmentResult: |
| 87 | """Cluster one segment's classified mask into Step 6-compatible artifacts. | 81 | """Cluster one segment's classified mask into Step 6-compatible artifacts. |
| 88 | 82 | ||
| 89 | Args: | 83 | Args: |
| 90 | segment_dir: Segment directory holding the Step 3 point NPZ chunks. | 84 | segment_dir: Segment directory holding the Step 3 point NPZ chunks. |
| 105 | the configured suffix, or *raster_metadata_path* does not exist. | 99 | the configured suffix, or *raster_metadata_path* does not exist. |
| 106 | ValueError: The segment holds no points, or the mask, sidecar and | 100 | ValueError: The segment holds no points, or the mask, sidecar and |
| 107 | reconstructed frame disagree. | 101 | reconstructed frame disagree. |
| 108 | """ | 102 | """ |
| 109 | started_at = datetime.now(UTC) | 103 | started_at = datetime.datetime.now(datetime.UTC) |
| 110 | cfg = MaskClusteringConfig.coerce(config) | 104 | cfg = _config.MaskClusteringConfig.coerce(config) |
| 111 | manifest_path = output_dir / cfg.output.manifest_filename | 105 | manifest_path = output_dir / cfg.output.manifest_filename |
| 112 | if manifest_path.exists(): | 106 | if manifest_path.exists(): |
| 113 | if not overwrite: | 107 | if not overwrite: |
| 114 | raise FileExistsError( | 108 | raise FileExistsError( |
| 115 | f"Completed mask-clustering manifest already exists: {manifest_path}" | 109 | f"Completed mask-clustering manifest already exists: {manifest_path}" |
| 116 | ) | 110 | ) |
| 117 | _remove_prior_artifacts(output_dir, manifest_path) | 111 | _remove_prior_artifacts(output_dir, manifest_path) |
| 118 | 112 | ||
| 119 | step3_paths = segment_record_files( | 113 | step3_paths = segments.segment_record_files( |
| 120 | segment_dir, cfg.file_naming.segment_points_suffix, required=True | 114 | segment_dir, cfg.file_naming.segment_points_suffix, required=True |
| 121 | ) | 115 | ) |
| 122 | metadata = read_stats(raster_metadata_path, require_dict=True) | 116 | segment_metadata = run_stats.read_stats(raster_metadata_path, require_dict=True) |
| 123 | segment_name = str( | 117 | segment_name = str( |
| 124 | metadata.get("segment_name", metadata.get("segment_id", segment_dir.name)) | 118 | segment_metadata.get("segment_name", segment_metadata.get("segment_id", segment_dir.name)) |
| 125 | ) | 119 | ) |
| 126 | 120 | ||
| 127 | source_count = 0 | 121 | source_count = 0 |
| 128 | xy_min = np.array([np.inf, np.inf]) | 122 | xy_min = np.array([np.inf, np.inf]) |
| 130 | z_min = np.inf | 124 | z_min = np.inf |
| 131 | z_max = -np.inf | 125 | z_max = -np.inf |
| 132 | channel_dtypes: dict[str, np.dtype] | None = None | 126 | channel_dtypes: dict[str, np.dtype] | None = None |
| 133 | for path in step3_paths: | 127 | for path in step3_paths: |
| 134 | chunk = load_step3_file(path) | 128 | chunk = input_io.load_step3_file(path) |
| 135 | source_count += len(chunk.points) | 129 | source_count += len(chunk.points) |
| 136 | chunk_dtypes = channel_dtypes_of(chunk.data) | 130 | chunk_dtypes = input_io.channel_dtypes_of(chunk.data) |
| 137 | if channel_dtypes is None: | 131 | if channel_dtypes is None: |
| 138 | channel_dtypes = chunk_dtypes | 132 | channel_dtypes = chunk_dtypes |
| 139 | else: | 133 | else: |
| 140 | channel_dtypes = { | 134 | channel_dtypes = { |
| 153 | if source_count == 0: | 147 | if source_count == 0: |
| 154 | raise ValueError(f"Step 3 source in {segment_dir} contains no points") | 148 | raise ValueError(f"Step 3 source in {segment_dir} contains no points") |
| 155 | assert channel_dtypes is not None | 149 | assert channel_dtypes is not None |
| 156 | 150 | ||
| 157 | mask_shape = (int(metadata["pixel_y"]), int(metadata["pixel_x"])) | 151 | mask_shape = (int(segment_metadata["pixel_y"]), int(segment_metadata["pixel_x"])) |
| 158 | mask_cfg = cfg.mask | 152 | mask_cfg = cfg.mask |
| 159 | extrema = np.array( | 153 | extrema = np.array( |
| 160 | [[xy_min[0], xy_min[1], z_min], [xy_max[0], xy_max[1], z_max]], | 154 | [[xy_min[0], xy_min[1], z_min], [xy_max[0], xy_max[1], z_max]], |
| 161 | dtype=np.float64, | 155 | dtype=np.float64, |
| 162 | ) | 156 | ) |
| 163 | frame = reconstruct_raster_frame( | 157 | frame = raster_frame.reconstruct_raster_frame( |
| 164 | extrema, | 158 | extrema, |
| 165 | metadata, | 159 | segment_metadata, |
| 166 | mask_shape=mask_shape, | 160 | mask_shape=mask_shape, |
| 167 | margin_pixels=cfg.raster_frame.margin_pixels, | 161 | margin_pixels=cfg.raster_frame.margin_pixels, |
| 168 | origin_tolerance_pixels=cfg.raster_frame.metadata_origin_tolerance_pixels, | 162 | origin_tolerance_pixels=cfg.raster_frame.metadata_origin_tolerance_pixels, |
| 169 | ) | 163 | ) |
| 170 | classified_mask = load_classified_mask( | 164 | classified_mask = mask_components.load_classified_mask( |
| 171 | mask_or_vectors_path, | 165 | mask_or_vectors_path, |
| 172 | shape=mask_shape, | 166 | shape=mask_shape, |
| 173 | background_class=mask_cfg.background_class, | 167 | background_class=mask_cfg.background_class, |
| 174 | solid_class=mask_cfg.solid_class, | 168 | solid_class=mask_cfg.solid_class, |
| 175 | dashed_class=mask_cfg.dashed_class, | 169 | dashed_class=mask_cfg.dashed_class, |
| 176 | vector_stroke_px=mask_cfg.vector_stroke_px, | 170 | vector_stroke_px=mask_cfg.vector_stroke_px, |
| 177 | ) | 171 | ) |
| 178 | components, label_image = label_components( | 172 | components, label_image = mask_components.label_components( |
| 179 | classified_mask, | 173 | classified_mask, |
| 180 | solid_class=mask_cfg.solid_class, | 174 | solid_class=mask_cfg.solid_class, |
| 181 | dashed_class=mask_cfg.dashed_class, | 175 | dashed_class=mask_cfg.dashed_class, |
| 182 | connectivity=mask_cfg.connectivity, | 176 | connectivity=mask_cfg.connectivity, |
| 187 | segment_name, | 181 | segment_name, |
| 188 | extra={"segment_name": segment_name, "action": "empty_success"}, | 182 | extra={"segment_name": segment_name, "action": "empty_success"}, |
| 189 | ) | 183 | ) |
| 190 | 184 | ||
| 191 | assigned, out_of_frame = assign_point_chunks( | 185 | assigned, out_of_frame = point_projection.assign_point_chunks( |
| 192 | (load_step3_file(path) for path in step3_paths), | 186 | (input_io.load_step3_file(path) for path in step3_paths), |
| 193 | frame, | 187 | frame, |
| 194 | label_image, | 188 | label_image, |
| 195 | components, | 189 | components, |
| 196 | ) | 190 | ) |
| 209 | # point-cloud viewer. Written to a separate subdir so Step 6b never rasterises them. | 203 | # point-cloud viewer. Written to a separate subdir so Step 6b never rasterises them. |
| 210 | save_padded = separation_enabled and bool(separation_cfg.get("save_padded_clusters")) | 204 | save_padded = separation_enabled and bool(separation_cfg.get("save_padded_clusters")) |
| 211 | padding_debug_dir = output_dir / "padding_debug" | 205 | padding_debug_dir = output_dir / "padding_debug" |
| 212 | dilation_px = int(separation_cfg["dilation_px"]) | 206 | dilation_px = int(separation_cfg["dilation_px"]) |
| 213 | separations: list[ClusterSeparation] = [] | 207 | separations: list[intensity_separation.ClusterSeparation] = [] |
| 214 | # The asphalt sampling ring lies outside each component's mask, so it needs more | 208 | # The asphalt sampling ring lies outside each component's mask, so it needs more |
| 215 | # than the per-component assigned points. Stream-load and project the full segment | 209 | # than the per-component assigned points. Stream-load and project the full segment |
| 216 | # once, up front, when separation is enabled. (Cropping to the union of dilated | 210 | # once, up front, when separation is enabled. (Cropping to the union of dilated |
| 217 | # component masks would also work; full-segment streaming is simpler and the lean | 211 | # component masks would also work; full-segment streaming is simpler and the lean |
| 218 | # xyz/rows/cols/intensity buffers keep it affordable.) | 212 | # xyz/rows/cols/intensity buffers keep it affordable.) |
| 219 | sep_channels: ColorIntensityData | None = None | 213 | sep_channels: color_intensity_data.ColorIntensityData | None = None |
| 220 | if separation_enabled: | 214 | if separation_enabled: |
| 221 | sep = load_separation_arrays( | 215 | sep = input_io.load_separation_arrays( |
| 222 | step3_paths, | 216 | step3_paths, |
| 223 | frame, | 217 | frame, |
| 224 | source_count, | 218 | source_count, |
| 225 | channel_dtypes=channel_dtypes, | 219 | channel_dtypes=channel_dtypes, |
| 246 | warn_below_points=cluster_cfg.warn_below_points, | 240 | warn_below_points=cluster_cfg.warn_below_points, |
| 247 | min_points_per_cluster=cluster_cfg.min_points_per_cluster, | 241 | min_points_per_cluster=cluster_cfg.min_points_per_cluster, |
| 248 | logger=logger, | 242 | logger=logger, |
| 249 | ) | 243 | ) |
| 250 | separation: ClusterSeparation | None = None | 244 | separation: intensity_separation.ClusterSeparation | None = None |
| 251 | cluster_path: Path | None = None | 245 | cluster_path: pathlib.Path | None = None |
| 252 | padding_debug_path: Path | None = None | 246 | padding_debug_path: pathlib.Path | None = None |
| 253 | if status == "retained": | 247 | if status == "retained": |
| 254 | if separation_enabled: | 248 | if separation_enabled: |
| 255 | separation = separate_cluster( | 249 | separation = intensity_separation.separate_cluster( |
| 256 | component, | 250 | component, |
| 257 | channels.data.intensity, | 251 | channels.data.intensity, |
| 258 | points_xyz=sep_xyz, | 252 | points_xyz=sep_xyz, |
| 259 | rows=sep_rows, | 253 | rows=sep_rows, |
| 262 | cfg=separation_cfg, | 256 | cfg=separation_cfg, |
| 263 | ) | 257 | ) |
| 264 | separations.append(separation) | 258 | separations.append(separation) |
| 265 | if apply_filter and separation.cutoff is not None: | 259 | if apply_filter and separation.cutoff is not None: |
| 266 | channels = PointChannels( | 260 | channels = types.PointChannels( |
| 267 | points=channels.points[separation.paint_mask], | 261 | points=channels.points[separation.paint_mask], |
| 268 | data=channels.data.select_by_mask(separation.paint_mask), | 262 | data=channels.data.select_by_mask(separation.paint_mask), |
| 269 | ) | 263 | ) |
| 270 | cluster_path = output_dir / f"{cluster_prefix}{retained_index:03d}.npz" | 264 | cluster_path = output_dir / f"{cluster_prefix}{retained_index:03d}.npz" |
| 271 | write_cluster_npz(cluster_path, channels, component.segment_type) | 265 | cluster_io.write_cluster_npz(cluster_path, channels, component.segment_type) |
| 272 | if save_padded: | 266 | if save_padded: |
| 273 | assert sep_channels is not None | 267 | assert sep_channels is not None |
| 274 | padded_mask = dilate_mask(component.binary_mask, dilation_px) | 268 | padded_mask = intensity_separation.dilate_mask(component.binary_mask, dilation_px) |
| 275 | in_padded = padded_mask[sep_rows, sep_cols] | 269 | in_padded = padded_mask[sep_rows, sep_cols] |
| 276 | padded_channels = PointChannels( | 270 | padded_channels = types.PointChannels( |
| 277 | points=sep_xyz[in_padded], | 271 | points=sep_xyz[in_padded], |
| 278 | data=sep_channels.select_by_mask(in_padded), | 272 | data=sep_channels.select_by_mask(in_padded), |
| 279 | ) | 273 | ) |
| 280 | padding_debug_dir.mkdir(parents=True, exist_ok=True) | 274 | padding_debug_dir.mkdir(parents=True, exist_ok=True) |
| 281 | padding_debug_path = padding_debug_dir / f"padded_{retained_index:03d}.npz" | 275 | padding_debug_path = padding_debug_dir / f"padded_{retained_index:03d}.npz" |
| 282 | write_cluster_npz( | 276 | cluster_io.write_cluster_npz( |
| 283 | padding_debug_path, padded_channels, component.segment_type | 277 | padding_debug_path, padded_channels, component.segment_type |
| 284 | ) | 278 | ) |
| 285 | retained_index += 1 | 279 | retained_index += 1 |
| 286 | 280 | ||
| 287 | prism_path = geometry_dir / ( | 281 | prism_path = geometry_dir / ( |
| 288 | f"component_{component.component_id:04d}_{component.segment_type.value}_prism.ply" | 282 | f"component_{component.component_id:04d}_{component.segment_type.value}_prism.ply" |
| 289 | ) | 283 | ) |
| 290 | polygons = mask_polygons_xy( | 284 | polygons = geometry.mask_polygons_xy( |
| 291 | component.binary_mask, | 285 | component.binary_mask, |
| 292 | frame, | 286 | frame, |
| 293 | simplify_tolerance_px=geometry_cfg.simplify_tolerance_px, | 287 | simplify_tolerance_px=geometry_cfg.simplify_tolerance_px, |
| 294 | ) | 288 | ) |
| 295 | if geometry_cfg.write_ply: | 289 | if geometry_cfg.write_ply: |
| 296 | write_prism_ply( | 290 | geometry.write_prism_ply( |
| 297 | prism_path, | 291 | prism_path, |
| 298 | polygons, | 292 | polygons, |
| 299 | z_min=float(z_min), | 293 | z_min=float(z_min), |
| 300 | z_max=float(z_max), | 294 | z_max=float(z_max), |
| 341 | } | 335 | } |
| 342 | geojson_features.append( | 336 | geojson_features.append( |
| 343 | { | 337 | { |
| 344 | "type": "Feature", | 338 | "type": "Feature", |
| 345 | "geometry": component_geojson_geometry(polygons), | 339 | "geometry": geometry.component_geojson_geometry(polygons), |
| 346 | "properties": properties, | 340 | "properties": properties, |
| 347 | } | 341 | } |
| 348 | ) | 342 | ) |
| 349 | warning_state = point_count < cluster_cfg.warn_below_points | 343 | warning_state = point_count < cluster_cfg.warn_below_points |
| 361 | **separation_fields, | 355 | **separation_fields, |
| 362 | } | 356 | } |
| 363 | ) | 357 | ) |
| 364 | component_results.append( | 358 | component_results.append( |
| 365 | ComponentResult( | 359 | types.ComponentResult( |
| 366 | component_id=component.component_id, | 360 | component_id=component.component_id, |
| 367 | segment_type=component.segment_type, | 361 | segment_type=component.segment_type, |
| 368 | mask_pixel_count=component.pixel_count, | 362 | mask_pixel_count=component.pixel_count, |
| 369 | point_count=point_count, | 363 | point_count=point_count, |
| 422 | "configuration": cfg.raw, | 416 | "configuration": cfg.raw, |
| 423 | "counts": { | 417 | "counts": { |
| 424 | "foreground_components": len(components), | 418 | "foreground_components": len(components), |
| 425 | "solid_components": sum( | 419 | "solid_components": sum( |
| 426 | item.segment_type is SegmentType.SOLID for item in components | 420 | item.segment_type is types.SegmentType.SOLID for item in components |
| 427 | ), | 421 | ), |
| 428 | "dashed_components": sum( | 422 | "dashed_components": sum( |
| 429 | item.segment_type is SegmentType.DASHED for item in components | 423 | item.segment_type is types.SegmentType.DASHED for item in components |
| 430 | ), | 424 | ), |
| 431 | "retained_clusters": retained_index, | 425 | "retained_clusters": retained_index, |
| 432 | "dropped_clusters": dropped_count, | 426 | "dropped_clusters": dropped_count, |
| 433 | "out_of_frame_points": out_of_frame, | 427 | "out_of_frame_points": out_of_frame, |
| 452 | "dashed_count": manifest["counts"]["dashed_components"], | 446 | "dashed_count": manifest["counts"]["dashed_components"], |
| 453 | "warning_count": sum(item["warning"] for item in manifest_components), | 447 | "warning_count": sum(item["warning"] for item in manifest_components), |
| 454 | }, | 448 | }, |
| 455 | ) | 449 | ) |
| 456 | return SegmentResult( | 450 | return types.SegmentResult( |
| 457 | segment_name=segment_name, | 451 | segment_name=segment_name, |
| 458 | manifest_path=manifest_path, | 452 | manifest_path=manifest_path, |
| 459 | source_point_count=source_count, | 453 | source_point_count=source_count, |
| 460 | retained_count=retained_index, | 454 | retained_count=retained_index, |
| 462 | components=tuple(component_results), | 456 | components=tuple(component_results), |
| 463 | ) | 457 | ) |
| 464 | 458 | ||
| 465 | 459 | ||
| 466 | def _write_json_atomic(path: Path, value: dict[str, Any]) -> None: | 460 | def _write_json_atomic(path: pathlib.Path, value: dict[str, Any]) -> None: |
| 467 | path.parent.mkdir(parents=True, exist_ok=True) | 461 | path.parent.mkdir(parents=True, exist_ok=True) |
| 468 | descriptor, name = tempfile.mkstemp( | 462 | descriptor, name = tempfile.mkstemp( |
| 469 | prefix=f".{path.name}.", suffix=".tmp", dir=path.parent | 463 | prefix=f".{path.name}.", suffix=".tmp", dir=path.parent |
| 470 | ) | 464 | ) |
| 471 | os.close(descriptor) | 465 | os.close(descriptor) |
| 472 | try: | 466 | try: |
| 473 | write_stats(name, value, sort_keys=True, trailing_newline=True) | 467 | run_stats.write_stats(name, value, sort_keys=True, trailing_newline=True) |
| 474 | os.replace(name, path) | 468 | os.replace(name, path) |
| 475 | finally: | 469 | finally: |
| 476 | Path(name).unlink(missing_ok=True) | 470 | pathlib.Path(name).unlink(missing_ok=True) |
| 477 | 471 | ||
| 478 | 472 | ||
| 479 | def _remove_prior_artifacts(output_dir: Path, manifest_path: Path) -> None: | 473 | def _remove_prior_artifacts(output_dir: pathlib.Path, manifest_path: pathlib.Path) -> None: |
| 480 | previous = read_stats(manifest_path, require_dict=True) | 474 | previous = run_stats.read_stats(manifest_path, require_dict=True) |
| 481 | relative_paths: list[str] = [] | 475 | relative_paths: list[str] = [] |
| 482 | for component in previous.get("components", []): | 476 | for component in previous.get("components", []): |
| 483 | for key in ("cluster_path", "padding_debug_path", "prism_path"): | 477 | for key in ("cluster_path", "padding_debug_path", "prism_path"): |
| 484 | if component.get(key): | 478 | if component.get(key): |
| 1 | import dataclasses | 1 | import dataclasses |
| 2 | import logging | 2 | import logging |
| 3 | import pathlib | ||
| 3 | import zipfile | 4 | import zipfile |
| 4 | from pathlib import Path | ||
| 5 | from typing import BinaryIO | 5 | from typing import BinaryIO |
| 6 | 6 | ||
| 7 | import numpy as np | 7 | import numpy as np |
| 8 | import pytest | 8 | import pytest |
| 9 | from iolabs.common import atomic_io | 9 | from iolabs.common import atomic_io, color_intensity_data |
| 10 | from iolabs.common.color_intensity_data import ColorIntensityData | 10 | from iolabs_point_cloud_filtering_clusters import clustering_gpu_io |
| 11 | from iolabs_point_cloud_filtering_clusters.clustering_gpu_io import load_cluster_artifact_npz | ||
| 12 | 11 | ||
| 13 | from iolabs_point_cloud_mask_clustering import cluster_io | 12 | from iolabs_point_cloud_mask_clustering import cluster_io, types |
| 14 | from iolabs_point_cloud_mask_clustering.cluster_io import write_cluster_npz | ||
| 15 | from iolabs_point_cloud_mask_clustering.types import PointChannels, SegmentType | ||
| 16 | 13 | ||
| 17 | FIXED_MEMBERS = { | 14 | FIXED_MEMBERS = { |
| 18 | "points", | 15 | "points", |
| 19 | "scan_angle", | 16 | "scan_angle", |
| 43 | scan_angle: np.ndarray = dataclasses.field(default_factory=lambda: np.array([7, 8])) | 40 | scan_angle: np.ndarray = dataclasses.field(default_factory=lambda: np.array([7, 8])) |
| 44 | 41 | ||
| 45 | 42 | ||
| 46 | def _installed_channel_fields() -> set[str]: | 43 | def _installed_channel_fields() -> set[str]: |
| 47 | return {field.name for field in dataclasses.fields(ColorIntensityData)} | 44 | return {field.name for field in dataclasses.fields(color_intensity_data.ColorIntensityData)} |
| 48 | 45 | ||
| 49 | 46 | ||
| 50 | def _stub_channels(stub: type) -> PointChannels: | 47 | def _stub_channels(stub: type) -> types.PointChannels: |
| 51 | return PointChannels( | 48 | return types.PointChannels( |
| 52 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), | 49 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), |
| 53 | data=stub( | 50 | data=stub( |
| 54 | red=np.array([1.6, 2.4]), | 51 | red=np.array([1.6, 2.4]), |
| 55 | green=np.array([3.5, 4.5]), | 52 | green=np.array([3.5, 4.5]), |
| 59 | ), | 56 | ), |
| 60 | ) | 57 | ) |
| 61 | 58 | ||
| 62 | 59 | ||
| 63 | def _sample() -> PointChannels: | 60 | def _sample() -> types.PointChannels: |
| 64 | return PointChannels( | 61 | return types.PointChannels( |
| 65 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), | 62 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), |
| 66 | data=ColorIntensityData( | 63 | data=color_intensity_data.ColorIntensityData( |
| 67 | scan_angle_rank=np.array([-200.1, 127.4]), | 64 | scan_angle_rank=np.array([-200.1, 127.4]), |
| 68 | intensity=np.array([-1.0, 70000.0]), | 65 | intensity=np.array([-1.0, 70000.0]), |
| 69 | red=np.array([1.6, 2.4]), | 66 | red=np.array([1.6, 2.4]), |
| 70 | green=np.array([3.5, 4.5]), | 67 | green=np.array([3.5, 4.5]), |
| 72 | ), | 69 | ), |
| 73 | ) | 70 | ) |
| 74 | 71 | ||
| 75 | 72 | ||
| 76 | def test_writes_step6_compatible_dtypes_and_segment_type(tmp_path: Path) -> None: | 73 | def test_writes_step6_compatible_dtypes_and_segment_type(tmp_path: pathlib.Path) -> None: |
| 77 | path = tmp_path / "run6_cluster_000.npz" | 74 | path = tmp_path / "run6_cluster_000.npz" |
| 78 | source = _sample() | 75 | source = _sample() |
| 79 | write_cluster_npz(path, source, SegmentType.DASHED) | 76 | cluster_io.write_cluster_npz(path, source, types.SegmentType.DASHED) |
| 80 | with np.load(path) as payload: | 77 | with np.load(path) as payload: |
| 81 | assert payload["points"].dtype == np.float64 | 78 | assert payload["points"].dtype == np.float64 |
| 82 | assert payload["scan_angle"].dtype == np.int8 | 79 | assert payload["scan_angle"].dtype == np.int8 |
| 83 | assert payload["intensity"].dtype == np.uint16 | 80 | assert payload["intensity"].dtype == np.uint16 |
| 87 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) | 84 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) |
| 88 | np.testing.assert_array_equal(payload["intensity"], [0, 65535]) | 85 | np.testing.assert_array_equal(payload["intensity"], [0, 65535]) |
| 89 | np.testing.assert_array_equal(payload["red"], [2, 2]) | 86 | np.testing.assert_array_equal(payload["red"], [2, 2]) |
| 90 | np.testing.assert_array_equal(payload["green"], [4, 4]) | 87 | np.testing.assert_array_equal(payload["green"], [4, 4]) |
| 91 | artifact = load_cluster_artifact_npz(path) | 88 | artifact = clustering_gpu_io.load_cluster_artifact_npz(path) |
| 92 | np.testing.assert_array_equal(artifact.points, source.points) | 89 | np.testing.assert_array_equal(artifact.points, source.points) |
| 93 | 90 | ||
| 94 | 91 | ||
| 95 | def test_members_are_stored_not_deflated(tmp_path: Path) -> None: | 92 | def test_members_are_stored_not_deflated(tmp_path: pathlib.Path) -> None: |
| 96 | # Tablecloth chunk-streams cluster archives, which only works for STORED members. | 93 | # Tablecloth chunk-streams cluster archives, which only works for STORED members. |
| 97 | path = tmp_path / "run6_cluster_000.npz" | 94 | path = tmp_path / "run6_cluster_000.npz" |
| 98 | write_cluster_npz(path, _sample(), SegmentType.SOLID) | 95 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.SOLID) |
| 99 | with zipfile.ZipFile(path) as archive: | 96 | with zipfile.ZipFile(path) as archive: |
| 100 | methods = {info.filename: info.compress_type for info in archive.infolist()} | 97 | methods = {info.filename: info.compress_type for info in archive.infolist()} |
| 101 | assert methods | 98 | assert methods |
| 102 | assert set(methods.values()) == {zipfile.ZIP_STORED} | 99 | assert set(methods.values()) == {zipfile.ZIP_STORED} |
| 103 | 100 | ||
| 104 | 101 | ||
| 105 | def test_five_field_channels_write_only_the_fixed_members(tmp_path: Path) -> None: | 102 | def test_five_field_channels_write_only_the_fixed_members(tmp_path: pathlib.Path) -> None: |
| 106 | # A ColorIntensityData declaring only the explicit five (iolabs-common 0.7.0) | 103 | # A ColorIntensityData declaring only the explicit five (iolabs-common 0.7.0) |
| 107 | # must produce exactly the pre-AI3D-382 member set: no generic extras. | 104 | # must produce exactly the pre-AI3D-382 member set: no generic extras. |
| 108 | path = tmp_path / "run6_cluster_000.npz" | 105 | path = tmp_path / "run6_cluster_000.npz" |
| 109 | write_cluster_npz(path, _stub_channels(_FiveChannels), SegmentType.SOLID) | 106 | cluster_io.write_cluster_npz(path, _stub_channels(_FiveChannels), types.SegmentType.SOLID) |
| 110 | with np.load(path) as payload: | 107 | with np.load(path) as payload: |
| 111 | assert set(payload.files) == FIXED_MEMBERS | 108 | assert set(payload.files) == FIXED_MEMBERS |
| 112 | assert payload["scan_angle"].dtype == np.int8 | 109 | assert payload["scan_angle"].dtype == np.int8 |
| 113 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) | 110 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) |
| 114 | 111 | ||
| 115 | 112 | ||
| 116 | def test_member_set_follows_the_installed_channel_schema(tmp_path: Path) -> None: | 113 | def test_member_set_follows_the_installed_channel_schema(tmp_path: pathlib.Path) -> None: |
| 117 | # Holds on any iolabs-common: the extras are exactly the installed | 114 | # Holds on any iolabs-common: the extras are exactly the installed |
| 118 | # dataclass's non-explicit init fields. | 115 | # dataclass's non-explicit init fields. |
| 119 | path = tmp_path / "run6_cluster_000.npz" | 116 | path = tmp_path / "run6_cluster_000.npz" |
| 120 | write_cluster_npz(path, _sample(), SegmentType.SOLID) | 117 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.SOLID) |
| 121 | expected_extras = _installed_channel_fields() - cluster_io.EXPLICIT_CHANNEL_FIELDS | 118 | expected_extras = _installed_channel_fields() - cluster_io.EXPLICIT_CHANNEL_FIELDS |
| 122 | with np.load(path) as payload: | 119 | with np.load(path) as payload: |
| 123 | assert set(payload.files) == FIXED_MEMBERS | expected_extras | 120 | assert set(payload.files) == FIXED_MEMBERS | expected_extras |
| 124 | 121 |
| 126 | @pytest.mark.skipif( | 123 | @pytest.mark.skipif( |
| 127 | "number_of_returns" not in _installed_channel_fields(), | 124 | "number_of_returns" not in _installed_channel_fields(), |
| 128 | reason="installed ColorIntensityData predates number_of_returns", | 125 | reason="installed ColorIntensityData predates number_of_returns", |
| 129 | ) | 126 | ) |
| 130 | def test_number_of_returns_is_forwarded_aligned_with_points(tmp_path: Path) -> None: | 127 | def test_number_of_returns_is_forwarded_aligned_with_points(tmp_path: pathlib.Path) -> None: |
| 131 | source = _sample() | 128 | source = _sample() |
| 132 | source.data.number_of_returns = np.array([1, 3], dtype=np.uint8) | 129 | source.data.number_of_returns = np.array([1, 3], dtype=np.uint8) |
| 133 | path = tmp_path / "run6_cluster_000.npz" | 130 | path = tmp_path / "run6_cluster_000.npz" |
| 134 | write_cluster_npz(path, source, SegmentType.DASHED) | 131 | cluster_io.write_cluster_npz(path, source, types.SegmentType.DASHED) |
| 135 | with np.load(path) as payload: | 132 | with np.load(path) as payload: |
| 136 | returns = payload["number_of_returns"] | 133 | returns = payload["number_of_returns"] |
| 137 | assert returns.dtype == np.uint8 | 134 | assert returns.dtype == np.uint8 |
| 138 | assert len(returns) == len(source.points) | 135 | assert len(returns) == len(source.points) |
| 142 | @pytest.mark.skipif( | 139 | @pytest.mark.skipif( |
| 143 | "number_of_returns" not in _installed_channel_fields(), | 140 | "number_of_returns" not in _installed_channel_fields(), |
| 144 | reason="installed ColorIntensityData predates number_of_returns", | 141 | reason="installed ColorIntensityData predates number_of_returns", |
| 145 | ) | 142 | ) |
| 146 | def test_omitted_number_of_returns_is_forwarded_as_unknown_zeros(tmp_path: Path) -> None: | 143 | def test_omitted_number_of_returns_is_forwarded_as_unknown_zeros(tmp_path: pathlib.Path) -> None: |
| 147 | path = tmp_path / "run6_cluster_000.npz" | 144 | path = tmp_path / "run6_cluster_000.npz" |
| 148 | write_cluster_npz(path, _sample(), SegmentType.DASHED) | 145 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.DASHED) |
| 149 | with np.load(path) as payload: | 146 | with np.load(path) as payload: |
| 150 | returns = payload["number_of_returns"] | 147 | returns = payload["number_of_returns"] |
| 151 | assert returns.dtype == np.uint8 | 148 | assert returns.dtype == np.uint8 |
| 152 | np.testing.assert_array_equal(returns, [0, 0]) | 149 | np.testing.assert_array_equal(returns, [0, 0]) |
| 153 | 150 | ||
| 154 | 151 | ||
| 155 | def test_field_colliding_with_a_fixed_member_is_skipped_with_a_warning( | 152 | def test_field_colliding_with_a_fixed_member_is_skipped_with_a_warning( |
| 156 | tmp_path: Path, caplog: pytest.LogCaptureFixture | 153 | tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture |
| 157 | ) -> None: | 154 | ) -> None: |
| 158 | path = tmp_path / "run6_cluster_000.npz" | 155 | path = tmp_path / "run6_cluster_000.npz" |
| 159 | with caplog.at_level(logging.WARNING, logger=cluster_io.__name__): | 156 | with caplog.at_level(logging.WARNING, logger=cluster_io.__name__): |
| 160 | write_cluster_npz(path, _stub_channels(_CollidingChannels), SegmentType.SOLID) | 157 | cluster_io.write_cluster_npz( |
| 158 | path, _stub_channels(_CollidingChannels), types.SegmentType.SOLID | ||
| 159 | ) | ||
| 161 | assert any("scan_angle" in record.getMessage() for record in caplog.records) | 160 | assert any("scan_angle" in record.getMessage() for record in caplog.records) |
| 162 | with np.load(path) as payload: | 161 | with np.load(path) as payload: |
| 163 | assert set(payload.files) == FIXED_MEMBERS | 162 | assert set(payload.files) == FIXED_MEMBERS |
| 164 | # The fixed member survived: it is the clipped scan_angle_rank, not [7, 8]. | 163 | # The fixed member survived: it is the clipped scan_angle_rank, not [7, 8]. |
| 166 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) | 165 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) |
| 167 | 166 | ||
| 168 | 167 | ||
| 169 | def test_atomic_failure_leaves_no_partial_output( | 168 | def test_atomic_failure_leaves_no_partial_output( |
| 170 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch | 169 | tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch |
| 171 | ) -> None: | 170 | ) -> None: |
| 172 | path = tmp_path / "run6_cluster_000.npz" | 171 | path = tmp_path / "run6_cluster_000.npz" |
| 173 | 172 | ||
| 174 | def fail(handle: BinaryIO, **payload: np.ndarray) -> None: | 173 | def fail(handle: BinaryIO, **payload: np.ndarray) -> None: |
| 176 | raise RuntimeError("disk failure") | 175 | raise RuntimeError("disk failure") |
| 177 | 176 | ||
| 178 | monkeypatch.setattr(atomic_io.np, "savez", fail) | 177 | monkeypatch.setattr(atomic_io.np, "savez", fail) |
| 179 | with pytest.raises(RuntimeError): | 178 | with pytest.raises(RuntimeError): |
| 180 | write_cluster_npz(path, _sample(), SegmentType.SOLID) | 179 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.SOLID) |
| 181 | assert not path.exists() | 180 | assert not path.exists() |
| 182 | assert list(tmp_path.iterdir()) == [] | 181 | assert list(tmp_path.iterdir()) == [] |
| 6 | ``ColorIntensityData`` flows through this module with no code change here. | 6 | ``ColorIntensityData`` flows through this module with no code change here. |
| 7 | """ | 7 | """ |
| 8 | 8 | ||
| 9 | import dataclasses | 9 | import dataclasses |
| 10 | from collections.abc import Iterable, Mapping | 10 | import pathlib |
| 11 | from dataclasses import dataclass | 11 | from collections import abc |
| 12 | from pathlib import Path | ||
| 13 | 12 | ||
| 14 | import numpy as np | 13 | import numpy as np |
| 15 | from iolabs.common import segment_points_io | 14 | from iolabs import logstash |
| 16 | from iolabs.common.color_intensity_data import ColorIntensityData | 15 | from iolabs.common import color_intensity_data, segment_points_io |
| 17 | from iolabs.logstash import get_props_logger | ||
| 18 | 16 | ||
| 19 | from ._log_props import LOG_PROPS | 17 | from . import _log_props, types |
| 20 | from .types import PointChannels, RasterFrame | ||
| 21 | 18 | ||
| 22 | logger = get_props_logger(__name__, LOG_PROPS) | 19 | logger = logstash.get_props_logger(__name__, _log_props.LOG_PROPS) |
| 23 | 20 | ||
| 24 | #: Deprecated alias for the shared run3 point-record schema; use | 21 | #: Deprecated alias for the shared run3 point-record schema; use |
| 25 | #: :data:`iolabs.common.segment_points_io.POINT_RECORD_KEYS` instead. | 22 | #: :data:`iolabs.common.segment_points_io.POINT_RECORD_KEYS` instead. |
| 26 | REQUIRED_ARRAYS: tuple[str, ...] = segment_points_io.POINT_RECORD_KEYS | 23 | REQUIRED_ARRAYS: tuple[str, ...] = segment_points_io.POINT_RECORD_KEYS |
| 27 | 24 | ||
| 28 | #: Record keys whose name differs from the matching ``ColorIntensityData`` field. | 25 | #: Record keys whose name differs from the matching ``ColorIntensityData`` field. |
| 29 | _RECORD_KEY_TO_CHANNEL_FIELD: Mapping[str, str] = {"scan_angle": "scan_angle_rank"} | 26 | _RECORD_KEY_TO_CHANNEL_FIELD: abc.Mapping[str, str] = {"scan_angle": "scan_angle_rank"} |
| 30 | 27 | ||
| 31 | #: Keys the lean separation path reads directly from the npz. | 28 | #: Keys the lean separation path reads directly from the npz. |
| 32 | _LEAN_KEYS: tuple[str, ...] = ("points", "intensity") | 29 | _LEAN_KEYS: tuple[str, ...] = ("points", "intensity") |
| 33 | 30 |
| 43 | Derived from the installed ``iolabs.common`` contract: a key added to both | 40 | Derived from the installed ``iolabs.common`` contract: a key added to both |
| 44 | the point-record schema and ``ColorIntensityData`` appears here | 41 | the point-record schema and ``ColorIntensityData`` appears here |
| 45 | automatically, and keys the installed dataclass cannot hold are skipped. | 42 | automatically, and keys the installed dataclass cannot hold are skipped. |
| 46 | """ | 43 | """ |
| 47 | fields = {field.name for field in dataclasses.fields(ColorIntensityData)} | 44 | fields = {field.name for field in dataclasses.fields(color_intensity_data.ColorIntensityData)} |
| 48 | return tuple( | 45 | return tuple( |
| 49 | key | 46 | key |
| 50 | for key in segment_points_io.POINT_RECORD_KEYS | 47 | for key in segment_points_io.POINT_RECORD_KEYS |
| 51 | if key != "points" and _channel_field(key) in fields | 48 | if key != "points" and _channel_field(key) in fields |
| 52 | ) | 49 | ) |
| 53 | 50 | ||
| 54 | 51 | ||
| 55 | def channel_dtypes_of(data: ColorIntensityData) -> dict[str, np.dtype]: | 52 | def channel_dtypes_of(data: color_intensity_data.ColorIntensityData) -> dict[str, np.dtype]: |
| 56 | """Return *data*'s per-channel storage dtypes, keyed by run3 record key.""" | 53 | """Return *data*'s per-channel storage dtypes, keyed by run3 record key.""" |
| 57 | return { | 54 | return { |
| 58 | key: getattr(data, _channel_field(key)).dtype for key in channel_record_keys() | 55 | key: getattr(data, _channel_field(key)).dtype for key in channel_record_keys() |
| 59 | } | 56 | } |
| 60 | 57 | ||
| 61 | 58 | ||
| 62 | def _channels_from_record(record: Mapping[str, np.ndarray]) -> ColorIntensityData: | 59 | def _channels_from_record( |
| 60 | record: abc.Mapping[str, np.ndarray], | ||
| 61 | ) -> color_intensity_data.ColorIntensityData: | ||
| 63 | """Build channels from a run3 record without naming each field.""" | 62 | """Build channels from a run3 record without naming each field.""" |
| 64 | return ColorIntensityData( | 63 | return color_intensity_data.ColorIntensityData( |
| 65 | **{ | 64 | **{ |
| 66 | _channel_field(key): record[key] | 65 | _channel_field(key): record[key] |
| 67 | for key in channel_record_keys() | 66 | for key in channel_record_keys() |
| 68 | if key in record | 67 | if key in record |
| 69 | } | 68 | } |
| 70 | ) | 69 | ) |
| 71 | 70 | ||
| 72 | 71 | ||
| 73 | def _as_points(path: Path, raw: np.ndarray) -> np.ndarray: | 72 | def _as_points(path: pathlib.Path, raw: np.ndarray) -> np.ndarray: |
| 74 | points = np.asarray(raw, dtype=np.float64) | 73 | points = np.asarray(raw, dtype=np.float64) |
| 75 | if points.ndim != 2 or points.shape[1] != 3: | 74 | if points.ndim != 2 or points.shape[1] != 3: |
| 76 | raise ValueError(f"{path}: points must have shape (N, 3), got {points.shape}") | 75 | raise ValueError(f"{path}: points must have shape (N, 3), got {points.shape}") |
| 77 | return points | 76 | return points |
| 78 | 77 | ||
| 79 | 78 | ||
| 80 | def _as_channel(path: Path, name: str, raw: np.ndarray, count: int) -> np.ndarray: | 79 | def _as_channel(path: pathlib.Path, name: str, raw: np.ndarray, count: int) -> np.ndarray: |
| 81 | array = np.asarray(raw) | 80 | array = np.asarray(raw) |
| 82 | if array.ndim != 1 or len(array) != count: | 81 | if array.ndim != 1 or len(array) != count: |
| 83 | raise ValueError( | 82 | raise ValueError( |
| 84 | f"{path}: {name} must have shape ({count},), got {array.shape}" | 83 | f"{path}: {name} must have shape ({count},), got {array.shape}" |
| 85 | ) | 84 | ) |
| 86 | return array | 85 | return array |
| 87 | 86 | ||
| 88 | 87 | ||
| 89 | @dataclass(frozen=True) | 88 | @dataclasses.dataclass(frozen=True) |
| 90 | class SeparationArrays: | 89 | class SeparationArrays: |
| 91 | """Preallocated full-segment arrays for intensity separation.""" | 90 | """Preallocated full-segment arrays for intensity separation.""" |
| 92 | 91 | ||
| 93 | xyz: np.ndarray | 92 | xyz: np.ndarray |
| 94 | rows: np.ndarray | 93 | rows: np.ndarray |
| 95 | cols: np.ndarray | 94 | cols: np.ndarray |
| 96 | intensity: np.ndarray | 95 | intensity: np.ndarray |
| 97 | channels: ColorIntensityData | None | 96 | channels: color_intensity_data.ColorIntensityData | None |
| 98 | 97 | ||
| 99 | 98 | ||
| 100 | def load_step3_file(path: Path) -> PointChannels: | 99 | def load_step3_file(path: pathlib.Path) -> types.PointChannels: |
| 101 | """Load one Step 3 chunk with all its channels. | 100 | """Load one Step 3 chunk with all its channels. |
| 102 | 101 | ||
| 103 | The record is read and schema-validated by | 102 | The record is read and schema-validated by |
| 104 | :func:`iolabs.common.segment_points_io.load_points_npz`; XYZ is then cast to | 103 | :func:`iolabs.common.segment_points_io.load_points_npz`; XYZ is then cast to |
| 113 | Raises: | 112 | Raises: |
| 114 | ValueError: A required array is missing or has the wrong shape. | 113 | ValueError: A required array is missing or has the wrong shape. |
| 115 | """ | 114 | """ |
| 116 | record = segment_points_io.load_points_npz(path) | 115 | record = segment_points_io.load_points_npz(path) |
| 117 | return PointChannels( | 116 | return types.PointChannels( |
| 118 | points=_as_points(path, record["points"]), | 117 | points=_as_points(path, record["points"]), |
| 119 | data=_channels_from_record(record), | 118 | data=_channels_from_record(record), |
| 120 | ) | 119 | ) |
| 121 | 120 | ||
| 122 | 121 | ||
| 123 | def _load_separation_chunk( | 122 | def _load_separation_chunk( |
| 124 | path: Path, *, with_channels: bool | 123 | path: pathlib.Path, *, with_channels: bool |
| 125 | ) -> tuple[np.ndarray, np.ndarray, ColorIntensityData | None]: | 124 | ) -> tuple[np.ndarray, np.ndarray, color_intensity_data.ColorIntensityData | None]: |
| 126 | """Load one Step 3 chunk for separation. | 125 | """Load one Step 3 chunk for separation. |
| 127 | 126 | ||
| 128 | When ``with_channels`` is false, only ``points`` and ``intensity`` are read | 127 | When ``with_channels`` is false, only ``points`` and ``intensity`` are read |
| 129 | from the npz; the other record keys are neither materialized nor validated | 128 | from the npz; the other record keys are neither materialized nor validated |
| 146 | intensity = _as_channel(path, "intensity", payload["intensity"], len(points)) | 145 | intensity = _as_channel(path, "intensity", payload["intensity"], len(points)) |
| 147 | return points, intensity, None | 146 | return points, intensity, None |
| 148 | 147 | ||
| 149 | 148 | ||
| 150 | def load_step3_points(paths: Iterable[Path]) -> PointChannels: | 149 | def load_step3_points(paths: abc.Iterable[pathlib.Path]) -> types.PointChannels: |
| 151 | """Load and concatenate several Step 3 chunks. | 150 | """Load and concatenate several Step 3 chunks. |
| 152 | 151 | ||
| 153 | Args: | 152 | Args: |
| 154 | paths: Step 3 ``.npz`` chunks; loaded in sorted order. | 153 | paths: Step 3 ``.npz`` chunks; loaded in sorted order. |
| 165 | points = np.concatenate([chunk.points for chunk in chunks]) | 164 | points = np.concatenate([chunk.points for chunk in chunks]) |
| 166 | data = chunks[0].data | 165 | data = chunks[0].data |
| 167 | for chunk in chunks[1:]: | 166 | for chunk in chunks[1:]: |
| 168 | data = data.append(chunk.data) | 167 | data = data.append(chunk.data) |
| 169 | return PointChannels(points=points, data=data) | 168 | return types.PointChannels(points=points, data=data) |
| 170 | 169 | ||
| 171 | 170 | ||
| 172 | def load_separation_arrays( | 171 | def load_separation_arrays( |
| 173 | paths: Iterable[Path], | 172 | paths: abc.Iterable[pathlib.Path], |
| 174 | frame: RasterFrame, | 173 | frame: types.RasterFrame, |
| 175 | total_count: int, | 174 | total_count: int, |
| 176 | *, | 175 | *, |
| 177 | channel_dtypes: Mapping[str, np.dtype], | 176 | channel_dtypes: abc.Mapping[str, np.dtype], |
| 178 | with_channels: bool, | 177 | with_channels: bool, |
| 179 | ) -> SeparationArrays: | 178 | ) -> SeparationArrays: |
| 180 | """Stream Step 3 chunks into preallocated separation arrays. | 179 | """Stream Step 3 chunks into preallocated separation arrays. |
| 181 | 180 |
| 250 | ) | 249 | ) |
| 251 | 250 | ||
| 252 | channels = None | 251 | channels = None |
| 253 | if with_channels: | 252 | if with_channels: |
| 254 | channels = ColorIntensityData( | 253 | channels = color_intensity_data.ColorIntensityData( |
| 255 | intensity=intensity, | 254 | intensity=intensity, |
| 256 | **{_channel_field(key): buffers[key] for key in extra_keys}, | 255 | **{_channel_field(key): buffers[key] for key in extra_keys}, |
| 257 | ) | 256 | ) |
| 258 | return SeparationArrays( | 257 | return SeparationArrays( |
| 1 | """Run the mask-clustering step end to end for one segment.""" | 1 | """Run the mask-clustering step end to end for one segment.""" |
| 2 | 2 | ||
| 3 | import datetime | ||
| 3 | import os | 4 | import os |
| 5 | import pathlib | ||
| 4 | import tempfile | 6 | import tempfile |
| 5 | from datetime import UTC, datetime | 7 | from importlib import metadata |
| 6 | from importlib.metadata import PackageNotFoundError, version | ||
| 7 | from pathlib import Path | ||
| 8 | from typing import Any | 8 | from typing import Any |
| 9 | 9 | ||
| 10 | import numpy as np | 10 | import numpy as np |
| 11 | from iolabs.common.color_intensity_data import ColorIntensityData | 11 | from iolabs import logstash |
| 12 | from iolabs.common.run_stats import read_stats, write_stats | 12 | from iolabs.common import color_intensity_data, run_stats, segments, version_info |
| 13 | from iolabs.common.segments import segment_record_files | ||
| 14 | from iolabs.common.version_info import save_version_json | ||
| 15 | from iolabs.logstash import get_props_logger | ||
| 16 | 13 | ||
| 17 | from ._config import MaskClusteringConfig | 14 | from . import ( |
| 18 | from ._log_props import LOG_PROPS | 15 | _config, |
| 19 | from .cluster_io import write_cluster_npz | 16 | _log_props, |
| 20 | from .geometry import component_geojson_geometry, mask_polygons_xy, write_prism_ply | 17 | cluster_io, |
| 21 | from .input_io import channel_dtypes_of, load_separation_arrays, load_step3_file | 18 | geometry, |
| 22 | from .intensity_separation import ( | 19 | input_io, |
| 23 | ClusterSeparation, | 20 | intensity_separation, |
| 24 | dilate_mask, | 21 | mask_components, |
| 25 | separate_cluster, | 22 | point_projection, |
| 26 | write_separation_pdf, | 23 | raster_frame, |
| 24 | types, | ||
| 27 | ) | 25 | ) |
| 28 | from .mask_components import label_components, load_classified_mask | ||
| 29 | from .point_projection import assign_point_chunks | ||
| 30 | from .raster_frame import reconstruct_raster_frame | ||
| 31 | from .types import ComponentResult, PointChannels, SegmentResult, SegmentType | ||
| 32 | 26 | ||
| 33 | logger = get_props_logger(__name__, LOG_PROPS) | 27 | logger = logstash.get_props_logger(__name__, _log_props.LOG_PROPS) |
| 34 | 28 | ||
| 35 | 29 | ||
| 36 | def apply_sparse_policy( | 30 | def apply_sparse_policy( |
| 37 | *, | 31 | *, |
| 38 | segment_name: str, | 32 | segment_name: str, |
| 39 | component_id: int, | 33 | component_id: int, |
| 40 | segment_type: SegmentType, | 34 | segment_type: types.SegmentType, |
| 41 | point_count: int, | 35 | point_count: int, |
| 42 | warn_below_points: int, | 36 | warn_below_points: int, |
| 43 | min_points_per_cluster: int, | 37 | min_points_per_cluster: int, |
| 44 | logger: Any, | 38 | logger: Any, |
| 76 | 70 | ||
| 77 | 71 | ||
| 78 | def process_segment( | 72 | def process_segment( |
| 79 | *, | 73 | *, |
| 80 | segment_dir: Path, | 74 | segment_dir: pathlib.Path, |
| 81 | mask_or_vectors_path: Path, | 75 | mask_or_vectors_path: pathlib.Path, |
| 82 | raster_metadata_path: Path, | 76 | raster_metadata_path: pathlib.Path, |
| 83 | output_dir: Path, | 77 | output_dir: pathlib.Path, |
| 84 | config: MaskClusteringConfig | dict[str, Any], | 78 | config: _config.MaskClusteringConfig | dict[str, Any], |
| 85 | overwrite: bool = False, | 79 | overwrite: bool = False, |
| 86 | ) -> SegmentResult: | 80 | ) -> types.SegmentResult: |
| 87 | """Cluster one segment's classified mask into Step 6-compatible artifacts. | 81 | """Cluster one segment's classified mask into Step 6-compatible artifacts. |
| 88 | 82 | ||
| 89 | Args: | 83 | Args: |
| 90 | segment_dir: Segment directory holding the Step 3 point NPZ chunks. | 84 | segment_dir: Segment directory holding the Step 3 point NPZ chunks. |
| 105 | the configured suffix, or *raster_metadata_path* does not exist. | 99 | the configured suffix, or *raster_metadata_path* does not exist. |
| 106 | ValueError: The segment holds no points, or the mask, sidecar and | 100 | ValueError: The segment holds no points, or the mask, sidecar and |
| 107 | reconstructed frame disagree. | 101 | reconstructed frame disagree. |
| 108 | """ | 102 | """ |
| 109 | started_at = datetime.now(UTC) | 103 | started_at = datetime.datetime.now(datetime.UTC) |
| 110 | cfg = MaskClusteringConfig.coerce(config) | 104 | cfg = _config.MaskClusteringConfig.coerce(config) |
| 111 | manifest_path = output_dir / cfg.output.manifest_filename | 105 | manifest_path = output_dir / cfg.output.manifest_filename |
| 112 | if manifest_path.exists(): | 106 | if manifest_path.exists(): |
| 113 | if not overwrite: | 107 | if not overwrite: |
| 114 | raise FileExistsError( | 108 | raise FileExistsError( |
| 115 | f"Completed mask-clustering manifest already exists: {manifest_path}" | 109 | f"Completed mask-clustering manifest already exists: {manifest_path}" |
| 116 | ) | 110 | ) |
| 117 | _remove_prior_artifacts(output_dir, manifest_path) | 111 | _remove_prior_artifacts(output_dir, manifest_path) |
| 118 | 112 | ||
| 119 | step3_paths = segment_record_files( | 113 | step3_paths = segments.segment_record_files( |
| 120 | segment_dir, cfg.file_naming.segment_points_suffix, required=True | 114 | segment_dir, cfg.file_naming.segment_points_suffix, required=True |
| 121 | ) | 115 | ) |
| 122 | metadata = read_stats(raster_metadata_path, require_dict=True) | 116 | segment_metadata = run_stats.read_stats(raster_metadata_path, require_dict=True) |
| 123 | segment_name = str( | 117 | segment_name = str( |
| 124 | metadata.get("segment_name", metadata.get("segment_id", segment_dir.name)) | 118 | segment_metadata.get("segment_name", segment_metadata.get("segment_id", segment_dir.name)) |
| 125 | ) | 119 | ) |
| 126 | 120 | ||
| 127 | source_count = 0 | 121 | source_count = 0 |
| 128 | xy_min = np.array([np.inf, np.inf]) | 122 | xy_min = np.array([np.inf, np.inf]) |
| 130 | z_min = np.inf | 124 | z_min = np.inf |
| 131 | z_max = -np.inf | 125 | z_max = -np.inf |
| 132 | channel_dtypes: dict[str, np.dtype] | None = None | 126 | channel_dtypes: dict[str, np.dtype] | None = None |
| 133 | for path in step3_paths: | 127 | for path in step3_paths: |
| 134 | chunk = load_step3_file(path) | 128 | chunk = input_io.load_step3_file(path) |
| 135 | source_count += len(chunk.points) | 129 | source_count += len(chunk.points) |
| 136 | chunk_dtypes = channel_dtypes_of(chunk.data) | 130 | chunk_dtypes = input_io.channel_dtypes_of(chunk.data) |
| 137 | if channel_dtypes is None: | 131 | if channel_dtypes is None: |
| 138 | channel_dtypes = chunk_dtypes | 132 | channel_dtypes = chunk_dtypes |
| 139 | else: | 133 | else: |
| 140 | channel_dtypes = { | 134 | channel_dtypes = { |
| 153 | if source_count == 0: | 147 | if source_count == 0: |
| 154 | raise ValueError(f"Step 3 source in {segment_dir} contains no points") | 148 | raise ValueError(f"Step 3 source in {segment_dir} contains no points") |
| 155 | assert channel_dtypes is not None | 149 | assert channel_dtypes is not None |
| 156 | 150 | ||
| 157 | mask_shape = (int(metadata["pixel_y"]), int(metadata["pixel_x"])) | 151 | mask_shape = (int(segment_metadata["pixel_y"]), int(segment_metadata["pixel_x"])) |
| 158 | mask_cfg = cfg.mask | 152 | mask_cfg = cfg.mask |
| 159 | extrema = np.array( | 153 | extrema = np.array( |
| 160 | [[xy_min[0], xy_min[1], z_min], [xy_max[0], xy_max[1], z_max]], | 154 | [[xy_min[0], xy_min[1], z_min], [xy_max[0], xy_max[1], z_max]], |
| 161 | dtype=np.float64, | 155 | dtype=np.float64, |
| 162 | ) | 156 | ) |
| 163 | frame = reconstruct_raster_frame( | 157 | frame = raster_frame.reconstruct_raster_frame( |
| 164 | extrema, | 158 | extrema, |
| 165 | metadata, | 159 | segment_metadata, |
| 166 | mask_shape=mask_shape, | 160 | mask_shape=mask_shape, |
| 167 | margin_pixels=cfg.raster_frame.margin_pixels, | 161 | margin_pixels=cfg.raster_frame.margin_pixels, |
| 168 | origin_tolerance_pixels=cfg.raster_frame.metadata_origin_tolerance_pixels, | 162 | origin_tolerance_pixels=cfg.raster_frame.metadata_origin_tolerance_pixels, |
| 169 | ) | 163 | ) |
| 170 | classified_mask = load_classified_mask( | 164 | classified_mask = mask_components.load_classified_mask( |
| 171 | mask_or_vectors_path, | 165 | mask_or_vectors_path, |
| 172 | shape=mask_shape, | 166 | shape=mask_shape, |
| 173 | background_class=mask_cfg.background_class, | 167 | background_class=mask_cfg.background_class, |
| 174 | solid_class=mask_cfg.solid_class, | 168 | solid_class=mask_cfg.solid_class, |
| 175 | dashed_class=mask_cfg.dashed_class, | 169 | dashed_class=mask_cfg.dashed_class, |
| 176 | vector_stroke_px=mask_cfg.vector_stroke_px, | 170 | vector_stroke_px=mask_cfg.vector_stroke_px, |
| 177 | ) | 171 | ) |
| 178 | components, label_image = label_components( | 172 | components, label_image = mask_components.label_components( |
| 179 | classified_mask, | 173 | classified_mask, |
| 180 | solid_class=mask_cfg.solid_class, | 174 | solid_class=mask_cfg.solid_class, |
| 181 | dashed_class=mask_cfg.dashed_class, | 175 | dashed_class=mask_cfg.dashed_class, |
| 182 | connectivity=mask_cfg.connectivity, | 176 | connectivity=mask_cfg.connectivity, |
| 187 | segment_name, | 181 | segment_name, |
| 188 | extra={"segment_name": segment_name, "action": "empty_success"}, | 182 | extra={"segment_name": segment_name, "action": "empty_success"}, |
| 189 | ) | 183 | ) |
| 190 | 184 | ||
| 191 | assigned, out_of_frame = assign_point_chunks( | 185 | assigned, out_of_frame = point_projection.assign_point_chunks( |
| 192 | (load_step3_file(path) for path in step3_paths), | 186 | (input_io.load_step3_file(path) for path in step3_paths), |
| 193 | frame, | 187 | frame, |
| 194 | label_image, | 188 | label_image, |
| 195 | components, | 189 | components, |
| 196 | ) | 190 | ) |
| 209 | # point-cloud viewer. Written to a separate subdir so Step 6b never rasterises them. | 203 | # point-cloud viewer. Written to a separate subdir so Step 6b never rasterises them. |
| 210 | save_padded = separation_enabled and bool(separation_cfg.get("save_padded_clusters")) | 204 | save_padded = separation_enabled and bool(separation_cfg.get("save_padded_clusters")) |
| 211 | padding_debug_dir = output_dir / "padding_debug" | 205 | padding_debug_dir = output_dir / "padding_debug" |
| 212 | dilation_px = int(separation_cfg["dilation_px"]) | 206 | dilation_px = int(separation_cfg["dilation_px"]) |
| 213 | separations: list[ClusterSeparation] = [] | 207 | separations: list[intensity_separation.ClusterSeparation] = [] |
| 214 | # The asphalt sampling ring lies outside each component's mask, so it needs more | 208 | # The asphalt sampling ring lies outside each component's mask, so it needs more |
| 215 | # than the per-component assigned points. Stream-load and project the full segment | 209 | # than the per-component assigned points. Stream-load and project the full segment |
| 216 | # once, up front, when separation is enabled. (Cropping to the union of dilated | 210 | # once, up front, when separation is enabled. (Cropping to the union of dilated |
| 217 | # component masks would also work; full-segment streaming is simpler and the lean | 211 | # component masks would also work; full-segment streaming is simpler and the lean |
| 218 | # xyz/rows/cols/intensity buffers keep it affordable.) | 212 | # xyz/rows/cols/intensity buffers keep it affordable.) |
| 219 | sep_channels: ColorIntensityData | None = None | 213 | sep_channels: color_intensity_data.ColorIntensityData | None = None |
| 220 | if separation_enabled: | 214 | if separation_enabled: |
| 221 | sep = load_separation_arrays( | 215 | sep = input_io.load_separation_arrays( |
| 222 | step3_paths, | 216 | step3_paths, |
| 223 | frame, | 217 | frame, |
| 224 | source_count, | 218 | source_count, |
| 225 | channel_dtypes=channel_dtypes, | 219 | channel_dtypes=channel_dtypes, |
| 246 | warn_below_points=cluster_cfg.warn_below_points, | 240 | warn_below_points=cluster_cfg.warn_below_points, |
| 247 | min_points_per_cluster=cluster_cfg.min_points_per_cluster, | 241 | min_points_per_cluster=cluster_cfg.min_points_per_cluster, |
| 248 | logger=logger, | 242 | logger=logger, |
| 249 | ) | 243 | ) |
| 250 | separation: ClusterSeparation | None = None | 244 | separation: intensity_separation.ClusterSeparation | None = None |
| 251 | cluster_path: Path | None = None | 245 | cluster_path: pathlib.Path | None = None |
| 252 | padding_debug_path: Path | None = None | 246 | padding_debug_path: pathlib.Path | None = None |
| 253 | if status == "retained": | 247 | if status == "retained": |
| 254 | if separation_enabled: | 248 | if separation_enabled: |
| 255 | separation = separate_cluster( | 249 | separation = intensity_separation.separate_cluster( |
| 256 | component, | 250 | component, |
| 257 | channels.data.intensity, | 251 | channels.data.intensity, |
| 258 | points_xyz=sep_xyz, | 252 | points_xyz=sep_xyz, |
| 259 | rows=sep_rows, | 253 | rows=sep_rows, |
| 262 | cfg=separation_cfg, | 256 | cfg=separation_cfg, |
| 263 | ) | 257 | ) |
| 264 | separations.append(separation) | 258 | separations.append(separation) |
| 265 | if apply_filter and separation.cutoff is not None: | 259 | if apply_filter and separation.cutoff is not None: |
| 266 | channels = PointChannels( | 260 | channels = types.PointChannels( |
| 267 | points=channels.points[separation.paint_mask], | 261 | points=channels.points[separation.paint_mask], |
| 268 | data=channels.data.select_by_mask(separation.paint_mask), | 262 | data=channels.data.select_by_mask(separation.paint_mask), |
| 269 | ) | 263 | ) |
| 270 | cluster_path = output_dir / f"{cluster_prefix}{retained_index:03d}.npz" | 264 | cluster_path = output_dir / f"{cluster_prefix}{retained_index:03d}.npz" |
| 271 | write_cluster_npz(cluster_path, channels, component.segment_type) | 265 | cluster_io.write_cluster_npz(cluster_path, channels, component.segment_type) |
| 272 | if save_padded: | 266 | if save_padded: |
| 273 | assert sep_channels is not None | 267 | assert sep_channels is not None |
| 274 | padded_mask = dilate_mask(component.binary_mask, dilation_px) | 268 | padded_mask = intensity_separation.dilate_mask(component.binary_mask, dilation_px) |
| 275 | in_padded = padded_mask[sep_rows, sep_cols] | 269 | in_padded = padded_mask[sep_rows, sep_cols] |
| 276 | padded_channels = PointChannels( | 270 | padded_channels = types.PointChannels( |
| 277 | points=sep_xyz[in_padded], | 271 | points=sep_xyz[in_padded], |
| 278 | data=sep_channels.select_by_mask(in_padded), | 272 | data=sep_channels.select_by_mask(in_padded), |
| 279 | ) | 273 | ) |
| 280 | padding_debug_dir.mkdir(parents=True, exist_ok=True) | 274 | padding_debug_dir.mkdir(parents=True, exist_ok=True) |
| 281 | padding_debug_path = padding_debug_dir / f"padded_{retained_index:03d}.npz" | 275 | padding_debug_path = padding_debug_dir / f"padded_{retained_index:03d}.npz" |
| 282 | write_cluster_npz( | 276 | cluster_io.write_cluster_npz( |
| 283 | padding_debug_path, padded_channels, component.segment_type | 277 | padding_debug_path, padded_channels, component.segment_type |
| 284 | ) | 278 | ) |
| 285 | retained_index += 1 | 279 | retained_index += 1 |
| 286 | 280 | ||
| 287 | prism_path = geometry_dir / ( | 281 | prism_path = geometry_dir / ( |
| 288 | f"component_{component.component_id:04d}_{component.segment_type.value}_prism.ply" | 282 | f"component_{component.component_id:04d}_{component.segment_type.value}_prism.ply" |
| 289 | ) | 283 | ) |
| 290 | polygons = mask_polygons_xy( | 284 | polygons = geometry.mask_polygons_xy( |
| 291 | component.binary_mask, | 285 | component.binary_mask, |
| 292 | frame, | 286 | frame, |
| 293 | simplify_tolerance_px=geometry_cfg.simplify_tolerance_px, | 287 | simplify_tolerance_px=geometry_cfg.simplify_tolerance_px, |
| 294 | ) | 288 | ) |
| 295 | if geometry_cfg.write_ply: | 289 | if geometry_cfg.write_ply: |
| 296 | write_prism_ply( | 290 | geometry.write_prism_ply( |
| 297 | prism_path, | 291 | prism_path, |
| 298 | polygons, | 292 | polygons, |
| 299 | z_min=float(z_min), | 293 | z_min=float(z_min), |
| 300 | z_max=float(z_max), | 294 | z_max=float(z_max), |
| 341 | } | 335 | } |
| 342 | geojson_features.append( | 336 | geojson_features.append( |
| 343 | { | 337 | { |
| 344 | "type": "Feature", | 338 | "type": "Feature", |
| 345 | "geometry": component_geojson_geometry(polygons), | 339 | "geometry": geometry.component_geojson_geometry(polygons), |
| 346 | "properties": properties, | 340 | "properties": properties, |
| 347 | } | 341 | } |
| 348 | ) | 342 | ) |
| 349 | warning_state = point_count < cluster_cfg.warn_below_points | 343 | warning_state = point_count < cluster_cfg.warn_below_points |
| 361 | **separation_fields, | 355 | **separation_fields, |
| 362 | } | 356 | } |
| 363 | ) | 357 | ) |
| 364 | component_results.append( | 358 | component_results.append( |
| 365 | ComponentResult( | 359 | types.ComponentResult( |
| 366 | component_id=component.component_id, | 360 | component_id=component.component_id, |
| 367 | segment_type=component.segment_type, | 361 | segment_type=component.segment_type, |
| 368 | mask_pixel_count=component.pixel_count, | 362 | mask_pixel_count=component.pixel_count, |
| 369 | point_count=point_count, | 363 | point_count=point_count, |
| 387 | 381 | ||
| 388 | separation_pdf_relative: str | None = None | 382 | separation_pdf_relative: str | None = None |
| 389 | if separation_enabled: | 383 | if separation_enabled: |
| 390 | separation_pdf_path = output_dir / str(separation_cfg["pdf_filename"]) | 384 | separation_pdf_path = output_dir / str(separation_cfg["pdf_filename"]) |
| 391 | write_separation_pdf(separations, separation_pdf_path, separation_cfg) | 385 | intensity_separation.write_separation_pdf(separations, separation_pdf_path, separation_cfg) |
| 392 | separation_pdf_relative = separation_pdf_path.relative_to(output_dir).as_posix() | 386 | separation_pdf_relative = separation_pdf_path.relative_to(output_dir).as_posix() |
| 393 | 387 | ||
| 394 | versions_path = output_dir / "run6c_versions.json" | 388 | versions_path = output_dir / "run6c_versions.json" |
| 395 | save_version_json(versions_path, "step6c_mask_clustering") | 389 | version_info.save_version_json(versions_path, "step6c_mask_clustering") |
| 396 | dropped_count = len(components) - retained_index | 390 | dropped_count = len(components) - retained_index |
| 397 | package_version = _package_version() | 391 | package_version = _package_version() |
| 398 | completed_at = datetime.now(UTC) | 392 | completed_at = datetime.datetime.now(datetime.UTC) |
| 399 | manifest = { | 393 | manifest = { |
| 400 | "schema": "iolabs.mask_clustering/1", | 394 | "schema": "iolabs.mask_clustering/1", |
| 401 | "package_version": package_version, | 395 | "package_version": package_version, |
| 402 | "timestamps": { | 396 | "timestamps": { |
| 422 | "configuration": cfg.raw, | 416 | "configuration": cfg.raw, |
| 423 | "counts": { | 417 | "counts": { |
| 424 | "foreground_components": len(components), | 418 | "foreground_components": len(components), |
| 425 | "solid_components": sum( | 419 | "solid_components": sum( |
| 426 | item.segment_type is SegmentType.SOLID for item in components | 420 | item.segment_type is types.SegmentType.SOLID for item in components |
| 427 | ), | 421 | ), |
| 428 | "dashed_components": sum( | 422 | "dashed_components": sum( |
| 429 | item.segment_type is SegmentType.DASHED for item in components | 423 | item.segment_type is types.SegmentType.DASHED for item in components |
| 430 | ), | 424 | ), |
| 431 | "retained_clusters": retained_index, | 425 | "retained_clusters": retained_index, |
| 432 | "dropped_clusters": dropped_count, | 426 | "dropped_clusters": dropped_count, |
| 433 | "out_of_frame_points": out_of_frame, | 427 | "out_of_frame_points": out_of_frame, |
| 452 | "dashed_count": manifest["counts"]["dashed_components"], | 446 | "dashed_count": manifest["counts"]["dashed_components"], |
| 453 | "warning_count": sum(item["warning"] for item in manifest_components), | 447 | "warning_count": sum(item["warning"] for item in manifest_components), |
| 454 | }, | 448 | }, |
| 455 | ) | 449 | ) |
| 456 | return SegmentResult( | 450 | return types.SegmentResult( |
| 457 | segment_name=segment_name, | 451 | segment_name=segment_name, |
| 458 | manifest_path=manifest_path, | 452 | manifest_path=manifest_path, |
| 459 | source_point_count=source_count, | 453 | source_point_count=source_count, |
| 460 | retained_count=retained_index, | 454 | retained_count=retained_index, |
| 462 | components=tuple(component_results), | 456 | components=tuple(component_results), |
| 463 | ) | 457 | ) |
| 464 | 458 | ||
| 465 | 459 | ||
| 466 | def _write_json_atomic(path: Path, value: dict[str, Any]) -> None: | 460 | def _write_json_atomic(path: pathlib.Path, value: dict[str, Any]) -> None: |
| 467 | path.parent.mkdir(parents=True, exist_ok=True) | 461 | path.parent.mkdir(parents=True, exist_ok=True) |
| 468 | descriptor, name = tempfile.mkstemp( | 462 | descriptor, name = tempfile.mkstemp( |
| 469 | prefix=f".{path.name}.", suffix=".tmp", dir=path.parent | 463 | prefix=f".{path.name}.", suffix=".tmp", dir=path.parent |
| 470 | ) | 464 | ) |
| 471 | os.close(descriptor) | 465 | os.close(descriptor) |
| 472 | try: | 466 | try: |
| 473 | write_stats(name, value, sort_keys=True, trailing_newline=True) | 467 | run_stats.write_stats(name, value, sort_keys=True, trailing_newline=True) |
| 474 | os.replace(name, path) | 468 | os.replace(name, path) |
| 475 | finally: | 469 | finally: |
| 476 | Path(name).unlink(missing_ok=True) | 470 | pathlib.Path(name).unlink(missing_ok=True) |
| 477 | 471 | ||
| 478 | 472 | ||
| 479 | def _remove_prior_artifacts(output_dir: Path, manifest_path: Path) -> None: | 473 | def _remove_prior_artifacts(output_dir: pathlib.Path, manifest_path: pathlib.Path) -> None: |
| 480 | previous = read_stats(manifest_path, require_dict=True) | 474 | previous = run_stats.read_stats(manifest_path, require_dict=True) |
| 481 | relative_paths: list[str] = [] | 475 | relative_paths: list[str] = [] |
| 482 | for component in previous.get("components", []): | 476 | for component in previous.get("components", []): |
| 483 | for key in ("cluster_path", "padding_debug_path", "prism_path"): | 477 | for key in ("cluster_path", "padding_debug_path", "prism_path"): |
| 484 | if component.get(key): | 478 | if component.get(key): |
| 497 | 491 | ||
| 498 | 492 | ||
| 499 | def _package_version() -> str: | 493 | def _package_version() -> str: |
| 500 | try: | 494 | try: |
| 501 | return version("iolabs-point-cloud-mask-clustering") | 495 | return metadata.version("iolabs-point-cloud-mask-clustering") |
| 502 | except PackageNotFoundError: | 496 | except metadata.PackageNotFoundError: |
| 503 | return "0.1.0" | 497 | return "0.1.0" |
| 1 | import dataclasses | 1 | import dataclasses |
| 2 | import logging | 2 | import logging |
| 3 | import pathlib | ||
| 3 | import zipfile | 4 | import zipfile |
| 4 | from pathlib import Path | ||
| 5 | from typing import BinaryIO | 5 | from typing import BinaryIO |
| 6 | 6 | ||
| 7 | import numpy as np | 7 | import numpy as np |
| 8 | import pytest | 8 | import pytest |
| 9 | from iolabs.common import atomic_io | 9 | from iolabs.common import atomic_io, color_intensity_data |
| 10 | from iolabs.common.color_intensity_data import ColorIntensityData | 10 | from iolabs_point_cloud_filtering_clusters import clustering_gpu_io |
| 11 | from iolabs_point_cloud_filtering_clusters.clustering_gpu_io import load_cluster_artifact_npz | ||
| 12 | 11 | ||
| 13 | from iolabs_point_cloud_mask_clustering import cluster_io | 12 | from iolabs_point_cloud_mask_clustering import cluster_io, types |
| 14 | from iolabs_point_cloud_mask_clustering.cluster_io import write_cluster_npz | ||
| 15 | from iolabs_point_cloud_mask_clustering.types import PointChannels, SegmentType | ||
| 16 | 13 | ||
| 17 | FIXED_MEMBERS = { | 14 | FIXED_MEMBERS = { |
| 18 | "points", | 15 | "points", |
| 19 | "scan_angle", | 16 | "scan_angle", |
| 43 | scan_angle: np.ndarray = dataclasses.field(default_factory=lambda: np.array([7, 8])) | 40 | scan_angle: np.ndarray = dataclasses.field(default_factory=lambda: np.array([7, 8])) |
| 44 | 41 | ||
| 45 | 42 | ||
| 46 | def _installed_channel_fields() -> set[str]: | 43 | def _installed_channel_fields() -> set[str]: |
| 47 | return {field.name for field in dataclasses.fields(ColorIntensityData)} | 44 | return {field.name for field in dataclasses.fields(color_intensity_data.ColorIntensityData)} |
| 48 | 45 | ||
| 49 | 46 | ||
| 50 | def _stub_channels(stub: type) -> PointChannels: | 47 | def _stub_channels(stub: type) -> types.PointChannels: |
| 51 | return PointChannels( | 48 | return types.PointChannels( |
| 52 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), | 49 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), |
| 53 | data=stub( | 50 | data=stub( |
| 54 | red=np.array([1.6, 2.4]), | 51 | red=np.array([1.6, 2.4]), |
| 55 | green=np.array([3.5, 4.5]), | 52 | green=np.array([3.5, 4.5]), |
| 59 | ), | 56 | ), |
| 60 | ) | 57 | ) |
| 61 | 58 | ||
| 62 | 59 | ||
| 63 | def _sample() -> PointChannels: | 60 | def _sample() -> types.PointChannels: |
| 64 | return PointChannels( | 61 | return types.PointChannels( |
| 65 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), | 62 | points=np.array([[1.123456789, 2, 3], [4, 5, 6]], dtype=np.float64), |
| 66 | data=ColorIntensityData( | 63 | data=color_intensity_data.ColorIntensityData( |
| 67 | scan_angle_rank=np.array([-200.1, 127.4]), | 64 | scan_angle_rank=np.array([-200.1, 127.4]), |
| 68 | intensity=np.array([-1.0, 70000.0]), | 65 | intensity=np.array([-1.0, 70000.0]), |
| 69 | red=np.array([1.6, 2.4]), | 66 | red=np.array([1.6, 2.4]), |
| 70 | green=np.array([3.5, 4.5]), | 67 | green=np.array([3.5, 4.5]), |
| 72 | ), | 69 | ), |
| 73 | ) | 70 | ) |
| 74 | 71 | ||
| 75 | 72 | ||
| 76 | def test_writes_step6_compatible_dtypes_and_segment_type(tmp_path: Path) -> None: | 73 | def test_writes_step6_compatible_dtypes_and_segment_type(tmp_path: pathlib.Path) -> None: |
| 77 | path = tmp_path / "run6_cluster_000.npz" | 74 | path = tmp_path / "run6_cluster_000.npz" |
| 78 | source = _sample() | 75 | source = _sample() |
| 79 | write_cluster_npz(path, source, SegmentType.DASHED) | 76 | cluster_io.write_cluster_npz(path, source, types.SegmentType.DASHED) |
| 80 | with np.load(path) as payload: | 77 | with np.load(path) as payload: |
| 81 | assert payload["points"].dtype == np.float64 | 78 | assert payload["points"].dtype == np.float64 |
| 82 | assert payload["scan_angle"].dtype == np.int8 | 79 | assert payload["scan_angle"].dtype == np.int8 |
| 83 | assert payload["intensity"].dtype == np.uint16 | 80 | assert payload["intensity"].dtype == np.uint16 |
| 87 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) | 84 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) |
| 88 | np.testing.assert_array_equal(payload["intensity"], [0, 65535]) | 85 | np.testing.assert_array_equal(payload["intensity"], [0, 65535]) |
| 89 | np.testing.assert_array_equal(payload["red"], [2, 2]) | 86 | np.testing.assert_array_equal(payload["red"], [2, 2]) |
| 90 | np.testing.assert_array_equal(payload["green"], [4, 4]) | 87 | np.testing.assert_array_equal(payload["green"], [4, 4]) |
| 91 | artifact = load_cluster_artifact_npz(path) | 88 | artifact = clustering_gpu_io.load_cluster_artifact_npz(path) |
| 92 | np.testing.assert_array_equal(artifact.points, source.points) | 89 | np.testing.assert_array_equal(artifact.points, source.points) |
| 93 | 90 | ||
| 94 | 91 | ||
| 95 | def test_members_are_stored_not_deflated(tmp_path: Path) -> None: | 92 | def test_members_are_stored_not_deflated(tmp_path: pathlib.Path) -> None: |
| 96 | # Tablecloth chunk-streams cluster archives, which only works for STORED members. | 93 | # Tablecloth chunk-streams cluster archives, which only works for STORED members. |
| 97 | path = tmp_path / "run6_cluster_000.npz" | 94 | path = tmp_path / "run6_cluster_000.npz" |
| 98 | write_cluster_npz(path, _sample(), SegmentType.SOLID) | 95 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.SOLID) |
| 99 | with zipfile.ZipFile(path) as archive: | 96 | with zipfile.ZipFile(path) as archive: |
| 100 | methods = {info.filename: info.compress_type for info in archive.infolist()} | 97 | methods = {info.filename: info.compress_type for info in archive.infolist()} |
| 101 | assert methods | 98 | assert methods |
| 102 | assert set(methods.values()) == {zipfile.ZIP_STORED} | 99 | assert set(methods.values()) == {zipfile.ZIP_STORED} |
| 103 | 100 | ||
| 104 | 101 | ||
| 105 | def test_five_field_channels_write_only_the_fixed_members(tmp_path: Path) -> None: | 102 | def test_five_field_channels_write_only_the_fixed_members(tmp_path: pathlib.Path) -> None: |
| 106 | # A ColorIntensityData declaring only the explicit five (iolabs-common 0.7.0) | 103 | # A ColorIntensityData declaring only the explicit five (iolabs-common 0.7.0) |
| 107 | # must produce exactly the pre-AI3D-382 member set: no generic extras. | 104 | # must produce exactly the pre-AI3D-382 member set: no generic extras. |
| 108 | path = tmp_path / "run6_cluster_000.npz" | 105 | path = tmp_path / "run6_cluster_000.npz" |
| 109 | write_cluster_npz(path, _stub_channels(_FiveChannels), SegmentType.SOLID) | 106 | cluster_io.write_cluster_npz(path, _stub_channels(_FiveChannels), types.SegmentType.SOLID) |
| 110 | with np.load(path) as payload: | 107 | with np.load(path) as payload: |
| 111 | assert set(payload.files) == FIXED_MEMBERS | 108 | assert set(payload.files) == FIXED_MEMBERS |
| 112 | assert payload["scan_angle"].dtype == np.int8 | 109 | assert payload["scan_angle"].dtype == np.int8 |
| 113 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) | 110 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) |
| 114 | 111 | ||
| 115 | 112 | ||
| 116 | def test_member_set_follows_the_installed_channel_schema(tmp_path: Path) -> None: | 113 | def test_member_set_follows_the_installed_channel_schema(tmp_path: pathlib.Path) -> None: |
| 117 | # Holds on any iolabs-common: the extras are exactly the installed | 114 | # Holds on any iolabs-common: the extras are exactly the installed |
| 118 | # dataclass's non-explicit init fields. | 115 | # dataclass's non-explicit init fields. |
| 119 | path = tmp_path / "run6_cluster_000.npz" | 116 | path = tmp_path / "run6_cluster_000.npz" |
| 120 | write_cluster_npz(path, _sample(), SegmentType.SOLID) | 117 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.SOLID) |
| 121 | expected_extras = _installed_channel_fields() - cluster_io.EXPLICIT_CHANNEL_FIELDS | 118 | expected_extras = _installed_channel_fields() - cluster_io.EXPLICIT_CHANNEL_FIELDS |
| 122 | with np.load(path) as payload: | 119 | with np.load(path) as payload: |
| 123 | assert set(payload.files) == FIXED_MEMBERS | expected_extras | 120 | assert set(payload.files) == FIXED_MEMBERS | expected_extras |
| 124 | 121 |
| 126 | @pytest.mark.skipif( | 123 | @pytest.mark.skipif( |
| 127 | "number_of_returns" not in _installed_channel_fields(), | 124 | "number_of_returns" not in _installed_channel_fields(), |
| 128 | reason="installed ColorIntensityData predates number_of_returns", | 125 | reason="installed ColorIntensityData predates number_of_returns", |
| 129 | ) | 126 | ) |
| 130 | def test_number_of_returns_is_forwarded_aligned_with_points(tmp_path: Path) -> None: | 127 | def test_number_of_returns_is_forwarded_aligned_with_points(tmp_path: pathlib.Path) -> None: |
| 131 | source = _sample() | 128 | source = _sample() |
| 132 | source.data.number_of_returns = np.array([1, 3], dtype=np.uint8) | 129 | source.data.number_of_returns = np.array([1, 3], dtype=np.uint8) |
| 133 | path = tmp_path / "run6_cluster_000.npz" | 130 | path = tmp_path / "run6_cluster_000.npz" |
| 134 | write_cluster_npz(path, source, SegmentType.DASHED) | 131 | cluster_io.write_cluster_npz(path, source, types.SegmentType.DASHED) |
| 135 | with np.load(path) as payload: | 132 | with np.load(path) as payload: |
| 136 | returns = payload["number_of_returns"] | 133 | returns = payload["number_of_returns"] |
| 137 | assert returns.dtype == np.uint8 | 134 | assert returns.dtype == np.uint8 |
| 138 | assert len(returns) == len(source.points) | 135 | assert len(returns) == len(source.points) |
| 142 | @pytest.mark.skipif( | 139 | @pytest.mark.skipif( |
| 143 | "number_of_returns" not in _installed_channel_fields(), | 140 | "number_of_returns" not in _installed_channel_fields(), |
| 144 | reason="installed ColorIntensityData predates number_of_returns", | 141 | reason="installed ColorIntensityData predates number_of_returns", |
| 145 | ) | 142 | ) |
| 146 | def test_omitted_number_of_returns_is_forwarded_as_unknown_zeros(tmp_path: Path) -> None: | 143 | def test_omitted_number_of_returns_is_forwarded_as_unknown_zeros(tmp_path: pathlib.Path) -> None: |
| 147 | path = tmp_path / "run6_cluster_000.npz" | 144 | path = tmp_path / "run6_cluster_000.npz" |
| 148 | write_cluster_npz(path, _sample(), SegmentType.DASHED) | 145 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.DASHED) |
| 149 | with np.load(path) as payload: | 146 | with np.load(path) as payload: |
| 150 | returns = payload["number_of_returns"] | 147 | returns = payload["number_of_returns"] |
| 151 | assert returns.dtype == np.uint8 | 148 | assert returns.dtype == np.uint8 |
| 152 | np.testing.assert_array_equal(returns, [0, 0]) | 149 | np.testing.assert_array_equal(returns, [0, 0]) |
| 153 | 150 | ||
| 154 | 151 | ||
| 155 | def test_field_colliding_with_a_fixed_member_is_skipped_with_a_warning( | 152 | def test_field_colliding_with_a_fixed_member_is_skipped_with_a_warning( |
| 156 | tmp_path: Path, caplog: pytest.LogCaptureFixture | 153 | tmp_path: pathlib.Path, caplog: pytest.LogCaptureFixture |
| 157 | ) -> None: | 154 | ) -> None: |
| 158 | path = tmp_path / "run6_cluster_000.npz" | 155 | path = tmp_path / "run6_cluster_000.npz" |
| 159 | with caplog.at_level(logging.WARNING, logger=cluster_io.__name__): | 156 | with caplog.at_level(logging.WARNING, logger=cluster_io.__name__): |
| 160 | write_cluster_npz(path, _stub_channels(_CollidingChannels), SegmentType.SOLID) | 157 | cluster_io.write_cluster_npz( |
| 158 | path, _stub_channels(_CollidingChannels), types.SegmentType.SOLID | ||
| 159 | ) | ||
| 161 | assert any("scan_angle" in record.getMessage() for record in caplog.records) | 160 | assert any("scan_angle" in record.getMessage() for record in caplog.records) |
| 162 | with np.load(path) as payload: | 161 | with np.load(path) as payload: |
| 163 | assert set(payload.files) == FIXED_MEMBERS | 162 | assert set(payload.files) == FIXED_MEMBERS |
| 164 | # The fixed member survived: it is the clipped scan_angle_rank, not [7, 8]. | 163 | # The fixed member survived: it is the clipped scan_angle_rank, not [7, 8]. |
| 166 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) | 165 | np.testing.assert_array_equal(payload["scan_angle"], [-128, 127]) |
| 167 | 166 | ||
| 168 | 167 | ||
| 169 | def test_atomic_failure_leaves_no_partial_output( | 168 | def test_atomic_failure_leaves_no_partial_output( |
| 170 | tmp_path: Path, monkeypatch: pytest.MonkeyPatch | 169 | tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch |
| 171 | ) -> None: | 170 | ) -> None: |
| 172 | path = tmp_path / "run6_cluster_000.npz" | 171 | path = tmp_path / "run6_cluster_000.npz" |
| 173 | 172 | ||
| 174 | def fail(handle: BinaryIO, **payload: np.ndarray) -> None: | 173 | def fail(handle: BinaryIO, **payload: np.ndarray) -> None: |
| 176 | raise RuntimeError("disk failure") | 175 | raise RuntimeError("disk failure") |
| 177 | 176 | ||
| 178 | monkeypatch.setattr(atomic_io.np, "savez", fail) | 177 | monkeypatch.setattr(atomic_io.np, "savez", fail) |
| 179 | with pytest.raises(RuntimeError): | 178 | with pytest.raises(RuntimeError): |
| 180 | write_cluster_npz(path, _sample(), SegmentType.SOLID) | 179 | cluster_io.write_cluster_npz(path, _sample(), types.SegmentType.SOLID) |
| 181 | assert not path.exists() | 180 | assert not path.exists() |
| 182 | assert list(tmp_path.iterdir()) == [] | 181 | assert list(tmp_path.iterdir()) == [] |