Skip to main content

locus_core/test_utils/
mod.rs

1use rand::prelude::*;
2use rand_distr::{Distribution, Normal};
3
4/// Generate a synthetic image containing a single AprilTag or ArUco tag.
5///
6/// This generates a tag with a white quiet zone, placed on a white background,
7/// matching the setup used in Python benchmarks.
8/// rotation (rad), translation (x,y), scaling (pixels).
9#[must_use]
10#[allow(
11    clippy::cast_possible_truncation,
12    clippy::cast_sign_loss,
13    clippy::missing_panics_doc
14)]
15pub fn generate_synthetic_test_image(
16    family: crate::config::TagFamily,
17    id: u16,
18    tag_size: usize,
19    canvas_size: usize,
20    noise_sigma: f32,
21) -> (Vec<u8>, [[f64; 2]; 4]) {
22    let mut data = vec![255u8; canvas_size * canvas_size];
23
24    // Calculate tag position (centered)
25    let margin = (canvas_size - tag_size) / 2;
26    let quiet_zone = tag_size / 5;
27
28    // Draw white quiet zone (optional since background is 255, but adds robustness)
29    for y in margin.saturating_sub(quiet_zone)..(margin + tag_size + quiet_zone).min(canvas_size) {
30        for x in
31            margin.saturating_sub(quiet_zone)..(margin + tag_size + quiet_zone).min(canvas_size)
32        {
33            data[y * canvas_size + x] = 255;
34        }
35    }
36
37    // Generate tag pattern (bits)
38    let decoder = crate::decoder::family_to_decoder(family);
39    let code = decoder.get_code(id).expect("Invalid tag ID for family");
40    let dim = decoder.dimension();
41
42    let cell_size = tag_size / (dim + 2); // dim + 2 for black border
43    let actual_tag_size = cell_size * (dim + 2);
44    let start_x = margin + (tag_size - actual_tag_size) / 2;
45    let start_y = margin + (tag_size - actual_tag_size) / 2;
46
47    // Draw black border
48    for y in 0..(dim + 2) {
49        for x in 0..(dim + 2) {
50            if x == 0 || x == dim + 1 || y == 0 || y == dim + 1 {
51                draw_cell(&mut data, canvas_size, start_x, start_y, x, y, cell_size, 0);
52            } else {
53                let row = y - 1;
54                let col = x - 1;
55                let bit = (code >> (row * dim + col)) & 1;
56                let val = if bit != 0 { 255 } else { 0 };
57                draw_cell(
58                    &mut data,
59                    canvas_size,
60                    start_x,
61                    start_y,
62                    x,
63                    y,
64                    cell_size,
65                    val,
66                );
67            }
68        }
69    }
70
71    if noise_sigma > 0.0 {
72        let mut rng = thread_rng();
73        let normal = Normal::new(0.0, f64::from(noise_sigma)).expect("Invalid noise params");
74
75        for pixel in &mut data {
76            let noise = normal.sample(&mut rng) as i32;
77            let val = (i32::from(*pixel) + noise).clamp(0, 255);
78            *pixel = val as u8;
79        }
80    }
81
82    let gt_corners = [
83        [start_x as f64, start_y as f64],
84        [(start_x + actual_tag_size) as f64, start_y as f64],
85        [
86            (start_x + actual_tag_size) as f64,
87            (start_y + actual_tag_size) as f64,
88        ],
89        [start_x as f64, (start_y + actual_tag_size) as f64],
90    ];
91
92    (data, gt_corners)
93}
94
95#[allow(clippy::too_many_arguments)]
96fn draw_cell(
97    data: &mut [u8],
98    stride: usize,
99    start_x: usize,
100    start_y: usize,
101    cx: usize,
102    cy: usize,
103    size: usize,
104    val: u8,
105) {
106    let px = start_x + cx * size;
107    let py = start_y + cy * size;
108    for y in py..(py + size) {
109        for x in px..(px + size) {
110            data[y * stride + x] = val;
111        }
112    }
113}
114
115/// Compute mean Euclidean distance between detected and ground truth corners.
116/// Handles 4 rotations to find the minimum error.
117#[must_use]
118pub fn compute_corner_error(detected: &[[f64; 2]; 4], ground_truth: &[[f64; 2]; 4]) -> f64 {
119    let mut min_error = f64::MAX;
120
121    // Try all 4 rotations
122    for rot in 0..4 {
123        let mut sum_dist = 0.0;
124        for i in 0..4 {
125            let d = &detected[(i + rot) % 4];
126            let g = &ground_truth[i];
127            let dx = d[0] - g[0];
128            let dy = d[1] - g[1];
129            sum_dist += (dx * dx + dy * dy).sqrt();
130        }
131        let avg_dist = sum_dist / 4.0;
132        if avg_dist < min_error {
133            min_error = avg_dist;
134        }
135    }
136
137    min_error
138}
139
140/// Compute RMSE corner error between detected and ground truth corners.
141///
142/// Unlike `compute_corner_error`, this does NOT try rotations - it assumes corners
143/// are already in the correct order. Use this when comparing against ground truth
144/// datasets with known corner ordering conventions.
145///
146/// Formula: sqrt( (sum of squared distances for all 4 corners) / 4 )
147#[must_use]
148pub fn compute_rmse(detected: &[[f64; 2]; 4], ground_truth: &[[f64; 2]; 4]) -> f64 {
149    let mut sum_sq = 0.0;
150    for i in 0..4 {
151        let dx = detected[i][0] - ground_truth[i][0];
152        let dy = detected[i][1] - ground_truth[i][1];
153        sum_sq += dx * dx + dy * dy;
154    }
155    (sum_sq / 4.0).sqrt()
156}
157
158// ============================================================================
159// ROBUSTNESS TEST UTILITIES
160// ============================================================================
161
162/// Parameters for generating test images with photometric variations.
163#[derive(Clone, Debug)]
164pub struct TestImageParams {
165    /// Tag family to generate.
166    pub family: crate::config::TagFamily,
167    /// Tag ID to generate.
168    pub id: u16,
169    /// Tag size in pixels.
170    pub tag_size: usize,
171    /// Canvas size in pixels.
172    pub canvas_size: usize,
173    /// Gaussian noise standard deviation (0.0 = no noise).
174    pub noise_sigma: f32,
175    /// Brightness offset (-255 to +255).
176    pub brightness_offset: i16,
177    /// Contrast scale (1.0 = no change, 0.5 = reduce, 1.5 = increase).
178    pub contrast_scale: f32,
179}
180
181impl Default for TestImageParams {
182    fn default() -> Self {
183        Self {
184            family: crate::config::TagFamily::AprilTag36h11,
185            id: 0,
186            tag_size: 100,
187            canvas_size: 320,
188            noise_sigma: 0.0,
189            brightness_offset: 0,
190            contrast_scale: 1.0,
191        }
192    }
193}
194
195/// Generate a test image based on the provided parameters.
196/// Includes tag generation, placement, and photometric adjustments.
197#[must_use]
198#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
199pub fn generate_test_image_with_params(params: &TestImageParams) -> (Vec<u8>, [[f64; 2]; 4]) {
200    // First generate base image
201    let (mut data, corners) = generate_synthetic_test_image(
202        params.family,
203        params.id,
204        params.tag_size,
205        params.canvas_size,
206        params.noise_sigma,
207    );
208
209    // Apply brightness and contrast adjustments
210    if params.brightness_offset != 0 || (params.contrast_scale - 1.0).abs() > 0.001 {
211        apply_brightness_contrast(
212            &mut data,
213            i32::from(params.brightness_offset),
214            params.contrast_scale,
215        );
216    }
217
218    (data, corners)
219}
220
221/// Apply brightness and contrast to an image.
222/// `brightness`: -255 to +255
223/// `contrast`: 0.0 to 127.0 (1.0 = no change)
224#[allow(clippy::cast_possible_wrap, clippy::cast_sign_loss)]
225pub fn apply_brightness_contrast(image: &mut [u8], brightness: i32, contrast: f32) {
226    for pixel in image.iter_mut() {
227        let b = f32::from(*pixel);
228        let with_contrast = (b - 128.0) * contrast + 128.0;
229        let with_brightness = with_contrast as i32 + brightness;
230        *pixel = with_brightness.clamp(0, 255) as u8;
231    }
232}
233
234/// Count black pixels in binary data.
235#[must_use]
236#[allow(clippy::naive_bytecount)]
237pub fn count_black_pixels(data: &[u8]) -> usize {
238    data.iter().filter(|&&p| p == 0).count()
239}
240
241/// Check if the tag's outer black border is correctly binarized.
242/// Returns the ratio of correctly black pixels in the 1-cell-wide border (0.0 to 1.0).
243#[must_use]
244#[allow(clippy::cast_sign_loss)]
245pub fn measure_border_integrity(binary: &[u8], width: usize, corners: &[[f64; 2]; 4]) -> f64 {
246    let min_x = corners
247        .iter()
248        .map(|c| c[0])
249        .fold(f64::MAX, f64::min)
250        .max(0.0) as usize;
251    let max_x = corners
252        .iter()
253        .map(|c| c[0])
254        .fold(f64::MIN, f64::max)
255        .max(0.0) as usize;
256    let min_y = corners
257        .iter()
258        .map(|c| c[1])
259        .fold(f64::MAX, f64::min)
260        .max(0.0) as usize;
261    let max_y = corners
262        .iter()
263        .map(|c| c[1])
264        .fold(f64::MIN, f64::max)
265        .max(0.0) as usize;
266
267    let height = binary.len() / width;
268    let min_x = min_x.min(width.saturating_sub(1));
269    let max_x = max_x.min(width.saturating_sub(1));
270    let min_y = min_y.min(height.saturating_sub(1));
271    let max_y = max_y.min(height.saturating_sub(1));
272
273    if max_x <= min_x || max_y <= min_y {
274        return 0.0;
275    }
276
277    let tag_width = max_x - min_x;
278    let tag_height = max_y - min_y;
279
280    // For AprilTag 36h11: 8 cells total, border is 1 cell = 1/8 of tag
281    let cell_size_x = tag_width / 8;
282    let cell_size_y = tag_height / 8;
283
284    if cell_size_x == 0 || cell_size_y == 0 {
285        return 0.0;
286    }
287
288    let mut black_count = 0usize;
289    let mut total_count = 0usize;
290
291    // Top border row
292    for y in min_y..(min_y + cell_size_y).min(max_y) {
293        for x in min_x..=max_x {
294            if y < height && x < width {
295                total_count += 1;
296                if binary[y * width + x] == 0 {
297                    black_count += 1;
298                }
299            }
300        }
301    }
302
303    // Bottom border row
304    let bottom_start = max_y.saturating_sub(cell_size_y);
305    for y in bottom_start..=max_y {
306        for x in min_x..=max_x {
307            if y < height && x < width {
308                total_count += 1;
309                if binary[y * width + x] == 0 {
310                    black_count += 1;
311                }
312            }
313        }
314    }
315
316    // Left border column (excluding corners)
317    for y in (min_y + cell_size_y)..(max_y.saturating_sub(cell_size_y)) {
318        for x in min_x..(min_x + cell_size_x).min(max_x) {
319            if y < height && x < width {
320                total_count += 1;
321                if binary[y * width + x] == 0 {
322                    black_count += 1;
323                }
324            }
325        }
326    }
327
328    // Right border column (excluding corners)
329    let right_start = max_x.saturating_sub(cell_size_x);
330    for y in (min_y + cell_size_y)..(max_y.saturating_sub(cell_size_y)) {
331        for x in right_start..=max_x {
332            if y < height && x < width {
333                total_count += 1;
334                if binary[y * width + x] == 0 {
335                    black_count += 1;
336                }
337            }
338        }
339    }
340
341    if total_count == 0 {
342        0.0
343    } else {
344        black_count as f64 / total_count as f64
345    }
346}
347
348/// Generates a checkered pattern image for benchmarking.
349#[must_use]
350pub fn generate_checkered(width: usize, height: usize) -> Vec<u8> {
351    let mut data = vec![200u8; width * height];
352    for y in (0..height).step_by(16) {
353        for x in (0..width).step_by(16) {
354            if ((x / 16) + (y / 16)) % 2 == 0 {
355                for dy in 0..16 {
356                    if y + dy < height {
357                        let row_off = (y + dy) * width;
358                        for dx in 0..16 {
359                            if x + dx < width {
360                                data[row_off + x + dx] = 50;
361                            }
362                        }
363                    }
364                }
365            }
366        }
367    }
368    data
369}
370
371/// Complex multi-tag scene generation for integration testing.
372#[cfg(any(feature = "extended-tests", feature = "extended-bench"))]
373pub mod scene;
374#[cfg(any(feature = "extended-tests", feature = "extended-bench"))]
375pub use scene::{SceneBuilder, TagPlacement};