Miroslav Simko <ms@iolabs.ch> 2026-09-01T10:48:35+02:00
Commit #94 ยท 13 snippets
.../segment_mapper.py | 87 +++++++++++++++++++++- 1 file changed, 86 insertions(+), 1 deletion(-)
| 30 | 30 | ||
| 31 | OVERFLOW_RETRY_POINT_THRESHOLD = 100 | 31 | OVERFLOW_RETRY_POINT_THRESHOLD = 100 |
| 32 | MAX_OVERFLOW_RETRIES = 3 | 32 | MAX_OVERFLOW_RETRIES = 3 |
| 33 | 33 | ||
| 34 | #: Per-point LAS return-count member of the run3 NPZ contract (AI3D-382). | ||
| 35 | #: The schema SSOT is `iolabs.common.segment_points_io`; the name and dtype are | ||
| 36 | #: mirrored here so this module stays importable against any released | ||
| 37 | #: `iolabs-common`. | ||
| 38 | NUMBER_OF_RETURNS_KEY = "number_of_returns" | ||
| 39 | |||
| 40 | #: LAS carries the return count in 3 bits (valid values 1-7), so uint8 stores it | ||
| 41 | #: 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. | ||
| 43 | NUMBER_OF_RETURNS_DTYPE = np.uint8 | ||
| 44 | |||
| 45 | |||
| 46 | def _record_number_of_returns( | ||
| 47 | record: Any, | ||
| 48 | *, | ||
| 49 | point_count: int, | ||
| 50 | source: str, | ||
| 51 | logger: logging.Logger | None = None, | ||
| 52 | ) -> np.ndarray: | ||
| 53 | """Return per-point LAS return counts as uint8, zero-filled when absent. | ||
| 54 | |||
| 55 | Mirrors the scan-angle field fallback: a LAS record that does not expose | ||
| 56 | `number_of_returns` degrades to zeros rather than raising, and 0 is read | ||
| 57 | downstream as "unknown" (see :data:`NUMBER_OF_RETURNS_DTYPE`). | ||
| 58 | |||
| 59 | Args: | ||
| 60 | record: A laspy point record (whole-file `LasData` or a chunk). | ||
| 61 | point_count: Number of points in *record*, used to size the fallback. | ||
| 62 | source: Label used in the fallback log message (typically a file name). | ||
| 63 | logger: Optional logger for the fallback notice. | ||
| 64 | |||
| 65 | Returns: | ||
| 66 | A uint8 array of shape `(point_count,)`. | ||
| 67 | """ | ||
| 68 | values = getattr(record, NUMBER_OF_RETURNS_KEY, None) | ||
| 69 | if values is None: | ||
| 70 | if logger is not None: | ||
| 71 | logger.debug( | ||
| 72 | "No `%s` field for %s; filling %d zeros (unknown).", | ||
| 73 | NUMBER_OF_RETURNS_KEY, | ||
| 74 | source, | ||
| 75 | point_count, | ||
| 76 | ) | ||
| 77 | return np.zeros(point_count, dtype=NUMBER_OF_RETURNS_DTYPE) | ||
| 78 | return np.asarray(values).astype(NUMBER_OF_RETURNS_DTYPE, copy=False) | ||
| 79 | |||
| 34 | 80 | ||
| 35 | def _load_geoshift_from_json(geoshift_path: Path) -> np.ndarray: | 81 | def _load_geoshift_from_json(geoshift_path: Path) -> np.ndarray: |
| 36 | if not geoshift_path.exists(): | 82 | if not geoshift_path.exists(): |
| 37 | raise FileNotFoundError( | 83 | raise FileNotFoundError( |
| 262 | intensity: np.ndarray, | 308 | intensity: np.ndarray, |
| 263 | red: np.ndarray, | 309 | red: np.ndarray, |
| 264 | green: np.ndarray, | 310 | green: np.ndarray, |
| 265 | blue: np.ndarray, | 311 | blue: np.ndarray, |
| 312 | number_of_returns: np.ndarray, | ||
| 266 | ) -> None: | 313 | ) -> None: |
| 267 | point_count = int(points.shape[0]) | 314 | point_count = int(points.shape[0]) |
| 268 | if point_count == 0: | 315 | if point_count == 0: |
| 269 | return | 316 | return |
| 285 | arrays["intensity"][start:end] = intensity | 332 | arrays["intensity"][start:end] = intensity |
| 286 | arrays["red"][start:end] = red | 333 | arrays["red"][start:end] = red |
| 287 | arrays["green"][start:end] = green | 334 | arrays["green"][start:end] = green |
| 288 | arrays["blue"][start:end] = blue | 335 | arrays["blue"][start:end] = blue |
| 336 | arrays[NUMBER_OF_RETURNS_KEY][start:end] = number_of_returns | ||
| 289 | self._offsets[segment_idx] = end | 337 | self._offsets[segment_idx] = end |
| 290 | 338 | ||
| 291 | def finalize(self) -> None: | 339 | def finalize(self) -> None: |
| 292 | for segment_idx, expected_count in sorted(self.point_count_by_segment.items()): | 340 | for segment_idx, expected_count in sorted(self.point_count_by_segment.items()): |
| 328 | "intensity", | 376 | "intensity", |
| 329 | "red", | 377 | "red", |
| 330 | "green", | 378 | "green", |
| 331 | "blue", | 379 | "blue", |
| 380 | NUMBER_OF_RETURNS_KEY, | ||
| 332 | ): | 381 | ): |
| 333 | archive.write( | 382 | archive.write( |
| 334 | segment_tmp_dir / f"{field_name}.npy", | 383 | segment_tmp_dir / f"{field_name}.npy", |
| 335 | arcname=f"{field_name}.npy", | 384 | arcname=f"{field_name}.npy", |
| 385 | mode="w+", | 434 | mode="w+", |
| 386 | dtype=self.field_dtypes["blue"], | 435 | dtype=self.field_dtypes["blue"], |
| 387 | shape=(point_count,), | 436 | shape=(point_count,), |
| 388 | ), | 437 | ), |
| 438 | NUMBER_OF_RETURNS_KEY: np.lib.format.open_memmap( | ||
| 439 | segment_tmp_dir / f"{NUMBER_OF_RETURNS_KEY}.npy", | ||
| 440 | mode="w+", | ||
| 441 | dtype=self.field_dtypes.get( | ||
| 442 | NUMBER_OF_RETURNS_KEY, | ||
| 443 | np.dtype(NUMBER_OF_RETURNS_DTYPE), | ||
| 444 | ), | ||
| 445 | shape=(point_count,), | ||
| 446 | ), | ||
| 389 | } | 447 | } |
| 390 | self._arrays[segment_idx] = arrays | 448 | self._arrays[segment_idx] = arrays |
| 391 | return arrays | 449 | return arrays |
| 392 | 450 |
| 1895 | red, green, blue = self._chunk_rgb( | 1953 | red, green, blue = self._chunk_rgb( |
| 1896 | chunk=chunk, | 1954 | chunk=chunk, |
| 1897 | las_file=las_file, | 1955 | las_file=las_file, |
| 1898 | ) | 1956 | ) |
| 1957 | number_of_returns = _record_number_of_returns( | ||
| 1958 | chunk, | ||
| 1959 | point_count=int(points.shape[0]), | ||
| 1960 | source=f"{las_file.name} chunk {chunk_count}", | ||
| 1961 | logger=logger, | ||
| 1962 | ) | ||
| 1899 | 1963 | ||
| 1900 | processed_points += int(points.shape[0]) | 1964 | processed_points += int(points.shape[0]) |
| 1901 | 1965 | ||
| 1902 | if angle_limit is not None: | 1966 | if angle_limit is not None: |
| 1908 | intensity = intensity[angle_mask] | 1972 | intensity = intensity[angle_mask] |
| 1909 | red = red[angle_mask] | 1973 | red = red[angle_mask] |
| 1910 | green = green[angle_mask] | 1974 | green = green[angle_mask] |
| 1911 | blue = blue[angle_mask] | 1975 | blue = blue[angle_mask] |
| 1976 | number_of_returns = number_of_returns[angle_mask] | ||
| 1912 | 1977 | ||
| 1913 | if points.shape[0] == 0: | 1978 | if points.shape[0] == 0: |
| 1914 | continue | 1979 | continue |
| 1915 | 1980 |
| 1919 | field_dtypes.setdefault("intensity", np.dtype(intensity.dtype)) | 1984 | field_dtypes.setdefault("intensity", np.dtype(intensity.dtype)) |
| 1920 | field_dtypes.setdefault("red", np.dtype(red.dtype)) | 1985 | field_dtypes.setdefault("red", np.dtype(red.dtype)) |
| 1921 | field_dtypes.setdefault("green", np.dtype(green.dtype)) | 1986 | field_dtypes.setdefault("green", np.dtype(green.dtype)) |
| 1922 | field_dtypes.setdefault("blue", np.dtype(blue.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 | ) | ||
| 1923 | 1992 | ||
| 1924 | logger.debug( | 1993 | logger.debug( |
| 1925 | "Chunk %d: processing %d filtered points", | 1994 | "Chunk %d: processing %d filtered points", |
| 1926 | chunk_count, | 1995 | chunk_count, |
| 2019 | intensity_to_save = intensity[points_mask] | 2088 | intensity_to_save = intensity[points_mask] |
| 2020 | red_to_save = red[points_mask] | 2089 | red_to_save = red[points_mask] |
| 2021 | green_to_save = green[points_mask] | 2090 | green_to_save = green[points_mask] |
| 2022 | blue_to_save = blue[points_mask] | 2091 | blue_to_save = blue[points_mask] |
| 2092 | number_of_returns_to_save = number_of_returns[points_mask] | ||
| 2023 | if longitudinal_planes is not None: | 2093 | if longitudinal_planes is not None: |
| 2024 | left_plane, right_plane = longitudinal_planes | 2094 | left_plane, right_plane = longitudinal_planes |
| 2025 | longitudinal_mask = points_inside_longitudinal_limits( | 2095 | longitudinal_mask = points_inside_longitudinal_limits( |
| 2026 | points_to_save, | 2096 | points_to_save, |
| 2046 | intensity_to_save = intensity_to_save[longitudinal_mask] | 2116 | intensity_to_save = intensity_to_save[longitudinal_mask] |
| 2047 | red_to_save = red_to_save[longitudinal_mask] | 2117 | red_to_save = red_to_save[longitudinal_mask] |
| 2048 | green_to_save = green_to_save[longitudinal_mask] | 2118 | green_to_save = green_to_save[longitudinal_mask] |
| 2049 | blue_to_save = blue_to_save[longitudinal_mask] | 2119 | blue_to_save = blue_to_save[longitudinal_mask] |
| 2120 | number_of_returns_to_save = number_of_returns_to_save[ | ||
| 2121 | longitudinal_mask | ||
| 2122 | ] | ||
| 2050 | 2123 | ||
| 2051 | point_count = int(points_to_save.shape[0]) | 2124 | point_count = int(points_to_save.shape[0]) |
| 2052 | point_count_by_segment[segment_idx] = ( | 2125 | point_count_by_segment[segment_idx] = ( |
| 2053 | point_count_by_segment.get(segment_idx, 0) | 2126 | point_count_by_segment.get(segment_idx, 0) |
| 2061 | intensity=intensity_to_save, | 2134 | intensity=intensity_to_save, |
| 2062 | red=red_to_save, | 2135 | red=red_to_save, |
| 2063 | green=green_to_save, | 2136 | green=green_to_save, |
| 2064 | blue=blue_to_save, | 2137 | blue=blue_to_save, |
| 2138 | number_of_returns=number_of_returns_to_save, | ||
| 2065 | ) | 2139 | ) |
| 2066 | del ( | 2140 | del ( |
| 2067 | points, | 2141 | points, |
| 2068 | scan_angle, | 2142 | scan_angle, |
| 2069 | intensity, | 2143 | intensity, |
| 2070 | red, | 2144 | red, |
| 2071 | green, | 2145 | green, |
| 2072 | blue, | 2146 | blue, |
| 2147 | number_of_returns, | ||
| 2073 | mask, | 2148 | mask, |
| 2074 | seen_non_positive_plane, | 2149 | seen_non_positive_plane, |
| 2075 | non_prefix_plane_pattern, | 2150 | non_prefix_plane_pattern, |
| 2076 | plane_distance, | 2151 | plane_distance, |
| 2497 | return -1 | 2572 | return -1 |
| 2498 | 2573 | ||
| 2499 | @staticmethod | 2574 | @staticmethod |
| 2500 | def load_color_intensity_data(las: laspy.LasData) -> color_intensity_data.ColorIntensityData: | 2575 | def load_color_intensity_data(las: laspy.LasData) -> color_intensity_data.ColorIntensityData: |
| 2501 | """Loads the color and intensity data from the LAS file.""" | 2576 | """Loads the color, intensity, and return-count data from the LAS file. |
| 2577 | |||
| 2578 | A LAS without a `number_of_returns` field degrades to zeros (unknown) | ||
| 2579 | instead of raising, mirroring the chunked read path. | ||
| 2580 | """ | ||
| 2502 | if not (hasattr(las, "red") and hasattr(las, "green") and hasattr(las, "blue")): | 2581 | if not (hasattr(las, "red") and hasattr(las, "green") and hasattr(las, "blue")): |
| 2503 | raise ValueError("No color information found in the LAS file") | 2582 | raise ValueError("No color information found in the LAS file") |
| 2504 | if not hasattr(las, "intensity"): | 2583 | if not hasattr(las, "intensity"): |
| 2505 | raise ValueError("No intensity information found in the LAS file") | 2584 | raise ValueError("No intensity information found in the LAS file") |
| 2513 | intensity = np.asarray(las.intensity) | 2592 | intensity = np.asarray(las.intensity) |
| 2514 | red = np.asarray(las.red) | 2593 | red = np.asarray(las.red) |
| 2515 | green = np.asarray(las.green) | 2594 | green = np.asarray(las.green) |
| 2516 | blue = np.asarray(las.blue) | 2595 | blue = np.asarray(las.blue) |
| 2596 | number_of_returns = _record_number_of_returns( | ||
| 2597 | las, | ||
| 2598 | point_count=int(red.shape[0]), | ||
| 2599 | source="LAS file", | ||
| 2600 | ) | ||
| 2517 | 2601 | ||
| 2518 | return color_intensity_data.ColorIntensityData( | 2602 | return color_intensity_data.ColorIntensityData( |
| 2519 | red=red, | 2603 | red=red, |
| 2520 | green=green, | 2604 | green=green, |
| 2521 | blue=blue, | 2605 | blue=blue, |
| 2522 | intensity=intensity, | 2606 | intensity=intensity, |
| 2523 | scan_angle_rank=scan_angle, | 2607 | scan_angle_rank=scan_angle, |
| 2608 | number_of_returns=number_of_returns, | ||
| 2524 | ) | 2609 | ) |