Back to report index

Step 3 segmentationtrajectory d09a6ce: AI3D-382 Cover number_of_returns write path and zero fallback in tests

Miroslav Simko <ms@iolabs.ch> 2026-09-01T10:48:35+02:00

Commit #21 ยท 11 snippets

 tests/test_number_of_returns.py       | 116 ++++++++++++++++++++++++++++++++++
 tests/test_segment_mapper_overflow.py |   8 +++
 2 files changed, 124 insertions(+)

Tests for the producer write path and the zero fallback, plus fixture updates for the extra key.

Importance #1: tests/test_number_of_returns.py @@ -0,0 +1,116 @@
1import logging
2from pathlib import Path
3
4import numpy as np
5import pytest
6
7from iolabs_point_cloud_segmentation_trajectory import segment_mapper as sm
8
9DIVISION_PLANES = [
10 (np.array([0.0, 0.0, 0.0]), np.array([1.0, 0.0, 0.0])),
11 (np.array([100.0, 0.0, 0.0]), np.array([1.0, 0.0, 0.0])),
12]
13
14
15class _Chunk:
16 def __init__(self, number_of_returns: np.ndarray | None) -> None:
17 self.x = np.array([10.0, 20.0, 30.0], dtype=np.float64)
18 self.y = np.zeros(3, dtype=np.float64)
19 self.z = np.zeros(3, dtype=np.float64)
20 self.scan_angle_rank = np.zeros(3, dtype=np.int16)
21 self.intensity = np.array([10, 20, 30], dtype=np.uint16)
22 self.red = np.array([1, 2, 3], dtype=np.uint16)
23 self.green = np.array([4, 5, 6], dtype=np.uint16)
24 self.blue = np.array([7, 8, 9], dtype=np.uint16)
25 if number_of_returns is not None:
26 self.number_of_returns = number_of_returns
27
28
29class _LasReader:
30 def __init__(self, number_of_returns: np.ndarray | None) -> None:
31 self.header = type("Header", (), {"point_count": 3})()
32 self._number_of_returns = number_of_returns
33
34 def __enter__(self) -> "_LasReader":
35 return self
36
37 def __exit__(self, exc_type, exc, tb) -> None:
38 return None
39
40 def chunk_iterator(self, points_per_chunk: int):
41 yield _Chunk(self._number_of_returns)
42
43
44def _split_one_las(
45 monkeypatch: pytest.MonkeyPatch,
46 tmp_path: Path,
47 number_of_returns: np.ndarray | None,
48) -> dict[str, np.ndarray]:
49 monkeypatch.setattr(
50 sm.laspy,
51 "open",
52 lambda _path: _LasReader(number_of_returns),
53 )
54 mapper = sm.SegmentMapper.__new__(sm.SegmentMapper)
55 mapper.geoshift = np.zeros(3, dtype=np.float64)
56 logger = logging.getLogger("test")
57 las_file = Path("Record001.las")
58
59 result = mapper._collect_las_split_attempt(
60 las_file=las_file,
61 division_planes_np=DIVISION_PLANES,
62 starting_plane_idx=0,
63 angle_limit=None,
64 las_points_per_chunk=100,
65 logger=logger,
66 )
67 assert result.point_count_by_segment == {0: 3}
68
69 mapper._write_las_split_outputs(
70 las_file=las_file,
71 segments_base_dir=tmp_path,
72 points_suffix="_run3_points",
73 point_count_by_segment=result.point_count_by_segment,
74 field_dtypes=result.field_dtypes,
75 division_planes_np=DIVISION_PLANES,
76 starting_plane_idx=0,
77 angle_limit=None,
78 las_points_per_chunk=100,
79 logger=logger,
80 )
81
82 output = tmp_path / "segment_000" / "Record001_run3_points.npz"
83 assert output.exists()
84 with np.load(output) as data:
85 return {key: np.asarray(data[key]) for key in data.files}
86
87
88def test_number_of_returns_is_written_per_point(
89 monkeypatch: pytest.MonkeyPatch,
90 tmp_path: Path,
91) -> None:
92 record = _split_one_las(
93 monkeypatch,
94 tmp_path,
95 np.array([1, 2, 3], dtype=np.uint8),
96 )
97
98 assert "number_of_returns" in record
99 assert record["number_of_returns"].dtype == np.uint8
100 np.testing.assert_array_equal(
101 record["number_of_returns"],
102 np.array([1, 2, 3], dtype=np.uint8),
103 )
104
105
106def test_number_of_returns_falls_back_to_zeros_when_las_lacks_field(
107 monkeypatch: pytest.MonkeyPatch,
108 tmp_path: Path,
109) -> None:
110 record = _split_one_las(monkeypatch, tmp_path, None)
111
112 assert record["number_of_returns"].dtype == np.uint8
113 np.testing.assert_array_equal(
114 record["number_of_returns"],
115 np.zeros(3, dtype=np.uint8),
116 )
0
Importance #2: tests/test_segment_mapper_overflow.py @@ -80,8 +80,9 @@
80 "intensity": np.dtype(np.uint16),80 "intensity": np.dtype(np.uint16),
81 "red": np.dtype(np.uint16),81 "red": np.dtype(np.uint16),
82 "green": np.dtype(np.uint16),82 "green": np.dtype(np.uint16),
83 "blue": np.dtype(np.uint16),83 "blue": np.dtype(np.uint16),
84 "number_of_returns": np.dtype(np.uint8),
84 },85 },
85 )86 )
8687
87 with writer:88 with writer:
Importance #3: tests/test_segment_mapper_overflow.py @@ -92,8 +93,9 @@
92 intensity=np.array([100, 200], dtype=np.uint16),93 intensity=np.array([100, 200], dtype=np.uint16),
93 red=np.array([10, 20], dtype=np.uint16),94 red=np.array([10, 20], dtype=np.uint16),
94 green=np.array([30, 40], dtype=np.uint16),95 green=np.array([30, 40], dtype=np.uint16),
95 blue=np.array([50, 60], dtype=np.uint16),96 blue=np.array([50, 60], dtype=np.uint16),
97 number_of_returns=np.array([1, 2], dtype=np.uint8),
96 )98 )
97 writer.write(99 writer.write(
98 3,100 3,
99 points=np.array([[17.0, 28.0, 39.0]]),101 points=np.array([[17.0, 28.0, 39.0]]),
Importance #4: tests/test_segment_mapper_overflow.py @@ -101,8 +103,9 @@
101 intensity=np.array([300], dtype=np.uint16),103 intensity=np.array([300], dtype=np.uint16),
102 red=np.array([30], dtype=np.uint16),104 red=np.array([30], dtype=np.uint16),
103 green=np.array([50], dtype=np.uint16),105 green=np.array([50], dtype=np.uint16),
104 blue=np.array([70], dtype=np.uint16),106 blue=np.array([70], dtype=np.uint16),
107 number_of_returns=np.array([3], dtype=np.uint8),
105 )108 )
106 writer.finalize()109 writer.finalize()
107110
108 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"111 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"
Importance #5: tests/test_segment_mapper_overflow.py @@ -114,8 +117,9 @@
114 "intensity",117 "intensity",
115 "red",118 "red",
116 "green",119 "green",
117 "blue",120 "blue",
121 "number_of_returns",
118 }122 }
119 np.testing.assert_array_equal(123 np.testing.assert_array_equal(
120 data["points"],124 data["points"],
121 np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]),125 np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]),
Importance #6: tests/test_segment_mapper_overflow.py @@ -124,8 +128,12 @@
124 np.testing.assert_array_equal(data["intensity"], np.array([100, 200, 300], dtype=np.uint16))128 np.testing.assert_array_equal(data["intensity"], np.array([100, 200, 300], dtype=np.uint16))
125 np.testing.assert_array_equal(data["red"], np.array([10, 20, 30], dtype=np.uint16))129 np.testing.assert_array_equal(data["red"], np.array([10, 20, 30], dtype=np.uint16))
126 np.testing.assert_array_equal(data["green"], np.array([30, 40, 50], dtype=np.uint16))130 np.testing.assert_array_equal(data["green"], np.array([30, 40, 50], dtype=np.uint16))
127 np.testing.assert_array_equal(data["blue"], np.array([50, 60, 70], dtype=np.uint16))131 np.testing.assert_array_equal(data["blue"], np.array([50, 60, 70], dtype=np.uint16))
132 np.testing.assert_array_equal(
133 data["number_of_returns"],
134 np.array([1, 2, 3], dtype=np.uint8),
135 )
128136
129 temp_dirs = list((tmp_path / "lane_points").glob(".Record001_split_*"))137 temp_dirs = list((tmp_path / "lane_points").glob(".Record001_split_*"))
130 assert temp_dirs == []138 assert temp_dirs == []
131139
Importance #7: tests/test_segment_mapper_overflow.py @@ -80,8 +80,9 @@
80 "intensity": np.dtype(np.uint16),80 "intensity": np.dtype(np.uint16),
81 "red": np.dtype(np.uint16),81 "red": np.dtype(np.uint16),
82 "green": np.dtype(np.uint16),82 "green": np.dtype(np.uint16),
83 "blue": np.dtype(np.uint16),83 "blue": np.dtype(np.uint16),
84 "number_of_returns": np.dtype(np.uint8),
84 },85 },
85 )86 )
8687
87 with writer:88 with writer:
Importance #8: tests/test_segment_mapper_overflow.py @@ -92,8 +93,9 @@
92 intensity=np.array([100, 200], dtype=np.uint16),93 intensity=np.array([100, 200], dtype=np.uint16),
93 red=np.array([10, 20], dtype=np.uint16),94 red=np.array([10, 20], dtype=np.uint16),
94 green=np.array([30, 40], dtype=np.uint16),95 green=np.array([30, 40], dtype=np.uint16),
95 blue=np.array([50, 60], dtype=np.uint16),96 blue=np.array([50, 60], dtype=np.uint16),
97 number_of_returns=np.array([1, 2], dtype=np.uint8),
96 )98 )
97 writer.write(99 writer.write(
98 3,100 3,
99 points=np.array([[17.0, 28.0, 39.0]]),101 points=np.array([[17.0, 28.0, 39.0]]),
Importance #9: tests/test_segment_mapper_overflow.py @@ -101,8 +103,9 @@
101 intensity=np.array([300], dtype=np.uint16),103 intensity=np.array([300], dtype=np.uint16),
102 red=np.array([30], dtype=np.uint16),104 red=np.array([30], dtype=np.uint16),
103 green=np.array([50], dtype=np.uint16),105 green=np.array([50], dtype=np.uint16),
104 blue=np.array([70], dtype=np.uint16),106 blue=np.array([70], dtype=np.uint16),
107 number_of_returns=np.array([3], dtype=np.uint8),
105 )108 )
106 writer.finalize()109 writer.finalize()
107110
108 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"111 output = tmp_path / "lane_points" / "segment_003" / "Record001_run3_points.npz"
Importance #10: tests/test_segment_mapper_overflow.py @@ -114,8 +117,9 @@
114 "intensity",117 "intensity",
115 "red",118 "red",
116 "green",119 "green",
117 "blue",120 "blue",
121 "number_of_returns",
118 }122 }
119 np.testing.assert_array_equal(123 np.testing.assert_array_equal(
120 data["points"],124 data["points"],
121 np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]),125 np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]),
Importance #11: tests/test_segment_mapper_overflow.py @@ -124,8 +128,12 @@
124 np.testing.assert_array_equal(data["intensity"], np.array([100, 200, 300], dtype=np.uint16))128 np.testing.assert_array_equal(data["intensity"], np.array([100, 200, 300], dtype=np.uint16))
125 np.testing.assert_array_equal(data["red"], np.array([10, 20, 30], dtype=np.uint16))129 np.testing.assert_array_equal(data["red"], np.array([10, 20, 30], dtype=np.uint16))
126 np.testing.assert_array_equal(data["green"], np.array([30, 40, 50], dtype=np.uint16))130 np.testing.assert_array_equal(data["green"], np.array([30, 40, 50], dtype=np.uint16))
127 np.testing.assert_array_equal(data["blue"], np.array([50, 60, 70], dtype=np.uint16))131 np.testing.assert_array_equal(data["blue"], np.array([50, 60, 70], dtype=np.uint16))
132 np.testing.assert_array_equal(
133 data["number_of_returns"],
134 np.array([1, 2, 3], dtype=np.uint8),
135 )
128136
129 temp_dirs = list((tmp_path / "lane_points").glob(".Record001_split_*"))137 temp_dirs = list((tmp_path / "lane_points").glob(".Record001_split_*"))
130 assert temp_dirs == []138 assert temp_dirs == []
131139