locus_core/
batch.rs

1/// The maximum number of candidates in a single batch.
2pub(crate) const MAX_CANDIDATES: usize = 1024;
3
4/// A 2D point with subpixel precision (f32).
5#[derive(Debug, Clone, Copy, Default)]
6#[repr(C)]
7pub struct Point2f {
8    /// X coordinate.
9    pub x: f32,
10    /// Y coordinate.
11    pub y: f32,
12}
13
14/// A 3x3 homography matrix (f32).
15#[derive(Debug, Clone, Copy, Default)]
16#[repr(C, align(32))]
17pub struct Matrix3x3 {
18    /// The matrix elements in row-major or column-major format (internal use).
19    pub data: [f32; 9],
20    /// Padding to ensure 64-byte size (cache line) and alignment for SIMD.
21    pub padding: [f32; 7],
22}
23
24/// A 6D pose representing translation and rotation (unit quaternion).
25#[derive(Debug, Clone, Copy, Default)]
26#[repr(C, align(32))]
27pub struct Pose6D {
28    /// Translation (x, y, z) and Rotation as a unit quaternion (x, y, z, w).
29    pub data: [f32; 7],
30    /// Padding to 32-byte alignment.
31    pub padding: f32,
32}
33
34/// The lifecycle state of a candidate in the detection pipeline.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
36#[repr(u8)]
37pub enum CandidateState {
38    /// No candidate at this index.
39    #[default]
40    Empty = 0,
41    /// Quad vertices have been extracted.
42    Active = 1,
43    /// Failed to decode the marker payload.
44    FailedDecode = 2,
45    /// Successfully decoded and mathematically verified.
46    Valid = 3,
47}
48
49/// A batched state container for fiducial marker detections using a Structure of Arrays (SoA) layout.
50/// This structure is designed for high-performance SIMD processing and zero heap allocations.
51#[repr(C, align(32))]
52pub struct DetectionBatch {
53    /// Flattened array of sub-pixel quad vertices (4 corners per candidate).
54    pub corners: [[Point2f; 4]; MAX_CANDIDATES],
55    /// The 3x3 projection matrices.
56    pub homographies: [Matrix3x3; MAX_CANDIDATES],
57    /// The decoded IDs of the tags.
58    pub ids: [u32; MAX_CANDIDATES],
59    /// The extracted bitstrings.
60    pub payloads: [u64; MAX_CANDIDATES],
61    /// The MSE or Log-Likelihood Ratio confidence scores.
62    pub error_rates: [f32; MAX_CANDIDATES],
63    /// Translation vectors and unit quaternions.
64    pub poses: [Pose6D; MAX_CANDIDATES],
65    /// A dense byte-array tracking the lifecycle of each candidate.
66    pub status_mask: [CandidateState; MAX_CANDIDATES],
67}
68
69impl DetectionBatch {
70    /// Creates a new DetectionBatch with all fields initialized to zero (Empty state).
71    #[must_use]
72    #[allow(clippy::large_stack_arrays)]
73    pub fn new() -> Self {
74        *Box::new(Self {
75            corners: [[Point2f { x: 0.0, y: 0.0 }; 4]; MAX_CANDIDATES],
76            homographies: [Matrix3x3 {
77                data: [0.0; 9],
78                padding: [0.0; 7],
79            }; MAX_CANDIDATES],
80            ids: [0; MAX_CANDIDATES],
81            payloads: [0; MAX_CANDIDATES],
82            error_rates: [0.0; MAX_CANDIDATES],
83            poses: [Pose6D {
84                data: [0.0; 7],
85                padding: 0.0,
86            }; MAX_CANDIDATES],
87            status_mask: [CandidateState::Empty; MAX_CANDIDATES],
88        })
89    }
90    /// Returns the maximum capacity of the batch.
91    #[must_use]
92    pub fn capacity(&self) -> usize {
93        MAX_CANDIDATES
94    }
95
96    /// Partitions the batch so that all `Valid` candidates are at the front `[0..V]`.
97    /// Returns the number of valid candidates `V`.
98    pub fn partition(&mut self, n: usize) -> usize {
99        let mut v = 0;
100        let n_clamped = n.min(MAX_CANDIDATES);
101
102        // Two-pointer partition:
103        // Move Valid candidates to [0..v]
104        // Move everything else to [v..n_clamped]
105        for i in 0..n_clamped {
106            if self.status_mask[i] == CandidateState::Valid {
107                if i != v {
108                    // Swap index i with index v across all parallel arrays.
109                    self.corners.swap(i, v);
110                    self.homographies.swap(i, v);
111                    self.ids.swap(i, v);
112                    self.payloads.swap(i, v);
113                    self.error_rates.swap(i, v);
114                    self.poses.swap(i, v);
115                    self.status_mask.swap(i, v);
116                }
117                v += 1;
118            }
119        }
120        v
121    }
122
123    /// Reassemble the batched SoA data into a list of discrete `Detection` objects.
124    #[must_use]
125    #[allow(clippy::cast_sign_loss)]
126    pub fn reassemble(&self, v: usize) -> Vec<crate::Detection> {
127        let mut detections = Vec::with_capacity(v);
128        for i in 0..v {
129            let corners = [
130                [
131                    f64::from(self.corners[i][0].x),
132                    f64::from(self.corners[i][0].y),
133                ],
134                [
135                    f64::from(self.corners[i][1].x),
136                    f64::from(self.corners[i][1].y),
137                ],
138                [
139                    f64::from(self.corners[i][2].x),
140                    f64::from(self.corners[i][2].y),
141                ],
142                [
143                    f64::from(self.corners[i][3].x),
144                    f64::from(self.corners[i][3].y),
145                ],
146            ];
147
148            let center = [
149                (corners[0][0] + corners[1][0] + corners[2][0] + corners[3][0]) / 4.0,
150                (corners[0][1] + corners[1][1] + corners[2][1] + corners[3][1]) / 4.0,
151            ];
152
153            // Reconstruct Pose if available (Z translation > 0)
154            let pose = if self.poses[i].data[2] > 0.0 {
155                let d = self.poses[i].data;
156                // layout: [tx, ty, tz, qx, qy, qz, qw]
157                let t = nalgebra::Vector3::new(f64::from(d[0]), f64::from(d[1]), f64::from(d[2]));
158                let q = nalgebra::UnitQuaternion::from_quaternion(nalgebra::Quaternion::new(
159                    f64::from(d[6]), // w
160                    f64::from(d[3]), // x
161                    f64::from(d[4]), // y
162                    f64::from(d[5]), // z
163                ));
164                Some(crate::pose::Pose {
165                    rotation: q.to_rotation_matrix().into_inner(),
166                    translation: t,
167                })
168            } else {
169                None
170            };
171
172            detections.push(crate::Detection {
173                id: self.ids[i],
174                center,
175                corners,
176                hamming: self.error_rates[i] as u32,
177                rotation: 0,
178                decision_margin: 0.0,
179                bits: self.payloads[i],
180                pose,
181                pose_covariance: None,
182            });
183        }
184        detections
185    }
186}
187
188/// Helper function to partition a batch, moving all valid candidates to the front.
189pub fn partition_batch_soa(batch: &mut DetectionBatch, n: usize) -> usize {
190    batch.partition(n)
191}
192
193/// Metadata for zero-copy telemetry extraction.
194#[derive(Debug, Clone, Copy)]
195pub struct TelemetryPayload {
196    /// Pointer to the binarized image buffer.
197    pub binarized_ptr: *const u8,
198    /// Pointer to the threshold map buffer.
199    pub threshold_map_ptr: *const u8,
200    /// Pointer to subpixel jitter data [4 corners * 2 (dx, dy)] per valid candidate.
201    /// Allocated in the detector's arena.
202    pub subpixel_jitter_ptr: *const f32,
203    /// Number of valid candidates jitter data is available for.
204    pub num_jitter: usize,
205    /// Pointer to reprojection RMSE values per valid candidate.
206    /// Allocated in the detector's arena.
207    pub reprojection_errors_ptr: *const f32,
208    /// Number of valid candidates reprojection data is available for.
209    pub num_reprojection: usize,
210    /// Width of the buffers.
211    pub width: usize,
212    /// Height of the buffers.
213    pub height: usize,
214    /// Stride of the buffers.
215    pub stride: usize,
216}
217
218// SAFETY: TelemetryPayload contains raw pointers to the detector's arena.
219// These pointers are stable for the duration of the frame processing and
220// are only accessed from the Python thread while the arena is still alive.
221#[allow(unsafe_code)]
222unsafe impl Send for TelemetryPayload {}
223#[allow(unsafe_code)]
224unsafe impl Sync for TelemetryPayload {}
225
226/// A lightweight, borrowed view of the detection results.
227///
228/// This struct holds slices to the active elements in a [`DetectionBatch`].
229/// It avoids heap allocations and provides efficient access to detection data.
230#[derive(Debug, Clone, Copy)]
231pub struct DetectionBatchView<'a> {
232    /// Decoded IDs of the markers.
233    pub ids: &'a [u32],
234    /// Refined corners in image coordinates.
235    pub corners: &'a [[Point2f; 4]],
236    /// Computed homography matrices.
237    pub homographies: &'a [Matrix3x3],
238    /// Decoded bitstrings (paylods).
239    pub payloads: &'a [u64],
240    /// Confidence scores or error rates.
241    pub error_rates: &'a [f32],
242    /// 3D poses (rotation + translation).
243    pub poses: &'a [Pose6D],
244    /// Optional telemetry data for intermediate images.
245    pub telemetry: Option<TelemetryPayload>,
246    /// Corners of quads that were extracted but rejected during decoding or verification.
247    pub rejected_corners: &'a [[Point2f; 4]],
248    /// Error rates (e.g. Hamming distance) for rejected quads.
249    pub rejected_error_rates: &'a [f32],
250}
251
252impl DetectionBatchView<'_> {
253    /// Returns the number of detections in the view.
254    #[must_use]
255    pub fn len(&self) -> usize {
256        self.ids.len()
257    }
258
259    /// Returns true if the view contains no detections.
260    #[must_use]
261    pub fn is_empty(&self) -> bool {
262        self.ids.is_empty()
263    }
264}
265
266impl DetectionBatch {
267    /// Returns a borrowed view of the first `v` candidates in the batch.
268    #[must_use]
269    pub fn view(&self, v: usize) -> DetectionBatchView<'_> {
270        let n = v.min(MAX_CANDIDATES);
271        DetectionBatchView {
272            ids: &self.ids[..n],
273            corners: &self.corners[..n],
274            homographies: &self.homographies[..n],
275            payloads: &self.payloads[..n],
276            error_rates: &self.error_rates[..n],
277            poses: &self.poses[..n],
278            telemetry: None,
279            rejected_corners: &[],
280            rejected_error_rates: &[],
281        }
282    }
283
284    /// Returns a borrowed view with telemetry data.
285    #[must_use]
286    pub fn view_with_telemetry(
287        &self,
288        v: usize,
289        n: usize,
290        telemetry: Option<TelemetryPayload>,
291    ) -> DetectionBatchView<'_> {
292        let v_clamped = v.min(MAX_CANDIDATES);
293        let n_clamped = n.min(MAX_CANDIDATES);
294        DetectionBatchView {
295            ids: &self.ids[..v_clamped],
296            corners: &self.corners[..v_clamped],
297            homographies: &self.homographies[..v_clamped],
298            payloads: &self.payloads[..v_clamped],
299            error_rates: &self.error_rates[..v_clamped],
300            poses: &self.poses[..v_clamped],
301            telemetry,
302            rejected_corners: &self.corners[v_clamped..n_clamped],
303            rejected_error_rates: &self.error_rates[v_clamped..n_clamped],
304        }
305    }
306}
307
308impl Default for DetectionBatch {
309    fn default() -> Self {
310        Self::new()
311    }
312}
locus_core/batch.rs

locus_core/
batch.rs
