Back to report index

Step 5 maskclustering f9f1b08: AI3D-382 Load run3 records through common segment_points_io and carry channels by key

Miroslav Simko <ms@iolabs.ch> 2026-09-01T14:53:04+02:00

Commit #5 ยท 6 snippets

 src/iolabs_point_cloud_mask_clustering/input_io.py | 196 ++++++++++++---------
 1 file changed, 117 insertions(+), 79 deletions(-)

Step 5 stops parsing run3 NPZs by hand and loads them through common's segment_points_io, carrying ancillary channels by key. This is what makes the consumer schema-agnostic; REQUIRED_ARRAYS becomes a deprecated alias of POINT_RECORD_KEYS (member order changed, gains optional keys under new common).

Importance #1: src/iolabs_point_cloud_mask_clustering/input_io.py @@ -1,38 +1,74 @@

Consumer switches from hand-parsing NPZ members to segment_points_io.load_points_npz.

1"""Read Step 3 point NPZ chunks and stream them into separation buffers."""1"""Read Step 3 point NPZ chunks and stream them into separation buffers.
22
3Record loading goes through :mod:`iolabs.common.segment_points_io`, the SSOT for
4the ``*_run3_points.npz`` contract, and channels are carried by record key rather
5than field by field, so a key added to both that contract and
6``ColorIntensityData`` flows through this module with no code change here.
7"""
8
9import dataclasses
3from collections.abc import Iterable, Mapping10from collections.abc import Iterable, Mapping
4from dataclasses import dataclass11from dataclasses import dataclass
5from pathlib import Path12from pathlib import Path
613
7import numpy as np14import numpy as np
15from iolabs.common import segment_points_io
8from iolabs.common.color_intensity_data import ColorIntensityData16from iolabs.common.color_intensity_data import ColorIntensityData
9from iolabs.logstash import get_props_logger17from iolabs.logstash import get_props_logger
1018
11from ._log_props import LOG_PROPS19from ._log_props import LOG_PROPS
12from .types import PointChannels, RasterFrame20from .types import PointChannels, RasterFrame
1321
14logger = get_props_logger(__name__, LOG_PROPS)22logger = get_props_logger(__name__, LOG_PROPS)
1523
16REQUIRED_ARRAYS = ("points", "scan_angle", "intensity", "red", "green", "blue")24#: Deprecated alias for the shared run3 point-record schema; use
25#: :data:`iolabs.common.segment_points_io.POINT_RECORD_KEYS` instead.
26REQUIRED_ARRAYS: tuple[str, ...] = segment_points_io.POINT_RECORD_KEYS
1727
28#: Record keys whose name differs from the matching ``ColorIntensityData`` field.
29_RECORD_KEY_TO_CHANNEL_FIELD: Mapping[str, str] = {"scan_angle": "scan_angle_rank"}
1830
19@dataclass(frozen=True)31#: Keys the lean separation path reads directly from the npz.
20class SeparationArrays:32_LEAN_KEYS: tuple[str, ...] = ("points", "intensity")
21 """Preallocated full-segment arrays for intensity separation."""
2233
23 xyz: np.ndarray34
24 rows: np.ndarray35def _channel_field(record_key: str) -> str:
25 cols: np.ndarray36 """Map a run3 record key to its ``ColorIntensityData`` field name."""
26 intensity: np.ndarray37 return _RECORD_KEY_TO_CHANNEL_FIELD.get(record_key, record_key)
27 channels: ColorIntensityData | None38
39
40def channel_record_keys() -> tuple[str, ...]:
41 """Run3 record keys that ``ColorIntensityData`` can carry, in record order.
42
43 Derived from the installed ``iolabs.common`` contract: a key added to both
44 the point-record schema and ``ColorIntensityData`` appears here
45 automatically, and keys the installed dataclass cannot hold are skipped.
46 """
47 fields = {field.name for field in dataclasses.fields(ColorIntensityData)}
48 return tuple(
49 key
50 for key in segment_points_io.POINT_RECORD_KEYS
51 if key != "points" and _channel_field(key) in fields
52 )
53
54
55def channel_dtypes_of(data: ColorIntensityData) -> dict[str, np.dtype]:
56 """Return *data*'s per-channel storage dtypes, keyed by run3 record key."""
57 return {
58 key: getattr(data, _channel_field(key)).dtype for key in channel_record_keys()
59 }
2860
2961
30def _validate_step3_schema(path: Path, payload: object) -> None:62def _channels_from_record(record: Mapping[str, np.ndarray]) -> ColorIntensityData:
31 files = payload.files63 """Build channels from a run3 record without naming each field."""
32 missing = [name for name in REQUIRED_ARRAYS if name not in files]64 return ColorIntensityData(
33 if missing:65 **{
34 raise ValueError(f"{path} is missing required arrays: {missing}")66 _channel_field(key): record[key]
67 for key in channel_record_keys()
68 if key in record
69 }
70 )
3571
3672
37def _as_points(path: Path, raw: np.ndarray) -> np.ndarray:73def _as_points(path: Path, raw: np.ndarray) -> np.ndarray:
38 points = np.asarray(raw, dtype=np.float64)74 points = np.asarray(raw, dtype=np.float64)
Importance #2: src/iolabs_point_cloud_mask_clustering/input_io.py @@ -61,60 +112,40 @@
61112
62 Raises:113 Raises:
63 ValueError: A required array is missing or has the wrong shape.114 ValueError: A required array is missing or has the wrong shape.
64 """115 """
65 with np.load(path) as payload:116 record = segment_points_io.load_points_npz(path)
66 _validate_step3_schema(path, payload)
67 arrays = {name: np.asarray(payload[name]) for name in REQUIRED_ARRAYS}
68
69 points = _as_points(path, arrays["points"])
70 count = len(points)
71 channels = {
72 name: _as_channel(path, name, arrays[name], count)
73 for name in REQUIRED_ARRAYS[1:]
74 }
75 return PointChannels(117 return PointChannels(
76 points=points,118 points=_as_points(path, record["points"]),
77 data=ColorIntensityData(119 data=_channels_from_record(record),
78 red=channels["red"],
79 green=channels["green"],
80 blue=channels["blue"],
81 intensity=channels["intensity"],
82 scan_angle_rank=channels["scan_angle"],
83 ),
84 )120 )
85121
86122
87def _load_separation_chunk(123def _load_separation_chunk(
88 path: Path, *, with_channels: bool124 path: Path, *, with_channels: bool
89) -> tuple[np.ndarray, np.ndarray, ColorIntensityData | None]:125) -> tuple[np.ndarray, np.ndarray, ColorIntensityData | None]:
90 """Load one Step 3 chunk for separation.126 """Load one Step 3 chunk for separation.
91127
92 When ``with_channels`` is false, only ``points`` and ``intensity`` are read from128 When ``with_channels`` is false, only ``points`` and ``intensity`` are read
93 the npz (RGB/scan_angle keys are schema-checked but never materialized).129 from the npz; the other record keys are neither materialized nor validated
130 here, because every chunk is already validated up front through
131 :func:`load_step3_file`.
94 """132 """
133 if with_channels:
134 record = segment_points_io.load_points_npz(path)
135 return (
136 _as_points(path, record["points"]),
137 record["intensity"],
138 _channels_from_record(record),
139 )
140
95 with np.load(path) as payload:141 with np.load(path) as payload:
96 _validate_step3_schema(path, payload)142 missing = [name for name in _LEAN_KEYS if name not in payload.files]
143 if missing:
144 raise ValueError(f"{path} is missing required arrays: {missing}")
97 points = _as_points(path, payload["points"])145 points = _as_points(path, payload["points"])
98 count = len(points)146 intensity = _as_channel(path, "intensity", payload["intensity"], len(points))
99 intensity = _as_channel(path, "intensity", payload["intensity"], count)147 return points, intensity, None
100 if not with_channels:
101 return points, intensity, None
102 red = _as_channel(path, "red", payload["red"], count)
103 green = _as_channel(path, "green", payload["green"], count)
104 blue = _as_channel(path, "blue", payload["blue"], count)
105 scan_angle = _as_channel(path, "scan_angle", payload["scan_angle"], count)
106 return (
107 points,
108 intensity,
109 ColorIntensityData(
110 red=red,
111 green=green,
112 blue=blue,
113 intensity=intensity,
114 scan_angle_rank=scan_angle,
115 ),
116 )
117148
118149
119def load_step3_points(paths: Iterable[Path]) -> PointChannels:150def load_step3_points(paths: Iterable[Path]) -> PointChannels:
120 """Load and concatenate several Step 3 chunks.151 """Load and concatenate several Step 3 chunks.
Importance #3: src/iolabs_point_cloud_mask_clustering/input_io.py @@ -148,25 +179,41 @@
148) -> SeparationArrays:179) -> SeparationArrays:
149 """Stream Step 3 chunks into preallocated separation arrays.180 """Stream Step 3 chunks into preallocated separation arrays.
150181
151 Loads one chunk at a time so per-chunk arrays are never held alongside the182 Loads one chunk at a time so per-chunk arrays are never held alongside the
152 full-segment buffers. When ``with_channels`` is false (normal separation path),183 full-segment buffers. When ``with_channels`` is false (normal separation
153 RGB and scan_angle are neither read from disk nor allocated. They are loaded184 path), the ancillary channels are neither read from disk nor allocated. They
154 only when ``with_channels`` is true (save_padded_clusters debug path).185 are loaded only when ``with_channels`` is true (save_padded_clusters debug
186 path).
187
188 Args:
189 paths: Step 3 ``.npz`` chunks; streamed in sorted order.
190 frame: Raster frame used to project each chunk's XY.
191 total_count: Exact number of points held by the chunks together.
192 channel_dtypes: Buffer dtype per run3 record key, as produced by
193 :func:`channel_dtypes_of`. Keys other than ``intensity`` are
194 allocated only when *with_channels* is true.
195 with_channels: Whether to materialize the ancillary channels.
196
197 Returns:
198 The filled buffers; ``channels`` is set only when *with_channels*.
199
200 Raises:
201 ValueError: The chunks hold more or fewer points than *total_count*, or
202 a chunk is malformed.
155 """203 """
156 xyz = np.empty((total_count, 3), dtype=np.float64)204 xyz = np.empty((total_count, 3), dtype=np.float64)
157 rows = np.empty(total_count, dtype=np.int32)205 rows = np.empty(total_count, dtype=np.int32)
158 cols = np.empty(total_count, dtype=np.int32)206 cols = np.empty(total_count, dtype=np.int32)
159 intensity = np.empty(total_count, dtype=channel_dtypes["intensity"])207 intensity = np.empty(total_count, dtype=channel_dtypes["intensity"])
160 red: np.ndarray | None = None208 extra_keys = (
161 green: np.ndarray | None = None209 tuple(key for key in channel_dtypes if key != "intensity")
162 blue: np.ndarray | None = None210 if with_channels
163 scan_angle: np.ndarray | None = None211 else ()
164 if with_channels:212 )
165 red = np.empty(total_count, dtype=channel_dtypes["red"])213 buffers = {
166 green = np.empty(total_count, dtype=channel_dtypes["green"])214 key: np.empty(total_count, dtype=channel_dtypes[key]) for key in extra_keys
167 blue = np.empty(total_count, dtype=channel_dtypes["blue"])215 }
168 scan_angle = np.empty(total_count, dtype=channel_dtypes["scan_angle"])
169216
170 offset = 0217 offset = 0
171 for path in sorted(paths):218 for path in sorted(paths):
172 chunk_points, chunk_intensity, chunk_channels = _load_separation_chunk(219 chunk_points, chunk_intensity, chunk_channels = _load_separation_chunk(
Importance #4: src/iolabs_point_cloud_mask_clustering/input_io.py @@ -49,11 +85,26 @@
49 )85 )
50 return array86 return array
5187
5288
89@dataclass(frozen=True)
90class SeparationArrays:
91 """Preallocated full-segment arrays for intensity separation."""
92
93 xyz: np.ndarray
94 rows: np.ndarray
95 cols: np.ndarray
96 intensity: np.ndarray
97 channels: ColorIntensityData | None
98
99
53def load_step3_file(path: Path) -> PointChannels:100def load_step3_file(path: Path) -> PointChannels:
54 """Load one Step 3 chunk with all its channels.101 """Load one Step 3 chunk with all its channels.
55102
103 The record is read and schema-validated by
104 :func:`iolabs.common.segment_points_io.load_points_npz`; XYZ is then cast to
105 float64, as the rest of the pipeline expects.
106
56 Args:107 Args:
57 path: Step 3 ``.npz`` chunk.108 path: Step 3 ``.npz`` chunk.
58109
59 Returns:110 Returns:
Importance #5: src/iolabs_point_cloud_mask_clustering/input_io.py @@ -190,15 +237,11 @@
190 cols[offset:end] = chunk_cols237 cols[offset:end] = chunk_cols
191 del chunk_rows, chunk_cols238 del chunk_rows, chunk_cols
192 intensity[offset:end] = chunk_intensity239 intensity[offset:end] = chunk_intensity
193 if with_channels:240 if with_channels:
194 assert red is not None and green is not None
195 assert blue is not None and scan_angle is not None
196 assert chunk_channels is not None241 assert chunk_channels is not None
197 red[offset:end] = chunk_channels.red242 for key in extra_keys:
198 green[offset:end] = chunk_channels.green243 buffers[key][offset:end] = getattr(chunk_channels, _channel_field(key))
199 blue[offset:end] = chunk_channels.blue
200 scan_angle[offset:end] = chunk_channels.scan_angle_rank
201 del chunk_points, chunk_intensity, chunk_channels244 del chunk_points, chunk_intensity, chunk_channels
202 offset = end245 offset = end
203246
204 if offset != total_count:247 if offset != total_count:
Importance #6: src/iolabs_point_cloud_mask_clustering/input_io.py @@ -207,16 +250,11 @@
207 )250 )
208251
209 channels = None252 channels = None
210 if with_channels:253 if with_channels:
211 assert red is not None and green is not None
212 assert blue is not None and scan_angle is not None
213 channels = ColorIntensityData(254 channels = ColorIntensityData(
214 red=red,
215 green=green,
216 blue=blue,
217 intensity=intensity,255 intensity=intensity,
218 scan_angle_rank=scan_angle,256 **{_channel_field(key): buffers[key] for key in extra_keys},
219 )257 )
220 return SeparationArrays(258 return SeparationArrays(
221 xyz=xyz,259 xyz=xyz,
222 rows=rows,260 rows=rows,