Back to report index

Step 3 segmentationtrajectory 4e54928: AI3D-382 Thread run3 NPZ per-point fields as one schema-driven dict

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

Commit #98 ยท 13 snippets

 .../segment_mapper.py                              | 243 +++++++++++----------
 tests/test_segment_mapper_overflow.py              |  28 ++-
 2 files changed, 145 insertions(+), 126 deletions(-)
Importance #1: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -41,8 +41,37 @@
41#: exactly. 0 is not a legal LAS return count and means "unknown" -- it is what a41#: exactly. 0 is not a legal LAS return count and means "unknown" -- it is what a
42#: LAS without the field degrades to. Never substitute 1: that fabricates data.42#: LAS without the field degrades to. Never substitute 1: that fabricates data.
43NUMBER_OF_RETURNS_DTYPE = np.uint843NUMBER_OF_RETURNS_DTYPE = np.uint8
4444
45#: Name of the (N, 3) geometry member of the run3 NPZ contract. It is handled
46#: separately from the ancillary fields below because it is 2-D.
47POINTS_FIELD_NAME = "points"
48
49#: Ordered per-point ancillary members of the run3 NPZ contract. Everything that
50#: threads per-point arrays through the split path (extraction, masking, dtype
51#: registration, memmap allocation, archive member order) iterates this tuple, so
52#: adding a member means adding one entry here plus one extraction line in
53#: :meth:`SegmentMapper._chunk_field_arrays`.
54ANCILLARY_FIELD_NAMES: tuple[str, ...] = (
55 "scan_angle",
56 "intensity",
57 "red",
58 "green",
59 "blue",
60 NUMBER_OF_RETURNS_KEY,
61)
62
63#: Archive member order of a per-segment run3 `.npz`: geometry first, then the
64#: ancillary fields in declaration order.
65SEGMENT_NPZ_FIELD_NAMES: tuple[str, ...] = (POINTS_FIELD_NAME, *ANCILLARY_FIELD_NAMES)
66
67#: Dtypes used when a field is missing from the observed `field_dtypes` mapping.
68#: Fields absent here are looked up strictly (a missing entry is a bug).
69DEFAULT_FIELD_DTYPES: dict[str, np.dtype] = {
70 POINTS_FIELD_NAME: np.dtype(np.float64),
71 NUMBER_OF_RETURNS_KEY: np.dtype(NUMBER_OF_RETURNS_DTYPE),
72}
73
4574
46def _record_number_of_returns(75def _record_number_of_returns(
47 record: Any,76 record: Any,
48 *,77 *,
Importance #2: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -303,15 +332,19 @@
303 self,332 self,
304 segment_idx: int,333 segment_idx: int,
305 *,334 *,
306 points: np.ndarray,335 points: np.ndarray,
307 scan_angle: np.ndarray,336 fields: dict[str, np.ndarray],
308 intensity: np.ndarray,
309 red: np.ndarray,
310 green: np.ndarray,
311 blue: np.ndarray,
312 number_of_returns: np.ndarray,
313 ) -> None:337 ) -> None:
338 """Append one chunk slice to a segment buffer.
339
340 Args:
341 segment_idx: Target segment index.
342 points: `(N, 3)` geometry array in world coordinates; the geoshift is
343 subtracted on write.
344 fields: Per-point ancillary arrays keyed by
345 :data:`ANCILLARY_FIELD_NAMES`, each of length `N`.
346 """
314 point_count = int(points.shape[0])347 point_count = int(points.shape[0])
315 if point_count == 0:348 if point_count == 0:
316 return349 return
317 if segment_idx not in self.point_count_by_segment:350 if segment_idx not in self.point_count_by_segment:
Importance #3: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -326,15 +359,11 @@
326 f"Segment {segment_idx} for {self.las_file} received too many points: "359 f"Segment {segment_idx} for {self.las_file} received too many points: "
327 f"{end} > {expected}"360 f"{end} > {expected}"
328 )361 )
329362
330 arrays["points"][start:end] = points - self.geoshift363 arrays[POINTS_FIELD_NAME][start:end] = points - self.geoshift
331 arrays["scan_angle"][start:end] = scan_angle364 for field_name in ANCILLARY_FIELD_NAMES:
332 arrays["intensity"][start:end] = intensity365 arrays[field_name][start:end] = fields[field_name]
333 arrays["red"][start:end] = red
334 arrays["green"][start:end] = green
335 arrays["blue"][start:end] = blue
336 arrays[NUMBER_OF_RETURNS_KEY][start:end] = number_of_returns
337 self._offsets[segment_idx] = end366 self._offsets[segment_idx] = end
338367
339 def finalize(self) -> None:368 def finalize(self) -> None:
340 for segment_idx, expected_count in sorted(self.point_count_by_segment.items()):369 for segment_idx, expected_count in sorted(self.point_count_by_segment.items()):
Importance #4: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -369,17 +398,9 @@
369 mode="w",398 mode="w",
370 compression=zipfile.ZIP_STORED,399 compression=zipfile.ZIP_STORED,
371 allowZip64=True,400 allowZip64=True,
372 ) as archive:401 ) as archive:
373 for field_name in (402 for field_name in SEGMENT_NPZ_FIELD_NAMES:
374 "points",
375 "scan_angle",
376 "intensity",
377 "red",
378 "green",
379 "blue",
380 NUMBER_OF_RETURNS_KEY,
381 ):
382 archive.write(403 archive.write(
383 segment_tmp_dir / f"{field_name}.npy",404 segment_tmp_dir / f"{field_name}.npy",
384 arcname=f"{field_name}.npy",405 arcname=f"{field_name}.npy",
385 )406 )
Importance #5: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -389,8 +410,15 @@
389 tmp_npz_path.unlink()410 tmp_npz_path.unlink()
390 self._release_segment_arrays(segment_idx)411 self._release_segment_arrays(segment_idx)
391 shutil.rmtree(self.tmp_dir / f"segment_{segment_idx:03d}", ignore_errors=True)412 shutil.rmtree(self.tmp_dir / f"segment_{segment_idx:03d}", ignore_errors=True)
392413
414 def _dtype_for(self, field_name: str) -> np.dtype:
415 """Dtype observed for *field_name*, falling back to the schema default."""
416 default_dtype = DEFAULT_FIELD_DTYPES.get(field_name)
417 if default_dtype is None:
418 return self.field_dtypes[field_name]
419 return self.field_dtypes.get(field_name, default_dtype)
420
393 def _arrays_for_segment(self, segment_idx: int) -> dict[str, np.memmap]:421 def _arrays_for_segment(self, segment_idx: int) -> dict[str, np.memmap]:
394 arrays = self._arrays.get(segment_idx)422 arrays = self._arrays.get(segment_idx)
395 if arrays is not None:423 if arrays is not None:
396 return arrays424 return arrays
Importance #6: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -398,54 +426,22 @@
398 point_count = self.point_count_by_segment[segment_idx]426 point_count = self.point_count_by_segment[segment_idx]
399 segment_tmp_dir = self.tmp_dir / f"segment_{segment_idx:03d}"427 segment_tmp_dir = self.tmp_dir / f"segment_{segment_idx:03d}"
400 segment_tmp_dir.mkdir(parents=True, exist_ok=True)428 segment_tmp_dir.mkdir(parents=True, exist_ok=True)
401 arrays = {429 arrays = {
402 "points": np.lib.format.open_memmap(430 POINTS_FIELD_NAME: np.lib.format.open_memmap(
403 segment_tmp_dir / "points.npy",431 segment_tmp_dir / f"{POINTS_FIELD_NAME}.npy",
404 mode="w+",432 mode="w+",
405 dtype=self.field_dtypes.get("points", np.dtype(np.float64)),433 dtype=self._dtype_for(POINTS_FIELD_NAME),
406 shape=(point_count, 3),434 shape=(point_count, 3),
407 ),435 ),
408 "scan_angle": np.lib.format.open_memmap(436 }
409 segment_tmp_dir / "scan_angle.npy",437 for field_name in ANCILLARY_FIELD_NAMES:
410 mode="w+",438 arrays[field_name] = np.lib.format.open_memmap(
411 dtype=self.field_dtypes["scan_angle"],439 segment_tmp_dir / f"{field_name}.npy",
412 shape=(point_count,),
413 ),
414 "intensity": np.lib.format.open_memmap(
415 segment_tmp_dir / "intensity.npy",
416 mode="w+",
417 dtype=self.field_dtypes["intensity"],
418 shape=(point_count,),
419 ),
420 "red": np.lib.format.open_memmap(
421 segment_tmp_dir / "red.npy",
422 mode="w+",
423 dtype=self.field_dtypes["red"],
424 shape=(point_count,),
425 ),
426 "green": np.lib.format.open_memmap(
427 segment_tmp_dir / "green.npy",
428 mode="w+",
429 dtype=self.field_dtypes["green"],
430 shape=(point_count,),
431 ),
432 "blue": np.lib.format.open_memmap(
433 segment_tmp_dir / "blue.npy",
434 mode="w+",
435 dtype=self.field_dtypes["blue"],
436 shape=(point_count,),
437 ),
438 NUMBER_OF_RETURNS_KEY: np.lib.format.open_memmap(
439 segment_tmp_dir / f"{NUMBER_OF_RETURNS_KEY}.npy",
440 mode="w+",440 mode="w+",
441 dtype=self.field_dtypes.get(441 dtype=self._dtype_for(field_name),
442 NUMBER_OF_RETURNS_KEY,
443 np.dtype(NUMBER_OF_RETURNS_DTYPE),
444 ),
445 shape=(point_count,),442 shape=(point_count,),
446 ),443 )
447 }
448 self._arrays[segment_idx] = arrays444 self._arrays[segment_idx] = arrays
449 return arrays445 return arrays
450446
451447
Importance #7: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -1947,49 +1943,38 @@
1947 points = np.column_stack((chunk.x, chunk.y, chunk.z)).astype(1943 points = np.column_stack((chunk.x, chunk.y, chunk.z)).astype(
1948 np.float64,1944 np.float64,
1949 copy=False,1945 copy=False,
1950 )1946 )
1951 scan_angle = self._chunk_scan_angle(chunk, las_file=las_file, logger=logger)1947 fields = self._chunk_field_arrays(
1952 intensity = np.asarray(chunk.intensity)
1953 red, green, blue = self._chunk_rgb(
1954 chunk=chunk,
1955 las_file=las_file,
1956 )
1957 number_of_returns = _record_number_of_returns(
1958 chunk,1948 chunk,
1949 las_file=las_file,
1959 point_count=int(points.shape[0]),1950 point_count=int(points.shape[0]),
1960 source=f"{las_file.name} chunk {chunk_count}",1951 chunk_count=chunk_count,
1961 logger=logger,1952 logger=logger,
1962 )1953 )
19631954
1964 processed_points += int(points.shape[0])1955 processed_points += int(points.shape[0])
19651956
1966 if angle_limit is not None:1957 if angle_limit is not None:
1967 angle_mask = np.abs(scan_angle.astype(np.int16, copy=False)) < angle_limit1958 angle_mask = (
1959 np.abs(fields["scan_angle"].astype(np.int16, copy=False))
1960 < angle_limit
1961 )
1968 if not np.any(angle_mask):1962 if not np.any(angle_mask):
1969 continue1963 continue
1970 points = points[angle_mask]1964 points = points[angle_mask]
1971 scan_angle = scan_angle[angle_mask]1965 fields = {
1972 intensity = intensity[angle_mask]1966 field_name: array[angle_mask]
1973 red = red[angle_mask]1967 for field_name, array in fields.items()
1974 green = green[angle_mask]1968 }
1975 blue = blue[angle_mask]
1976 number_of_returns = number_of_returns[angle_mask]
19771969
1978 if points.shape[0] == 0:1970 if points.shape[0] == 0:
1979 continue1971 continue
19801972
1981 kept_points += int(points.shape[0])1973 kept_points += int(points.shape[0])
1982 field_dtypes.setdefault("points", np.dtype(points.dtype))1974 field_dtypes.setdefault(POINTS_FIELD_NAME, np.dtype(points.dtype))
1983 field_dtypes.setdefault("scan_angle", np.dtype(scan_angle.dtype))1975 for field_name, array in fields.items():
1984 field_dtypes.setdefault("intensity", np.dtype(intensity.dtype))1976 field_dtypes.setdefault(field_name, np.dtype(array.dtype))
1985 field_dtypes.setdefault("red", np.dtype(red.dtype))
1986 field_dtypes.setdefault("green", np.dtype(green.dtype))
1987 field_dtypes.setdefault("blue", np.dtype(blue.dtype))
1988 field_dtypes.setdefault(
1989 NUMBER_OF_RETURNS_KEY,
1990 np.dtype(number_of_returns.dtype),
1991 )
19921977
1993 logger.debug(1978 logger.debug(
1994 "Chunk %d: processing %d filtered points",1979 "Chunk %d: processing %d filtered points",
1995 chunk_count,1980 chunk_count,
Importance #8: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -2083,14 +2068,12 @@
2083 )2068 )
2084 continue2069 continue
20852070
2086 points_to_save = points[points_mask]2071 points_to_save = points[points_mask]
2087 scan_angle_to_save = scan_angle[points_mask]2072 fields_to_save = {
2088 intensity_to_save = intensity[points_mask]2073 field_name: array[points_mask]
2089 red_to_save = red[points_mask]2074 for field_name, array in fields.items()
2090 green_to_save = green[points_mask]2075 }
2091 blue_to_save = blue[points_mask]
2092 number_of_returns_to_save = number_of_returns[points_mask]
2093 if longitudinal_planes is not None:2076 if longitudinal_planes is not None:
2094 left_plane, right_plane = longitudinal_planes2077 left_plane, right_plane = longitudinal_planes
2095 longitudinal_mask = points_inside_longitudinal_limits(2078 longitudinal_mask = points_inside_longitudinal_limits(
2096 points_to_save,2079 points_to_save,
Importance #9: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -2111,16 +2094,12 @@
2111 )2094 )
2112 if not np.any(longitudinal_mask):2095 if not np.any(longitudinal_mask):
2113 continue2096 continue
2114 points_to_save = points_to_save[longitudinal_mask]2097 points_to_save = points_to_save[longitudinal_mask]
2115 scan_angle_to_save = scan_angle_to_save[longitudinal_mask]2098 fields_to_save = {
2116 intensity_to_save = intensity_to_save[longitudinal_mask]2099 field_name: array[longitudinal_mask]
2117 red_to_save = red_to_save[longitudinal_mask]2100 for field_name, array in fields_to_save.items()
2118 green_to_save = green_to_save[longitudinal_mask]2101 }
2119 blue_to_save = blue_to_save[longitudinal_mask]
2120 number_of_returns_to_save = number_of_returns_to_save[
2121 longitudinal_mask
2122 ]
21232102
2124 point_count = int(points_to_save.shape[0])2103 point_count = int(points_to_save.shape[0])
2125 point_count_by_segment[segment_idx] = (2104 point_count_by_segment[segment_idx] = (
2126 point_count_by_segment.get(segment_idx, 0)2105 point_count_by_segment.get(segment_idx, 0)
Importance #10: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -2129,23 +2108,13 @@
2129 if writer is not None:2108 if writer is not None:
2130 writer.write(2109 writer.write(
2131 segment_idx,2110 segment_idx,
2132 points=points_to_save,2111 points=points_to_save,
2133 scan_angle=scan_angle_to_save,2112 fields=fields_to_save,
2134 intensity=intensity_to_save,
2135 red=red_to_save,
2136 green=green_to_save,
2137 blue=blue_to_save,
2138 number_of_returns=number_of_returns_to_save,
2139 )2113 )
2140 del (2114 del (
2141 points,2115 points,
2142 scan_angle,2116 fields,
2143 intensity,
2144 red,
2145 green,
2146 blue,
2147 number_of_returns,
2148 mask,2117 mask,
2149 seen_non_positive_plane,2118 seen_non_positive_plane,
2150 non_prefix_plane_pattern,2119 non_prefix_plane_pattern,
2151 plane_distance,2120 plane_distance,
Importance #11: src/iolabs_point_cloud_segmentation_trajectory/segment_mapper.py @@ -2557,8 +2526,54 @@
2557 f"{las_file}: expected `red/green/blue`, "2526 f"{las_file}: expected `red/green/blue`, "
2558 f"available={SegmentMapper._chunk_field_names(chunk)}"2527 f"available={SegmentMapper._chunk_field_names(chunk)}"
2559 )2528 )
25602529
2530 @staticmethod
2531 def _chunk_field_arrays(
2532 chunk: Any,
2533 *,
2534 las_file: Path,
2535 point_count: int,
2536 chunk_count: int,
2537 logger: logging.Logger,
2538 ) -> dict[str, np.ndarray]:
2539 """Extract the ancillary per-point arrays of one LAS chunk.
2540
2541 Args:
2542 chunk: A laspy chunk record.
2543 las_file: LAS file the chunk came from, used in messages.
2544 point_count: Number of points in *chunk*, used to size fallbacks.
2545 chunk_count: 1-based chunk number, used in messages.
2546 logger: Logger for field-fallback notices.
2547
2548 Returns:
2549 One array per :data:`ANCILLARY_FIELD_NAMES` entry, in that order.
2550 Adding a schema field means adding an entry below plus one in the
2551 tuple; nothing else in the split path needs to change.
2552 """
2553 # Field-probing order matters for the error raised by a malformed chunk:
2554 # scan angle is validated before RGB, as in the pre-refactor code.
2555 scan_angle = SegmentMapper._chunk_scan_angle(
2556 chunk,
2557 las_file=las_file,
2558 logger=logger,
2559 )
2560 red, green, blue = SegmentMapper._chunk_rgb(chunk=chunk, las_file=las_file)
2561 arrays: dict[str, np.ndarray] = {
2562 "scan_angle": scan_angle,
2563 "intensity": np.asarray(chunk.intensity),
2564 "red": red,
2565 "green": green,
2566 "blue": blue,
2567 NUMBER_OF_RETURNS_KEY: _record_number_of_returns(
2568 chunk,
2569 point_count=point_count,
2570 source=f"{las_file.name} chunk {chunk_count}",
2571 logger=logger,
2572 ),
2573 }
2574 return {field_name: arrays[field_name] for field_name in ANCILLARY_FIELD_NAMES}
2575
2561 def _point_segment(self, point: np.ndarray) -> int:2576 def _point_segment(self, point: np.ndarray) -> int:
2562 """Returns the segment index where the point belongs. If no segment is found, return -1."""2577 """Returns the segment index where the point belongs. If no segment is found, return -1."""
2563 for segment_idx in range(len(self.planes) - 1):2578 for segment_idx in range(len(self.planes) - 1):
2564 plane1 = self.planes[segment_idx]2579 plane1 = self.planes[segment_idx]
Importance #12: tests/test_segment_mapper_overflow.py @@ -88,24 +88,28 @@
88 with writer:88 with writer:
89 writer.write(89 writer.write(
90 3,90 3,
91 points=np.array([[11.0, 22.0, 33.0], [14.0, 25.0, 36.0]]),91 points=np.array([[11.0, 22.0, 33.0], [14.0, 25.0, 36.0]]),
92 scan_angle=np.array([1, 2], dtype=np.int16),92 fields={
93 intensity=np.array([100, 200], dtype=np.uint16),93 "scan_angle": np.array([1, 2], dtype=np.int16),
94 red=np.array([10, 20], dtype=np.uint16),94 "intensity": np.array([100, 200], dtype=np.uint16),
95 green=np.array([30, 40], dtype=np.uint16),95 "red": np.array([10, 20], dtype=np.uint16),
96 blue=np.array([50, 60], dtype=np.uint16),96 "green": np.array([30, 40], dtype=np.uint16),
97 number_of_returns=np.array([1, 2], dtype=np.uint8),97 "blue": np.array([50, 60], dtype=np.uint16),
98 "number_of_returns": np.array([1, 2], dtype=np.uint8),
99 },
98 )100 )
99 writer.write(101 writer.write(
100 3,102 3,
101 points=np.array([[17.0, 28.0, 39.0]]),103 points=np.array([[17.0, 28.0, 39.0]]),
102 scan_angle=np.array([3], dtype=np.int16),104 fields={
103 intensity=np.array([300], dtype=np.uint16),105 "scan_angle": np.array([3], dtype=np.int16),
104 red=np.array([30], dtype=np.uint16),106 "intensity": np.array([300], dtype=np.uint16),
105 green=np.array([50], dtype=np.uint16),107 "red": np.array([30], dtype=np.uint16),
106 blue=np.array([70], dtype=np.uint16),108 "green": np.array([50], dtype=np.uint16),
107 number_of_returns=np.array([3], dtype=np.uint8),109 "blue": np.array([70], dtype=np.uint16),
110 "number_of_returns": np.array([3], dtype=np.uint8),
111 },
108 )112 )
109 writer.finalize()113 writer.finalize()
110114
111 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"115 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"
Importance #13: tests/test_segment_mapper_overflow.py @@ -88,24 +88,28 @@
88 with writer:88 with writer:
89 writer.write(89 writer.write(
90 3,90 3,
91 points=np.array([[11.0, 22.0, 33.0], [14.0, 25.0, 36.0]]),91 points=np.array([[11.0, 22.0, 33.0], [14.0, 25.0, 36.0]]),
92 scan_angle=np.array([1, 2], dtype=np.int16),92 fields={
93 intensity=np.array([100, 200], dtype=np.uint16),93 "scan_angle": np.array([1, 2], dtype=np.int16),
94 red=np.array([10, 20], dtype=np.uint16),94 "intensity": np.array([100, 200], dtype=np.uint16),
95 green=np.array([30, 40], dtype=np.uint16),95 "red": np.array([10, 20], dtype=np.uint16),
96 blue=np.array([50, 60], dtype=np.uint16),96 "green": np.array([30, 40], dtype=np.uint16),
97 number_of_returns=np.array([1, 2], dtype=np.uint8),97 "blue": np.array([50, 60], dtype=np.uint16),
98 "number_of_returns": np.array([1, 2], dtype=np.uint8),
99 },
98 )100 )
99 writer.write(101 writer.write(
100 3,102 3,
101 points=np.array([[17.0, 28.0, 39.0]]),103 points=np.array([[17.0, 28.0, 39.0]]),
102 scan_angle=np.array([3], dtype=np.int16),104 fields={
103 intensity=np.array([300], dtype=np.uint16),105 "scan_angle": np.array([3], dtype=np.int16),
104 red=np.array([30], dtype=np.uint16),106 "intensity": np.array([300], dtype=np.uint16),
105 green=np.array([50], dtype=np.uint16),107 "red": np.array([30], dtype=np.uint16),
106 blue=np.array([70], dtype=np.uint16),108 "green": np.array([50], dtype=np.uint16),
107 number_of_returns=np.array([3], dtype=np.uint8),109 "blue": np.array([70], dtype=np.uint16),
110 "number_of_returns": np.array([3], dtype=np.uint8),
111 },
108 )112 )
109 writer.finalize()113 writer.finalize()
110114
111 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"115 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"