solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4"""Tests for observe.aruco ArUco marker detection and masking."""
5
6import cv2
7import numpy as np
8from PIL import Image
9
10from solstone.observe.aruco import (
11 CORNER_TAG_IDS,
12 detect_markers,
13 mask_convey_region,
14 polygon_area,
15)
16
17
18def test_corner_tag_ids():
19 """Test that corner tag IDs match expected values."""
20 assert CORNER_TAG_IDS == {2, 4, 6, 7}
21
22
23def test_polygon_area_square():
24 """Test polygon area calculation for a square."""
25 # 100x100 square
26 polygon = [(0, 0), (100, 0), (100, 100), (0, 100)]
27 assert polygon_area(polygon) == 10000.0
28
29
30def test_polygon_area_triangle():
31 """Test polygon area calculation for a triangle."""
32 # Right triangle with legs 10 and 20
33 polygon = [(0, 0), (10, 0), (0, 20)]
34 assert polygon_area(polygon) == 100.0 # (10 * 20) / 2
35
36
37def test_polygon_area_empty():
38 """Test polygon area with insufficient points."""
39 assert polygon_area([]) == 0.0
40 assert polygon_area([(0, 0)]) == 0.0
41 assert polygon_area([(0, 0), (1, 1)]) == 0.0
42
43
44def test_detect_markers_no_markers():
45 """Test detect_markers returns None when no markers are present."""
46 img = Image.new("RGB", (640, 480), color="white")
47 result = detect_markers(img)
48 assert result is None
49
50
51def test_detect_markers_grayscale():
52 """Test detect_markers works with grayscale input."""
53 img = Image.new("L", (640, 480), color=128)
54 result = detect_markers(img)
55 assert result is None # No markers, but shouldn't crash
56
57
58def test_mask_convey_region():
59 """Test masking fills polygon with black."""
60 img = Image.new("RGB", (100, 100), color="white")
61
62 # Define a square polygon in the center
63 polygon = [(25, 25), (75, 25), (75, 75), (25, 75)]
64 mask_convey_region(img, polygon)
65
66 # Check corners are still white
67 assert img.getpixel((0, 0)) == (255, 255, 255)
68 assert img.getpixel((99, 99)) == (255, 255, 255)
69
70 # Check center is black
71 assert img.getpixel((50, 50)) == (0, 0, 0)
72
73
74def test_mask_convey_region_triangle():
75 """Test masking works with non-rectangular polygon."""
76 img = Image.new("RGB", (100, 100), color="white")
77
78 # Triangle
79 polygon = [(50, 10), (90, 90), (10, 90)]
80 mask_convey_region(img, polygon)
81
82 # Center should be black (inside triangle)
83 assert img.getpixel((50, 60)) == (0, 0, 0)
84
85 # Top corners should still be white (outside triangle)
86 assert img.getpixel((5, 5)) == (255, 255, 255)
87 assert img.getpixel((95, 5)) == (255, 255, 255)
88
89
90def test_detect_markers_with_all_corners():
91 """Test detect_markers returns full result with all 4 corner markers."""
92 # Create a test image
93 img_array = np.ones((480, 640, 3), dtype=np.uint8) * 255
94
95 # Generate and place the 4 corner markers
96 dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
97 marker_size = 50
98 pad = 20
99
100 # Generate and place markers
101 for tag_id in [6, 7, 4, 2]:
102 marker = cv2.aruco.generateImageMarker(dictionary, tag_id, marker_size)
103 marker_rgb = cv2.cvtColor(marker, cv2.COLOR_GRAY2RGB)
104
105 if tag_id == 6: # TL
106 img_array[pad : pad + marker_size, pad : pad + marker_size] = marker_rgb
107 elif tag_id == 7: # TR
108 img_array[pad : pad + marker_size, 640 - pad - marker_size : 640 - pad] = (
109 marker_rgb
110 )
111 elif tag_id == 4: # BL
112 img_array[480 - pad - marker_size : 480 - pad, pad : pad + marker_size] = (
113 marker_rgb
114 )
115 elif tag_id == 2: # BR
116 img_array[
117 480 - pad - marker_size : 480 - pad,
118 640 - pad - marker_size : 640 - pad,
119 ] = marker_rgb
120
121 pil_img = Image.fromarray(img_array)
122
123 result = detect_markers(pil_img)
124
125 # Should return dict with markers and polygon
126 assert result is not None
127 assert "markers" in result
128 assert "polygon" in result
129
130 # Should have 4 markers
131 assert len(result["markers"]) == 4
132
133 # Each marker should have id and corners
134 marker_ids = {m["id"] for m in result["markers"]}
135 assert marker_ids == {2, 4, 6, 7}
136
137 for marker in result["markers"]:
138 assert "id" in marker
139 assert "corners" in marker
140 assert len(marker["corners"]) == 4
141 for corner in marker["corners"]:
142 assert len(corner) == 2
143 assert isinstance(corner[0], (int, float))
144 assert isinstance(corner[1], (int, float))
145
146 # Polygon should be present (all 4 corners detected)
147 assert result["polygon"] is not None
148 assert len(result["polygon"]) == 4
149
150
151def test_detect_markers_partial():
152 """Test detect_markers returns markers but no polygon with partial detection."""
153 # Create a test image
154 img_array = np.ones((480, 640, 3), dtype=np.uint8) * 255
155
156 # Generate and place only 2 corner markers
157 dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
158 marker_size = 50
159 pad = 20
160
161 # Only place TL (6) and TR (7) markers
162 for tag_id, pos in [(6, (pad, pad)), (7, (pad, 640 - pad - marker_size))]:
163 marker = cv2.aruco.generateImageMarker(dictionary, tag_id, marker_size)
164 marker_rgb = cv2.cvtColor(marker, cv2.COLOR_GRAY2RGB)
165 y, x = pos
166 img_array[y : y + marker_size, x : x + marker_size] = marker_rgb
167
168 pil_img = Image.fromarray(img_array)
169
170 result = detect_markers(pil_img)
171
172 # Should return dict with markers but no polygon
173 assert result is not None
174 assert "markers" in result
175 assert "polygon" in result
176
177 # Should have 2 markers
178 assert len(result["markers"]) == 2
179 marker_ids = {m["id"] for m in result["markers"]}
180 assert marker_ids == {6, 7}
181
182 # Polygon should be None (only 2 of 4 corners)
183 assert result["polygon"] is None
184
185
186def test_detect_markers_extrapolated_tl():
187 """Test detect_markers extrapolates missing TL corner from 3 present markers."""
188 img_array = np.ones((480, 640, 3), dtype=np.uint8) * 255
189 dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
190 marker_size = 50
191 pad = 20
192
193 # Place TR (7), BR (2), BL (4) — omit TL (6)
194 for tag_id, (y, x) in [
195 (7, (pad, 640 - pad - marker_size)),
196 (2, (480 - pad - marker_size, 640 - pad - marker_size)),
197 (4, (480 - pad - marker_size, pad)),
198 ]:
199 marker = cv2.aruco.generateImageMarker(dictionary, tag_id, marker_size)
200 marker_rgb = cv2.cvtColor(marker, cv2.COLOR_GRAY2RGB)
201 img_array[y : y + marker_size, x : x + marker_size] = marker_rgb
202
203 result = detect_markers(Image.fromarray(img_array))
204
205 assert result is not None
206 assert result["polygon"] is not None
207 assert len(result["polygon"]) == 4
208 assert result.get("extrapolated") == 6
209
210 # Extrapolated TL should be within 2px of expected position
211 tl = result["polygon"][0]
212 assert abs(tl[0] - pad) <= 2
213 assert abs(tl[1] - pad) <= 2
214
215
216def test_detect_markers_two_missing_no_extrapolation():
217 """Test detect_markers returns no polygon when only 2 corner tags present."""
218 img_array = np.ones((480, 640, 3), dtype=np.uint8) * 255
219 dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
220 marker_size = 50
221 pad = 20
222
223 # Place only TL (6) and BR (2)
224 for tag_id, (y, x) in [
225 (6, (pad, pad)),
226 (2, (480 - pad - marker_size, 640 - pad - marker_size)),
227 ]:
228 marker = cv2.aruco.generateImageMarker(dictionary, tag_id, marker_size)
229 marker_rgb = cv2.cvtColor(marker, cv2.COLOR_GRAY2RGB)
230 img_array[y : y + marker_size, x : x + marker_size] = marker_rgb
231
232 result = detect_markers(Image.fromarray(img_array))
233
234 assert result is not None
235 assert result["polygon"] is None