edgefirst_decoder/tiling.rs
1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! SAHI-style tiled-inference postprocessing: lift per-tile detections to
5//! full-frame coordinates and merge them across tiles.
6//!
7//! A high-resolution frame is covered by a uniform overlapping grid of tiles
8//! (geometry lives in the `edgefirst-image` crate). Each tile is run through
9//! the small tile-input model and decoded independently to normalized `[0,1]`
10//! detections over the model input. This module lifts those to full-frame
11//! pixels ([`lift_tile_boxes`]) and merges duplicates at tile seams
12//! ([`merge_tiled_detections`]): a greedy, class-aware pass under the **IOS**
13//! (intersection-over-smaller) match metric that, by default, **keeps the
14//! highest-scoring box of each overlapping group unchanged** rather than
15//! growing it ([`MergeMode::KeepBest`]). [`TiledFrameAccumulator`] is a streaming
16//! collector so a pipelined runtime can push each tile's detections as
17//! inference completes and finalize once the frame's last tile arrives.
18//!
19//! IOS matters because an object split across a tile overlap appears as two
20//! partial boxes whose IoU is low but whose IoS is high, so IoS joins them
21//! where IoU leaves duplicates.
22//!
23//! # Merge mode (affects mAP)
24//!
25//! The original GREEDYNMM merge replaced each matched group with its
26//! **enclosing union** carrying the group's max score. On the Ocean Cleanup
27//! ADIS 4K validation set (5823 frames, one evaluator; TOP2-836) that union
28//! cost about **0.05 AP50 on every frame, tiled or not**: a whole frame with
29//! plain NMS scored 0.491 AP50, the same predictions after the union merge
30//! 0.437, and the same IoS-0.5 class-aware matching with the best box kept
31//! unchanged 0.490; a 28-tile pipeline scored 0.500 with keep-best against
32//! 0.442 with the union. Only ~5 % of boxes are touched, but a union inflates
33//! a true positive past IoU 0.5 and swallows the small, low-score true
34//! positives that lie mostly inside a larger same-class box. The seam case
35//! the union was built for is negligible next to that, so
36//! [`MergeMode::KeepBest`] is the default and [`MergeMode::Union`] is opt-in
37//! (it stays bit-exact with the original reference for parity work).
38//!
39//! # Per-tile decode guidance (affects mAP)
40//!
41//! Run the per-tile [`crate::Decoder`] with a **low score threshold** (e.g.
42//! 0.05) and **class-aware** NMS, and a modest per-tile `max_det`. The merge's
43//! own [`MergeConfig::score_threshold`] defaults to `0.0` precisely because
44//! per-tile decode is the real flood control — a high per-tile threshold
45//! discards true-positive fragments before the merge can join them, collapsing
46//! the recall the IOS design buys. Final score gating belongs in
47//! [`MergeConfig::score_threshold`].
48//!
49//! # Known limitations
50//!
51//! - **Objects larger than one tile** are not reconstructed: every tile sees
52//! only a fragment. Keep-best reports the highest-scoring fragment;
53//! [`MergeMode::Union`] grows the kept box over the fragments that match
54//! it, but with no whole-object box to anchor the group the fragments may
55//! not mutually pass the IOS threshold either. Choose a tile size that
56//! exceeds the largest expected object, or add the optional full-frame
57//! downscaled pass (push it as one extra tile into the accumulator).
58//!
59//! # Reference implementations
60//!
61//! - **Grid spacing (EvenDist):** the canonical authority is HAL's own
62//! `edgefirst_image::tile_grid` (ported from the adis-uav-model `sahi()`
63//! function). `overlap_ratio` is a *minimum*; realized overlap is never
64//! rounded below it. ModelPack's validator is slated to adopt this same grid.
65//! - **Merge:** ModelPack `metrics/tiled.py::merge_tiled_detections` — mirrored
66//! here numerically, mode for mode (ModelPack carries the same `merge_mode`
67//! switch with the same keep-best default; its original union-only merge is
68//! [`MergeMode::Union`]). The only deliberate difference is tie-breaking on
69//! exactly equal scores (ascending original index here vs NumPy's unstable
70//! `argsort`), which makes the streaming accumulator order-independent;
71//! results are identical on non-degenerate inputs.
72
73use crate::float::{box_area, intersection_area, ios_value, iou_value};
74use crate::{BoundingBox, DetectBox};
75// `pub use`: TilePlacement is part of this module's public surface --
76// `lift_tile_boxes` and `push_tile` take one, and the doctests below
77// import it from here. A private import compiles but leaves the type
78// unnameable by callers.
79pub use edgefirst_decoder_abi::TilePlacement;
80use edgefirst_tensor::unletter_norm;
81
82/// Overlap metric used by the tiled-detection merge to decide whether two boxes
83/// belong to the same object.
84#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
85pub enum MatchMetric {
86 /// Intersection-over-Union (standard NMS metric).
87 Iou,
88 /// Intersection-over-Smaller (default): `inter / min(area_a, area_b)`. A
89 /// seam-split object has low IoU but high IoS, so IoS merges the fragment.
90 #[default]
91 Ios,
92}
93
94impl MatchMetric {
95 /// Metric value in `[0, 1]` for two boxes.
96 #[inline]
97 pub fn value(self, a: &BoundingBox, b: &BoundingBox) -> f32 {
98 match self {
99 MatchMetric::Iou => iou_value(a, b),
100 MatchMetric::Ios => ios_value(a, b),
101 }
102 }
103}
104
105/// What the tiled-detection merge emits for a group of boxes the match
106/// metric has joined.
107///
108/// Measured on the Ocean Cleanup ADIS 4K validation set (TOP2-836), the
109/// enclosing union cost about 0.05 AP50 on every frame (whole frame: 0.491
110/// AP50 with plain NMS, 0.437 after the union merge, 0.490 with keep-best;
111/// 28-tile pipeline: 0.500 keep-best vs 0.442 union), so [`Self::KeepBest`]
112/// is the default. See the [module docs](self#merge-mode-affects-map).
113#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
114pub enum MergeMode {
115 /// Keep the group's highest-scoring box — its coordinates and its score
116 /// — and drop the boxes it matched. Greedy NMS under the configured
117 /// metric. Default.
118 ///
119 /// The kept box is never grown, rescored or relabelled. It is emitted
120 /// [canonicalized](crate::BoundingBox::to_canonical), which every box
121 /// entering the merge is, so it differs from the input only for a
122 /// degenerate input box whose `min` exceeded its `max`.
123 #[default]
124 KeepBest,
125 /// Replace the group with its enclosing union carrying the group's max
126 /// score: the original GREEDYNMM merge (ModelPack `metrics/tiled.py`
127 /// before `merge_mode`). Opt-in; inflates boxes and costs mAP.
128 Union,
129}
130
131/// Configuration for the tiled-detection merge.
132///
133/// The defaults are IoS at 0.5, class-aware, [`MergeMode::KeepBest`],
134/// `max_det` 300 and no score gate. Keep-best rather than the enclosing union
135/// because the union measured about −0.05 AP50 on every frame of the Ocean
136/// Cleanup ADIS 4K validation (TOP2-836); see [`MergeMode`].
137#[derive(Debug, Clone, Copy, PartialEq)]
138pub struct MergeConfig {
139 /// Overlap metric (default [`MatchMetric::Ios`]).
140 pub metric: MatchMetric,
141 /// Merge two boxes when `metric.value(a, b) >= threshold` (default 0.5).
142 pub threshold: f32,
143 /// Merge across classes when true (default false).
144 pub class_agnostic: bool,
145 /// Cap on returned detections after the merge (default 300).
146 pub max_det: usize,
147 /// Drop merged groups whose max score is below this (default 0.0 = keep
148 /// all; per-tile decode is the real flood control).
149 pub score_threshold: f32,
150 /// What a matched group becomes (default [`MergeMode::KeepBest`]).
151 pub mode: MergeMode,
152}
153
154impl Default for MergeConfig {
155 fn default() -> Self {
156 Self {
157 metric: MatchMetric::Ios,
158 threshold: 0.5,
159 class_agnostic: false,
160 max_det: 300,
161 score_threshold: 0.0,
162 mode: MergeMode::KeepBest,
163 }
164 }
165}
166
167/// Lift tile-local **normalized** `[0,1]` xyxy detections (over the model
168/// input) to full-frame **pixel** xyxy. Mirrors
169/// `metrics/tiled.py::lift_tile_boxes`: optionally invert the letterbox, then
170/// `full = origin + norm * crop_size`. Consumes and rewrites `boxes` in place.
171///
172/// # Examples
173/// ```
174/// use edgefirst_decoder::tiling::{lift_tile_boxes, TilePlacement};
175/// use edgefirst_decoder::{BoundingBox, DetectBox};
176///
177/// let placement = TilePlacement {
178/// index: 0, count: 1,
179/// origin: (100.0, 200.0), crop_size: (640.0, 640.0),
180/// letterbox: None, frame_dims: (3840.0, 2160.0),
181/// };
182/// let tile_local = DetectBox { bbox: BoundingBox::new(0.0, 0.0, 1.0, 1.0), score: 0.9, label: 0 };
183/// let lifted = lift_tile_boxes(vec![tile_local], &placement);
184/// assert_eq!(lifted[0].bbox, BoundingBox::new(100.0, 200.0, 740.0, 840.0));
185/// ```
186#[must_use]
187pub fn lift_tile_boxes(mut boxes: Vec<DetectBox>, placement: &TilePlacement) -> Vec<DetectBox> {
188 let _s = tracing::trace_span!("decoder.tiled.lift", boxes = boxes.len()).entered();
189 let (ox, oy) = placement.origin;
190 let (cw, ch) = placement.crop_size;
191 for d in &mut boxes {
192 let n = match placement.letterbox {
193 Some(lb) => unletter_norm(d.bbox, lb),
194 None => d.bbox,
195 };
196 d.bbox = BoundingBox {
197 xmin: ox + n.xmin * cw,
198 ymin: oy + n.ymin * ch,
199 xmax: ox + n.xmax * cw,
200 ymax: oy + n.ymax * ch,
201 };
202 }
203 boxes
204}
205
206/// Match metric using precomputed areas (avoids recomputing both operands'
207/// areas on every pair in the O(N^2) merge). Equivalent to
208/// [`MatchMetric::value`].
209#[inline]
210fn metric_value_with_areas(
211 metric: MatchMetric,
212 a: &BoundingBox,
213 area_a: f32,
214 b: &BoundingBox,
215 area_b: f32,
216) -> f32 {
217 let inter = intersection_area(a, b);
218 let denom = match metric {
219 MatchMetric::Iou => area_a + area_b - inter,
220 MatchMetric::Ios => area_a.min(area_b),
221 };
222 inter / denom.max(1e-9)
223}
224
225/// Greedy merge of lifted full-frame detections. Mirrors
226/// `metrics/tiled.py::merge_tiled_detections`:
227///
228/// 1. Sort descending by score (ties broken by ascending original index so the
229/// result is deterministic — this differs from NumPy's unstable `argsort`
230/// only on exact ties).
231/// 2. For each unused `base` in order, find later unused boxes (same class
232/// unless `class_agnostic`) whose `metric.value(base, cand) >= threshold` —
233/// matched against the **original** base box — and mark them used. With
234/// [`MergeMode::KeepBest`] (default) the base box is emitted unchanged
235/// (canonicalized, never grown or rescored); with [`MergeMode::Union`] the group is
236/// replaced by its **enclosing union** carrying the group's **max** score.
237/// Either way the output carries the base's label.
238/// 3. Drop groups below `score_threshold` and truncate to `max_det`.
239///
240/// Keep-best is the default because the union measured about −0.05 AP50 on
241/// every frame of the Ocean Cleanup ADIS 4K validation (TOP2-836): whole
242/// frame 0.491 AP50 with plain NMS, 0.437 after the union, 0.490 with
243/// keep-best. Operates in pixel space (the metric's `1e-9` epsilon is
244/// calibrated to pixel areas).
245///
246/// # Examples
247/// ```
248/// use edgefirst_decoder::tiling::{merge_tiled_detections, MatchMetric, MergeConfig, MergeMode};
249/// use edgefirst_decoder::{BoundingBox, DetectBox};
250///
251/// // A fragment (B) fully inside the full detection (A): IoS=1.0, IoU≈0.17.
252/// let a = DetectBox { bbox: BoundingBox::new(100.0, 100.0, 400.0, 300.0), score: 0.9, label: 0 };
253/// let b = DetectBox { bbox: BoundingBox::new(350.0, 100.0, 400.0, 300.0), score: 0.7, label: 0 };
254///
255/// // IOS matches the fragment; keep-best (the default) suppresses it and
256/// // leaves A's box unchanged…
257/// let ios = merge_tiled_detections(vec![a, b], &MergeConfig::default());
258/// assert_eq!(ios.len(), 1);
259/// assert_eq!(ios[0].bbox, a.bbox);
260/// assert_eq!(ios[0].score, 0.9);
261///
262/// // …while IOU leaves the two separate.
263/// let cfg = MergeConfig { metric: MatchMetric::Iou, ..MergeConfig::default() };
264/// assert_eq!(merge_tiled_detections(vec![a, b], &cfg).len(), 2);
265///
266/// // A box that extends past A at IoS 0.5: keep-best still returns A as is,
267/// // the legacy union mode grows A to the group's enclosing union.
268/// let c = DetectBox { bbox: BoundingBox::new(350.0, 100.0, 450.0, 300.0), score: 0.7, label: 0 };
269/// let keep = merge_tiled_detections(vec![a, c], &MergeConfig::default());
270/// assert_eq!(keep[0].bbox, a.bbox);
271/// let union = MergeConfig { mode: MergeMode::Union, ..MergeConfig::default() };
272/// let grown = merge_tiled_detections(vec![a, c], &union);
273/// assert_eq!(grown.len(), 1);
274/// assert_eq!(grown[0].bbox, BoundingBox::new(100.0, 100.0, 450.0, 300.0));
275/// ```
276#[must_use]
277pub fn merge_tiled_detections(dets: Vec<DetectBox>, cfg: &MergeConfig) -> Vec<DetectBox> {
278 if dets.is_empty() {
279 return dets;
280 }
281
282 // Descending score, ties by ascending original index (deterministic).
283 let mut order: Vec<usize> = (0..dets.len()).collect();
284 order.sort_by(|&i, &j| dets[j].score.total_cmp(&dets[i].score).then(i.cmp(&j)));
285
286 // Canonicalize once so degenerate (inverted) boxes have well-defined areas,
287 // and precompute each box's area once (the O(N^2) loop below would otherwise
288 // recompute both operands' areas on every pair).
289 let boxes: Vec<BoundingBox> = dets.iter().map(|d| d.bbox.to_canonical()).collect();
290 let areas: Vec<f32> = boxes.iter().map(box_area).collect();
291
292 let n = dets.len();
293 let mut used = vec![false; n];
294 let mut out: Vec<DetectBox> = Vec::with_capacity(n);
295
296 for oi in 0..n {
297 let i = order[oi];
298 if used[i] {
299 continue;
300 }
301 used[i] = true;
302 let base_box = boxes[i];
303 let base_area = areas[i];
304 let base_label = dets[i].label;
305 // What the group becomes. Keep-best never touches these; union grows
306 // the box over every matched candidate and takes the max score.
307 let mut out_box = base_box;
308 let mut out_score = dets[i].score;
309
310 for &j in &order[(oi + 1)..] {
311 if used[j] {
312 continue;
313 }
314 if !cfg.class_agnostic && dets[j].label != base_label {
315 continue;
316 }
317 // Metric measured against the ORIGINAL base box (parity with the
318 // reference), using cached areas.
319 if metric_value_with_areas(cfg.metric, &base_box, base_area, &boxes[j], areas[j])
320 >= cfg.threshold
321 {
322 used[j] = true;
323 match cfg.mode {
324 MergeMode::KeepBest => {}
325 MergeMode::Union => {
326 let c = boxes[j];
327 out_box.xmin = out_box.xmin.min(c.xmin);
328 out_box.ymin = out_box.ymin.min(c.ymin);
329 out_box.xmax = out_box.xmax.max(c.xmax);
330 out_box.ymax = out_box.ymax.max(c.ymax);
331 out_score = out_score.max(dets[j].score);
332 }
333 }
334 }
335 }
336
337 out.push(DetectBox {
338 bbox: out_box,
339 score: out_score,
340 label: base_label,
341 });
342 }
343
344 if cfg.score_threshold > 0.0 {
345 out.retain(|d| d.score >= cfg.score_threshold);
346 }
347 out.truncate(cfg.max_det);
348 out
349}
350
351/// Streaming collector for one frame's tiled detections. A pipelined runtime
352/// pushes each tile's per-tile-decoded boxes as inference completes (any
353/// order), then finalizes once every tile has arrived — the "collect after the
354/// final tile" fan-in. Not internally synchronized; keep one accumulator per
355/// in-flight frame.
356#[derive(Debug, Clone)]
357pub struct TiledFrameAccumulator {
358 frame_dims: (f32, f32),
359 tiles_total: usize,
360 /// Per-tile-index arrival flags — makes `push_tile` idempotent and the
361 /// completion fence robust to duplicate / out-of-range / retried pushes
362 /// (an at-least-once async pipeline can deliver the same tile twice).
363 seen: Vec<bool>,
364 received: usize,
365 dets: Vec<DetectBox>,
366 cfg: MergeConfig,
367}
368
369impl TiledFrameAccumulator {
370 /// Create an accumulator for a frame with `tiles_total` tiles. `frame_dims`
371 /// is `(frame_w, frame_h)` in pixels, used by [`Self::finalize_normalized`].
372 /// `est_per_tile` pre-reserves the detection buffer.
373 ///
374 /// # Examples
375 /// ```
376 /// use edgefirst_decoder::tiling::{MergeConfig, TiledFrameAccumulator};
377 /// let acc = TiledFrameAccumulator::new((1920.0, 1080.0), 12, MergeConfig::default(), 16);
378 /// assert_eq!(acc.remaining(), 12);
379 /// assert!(!acc.is_complete());
380 /// ```
381 pub fn new(
382 frame_dims: (f32, f32),
383 tiles_total: usize,
384 cfg: MergeConfig,
385 est_per_tile: usize,
386 ) -> Self {
387 Self {
388 frame_dims,
389 tiles_total,
390 seen: vec![false; tiles_total],
391 received: 0,
392 dets: Vec::with_capacity(tiles_total.saturating_mul(est_per_tile)),
393 cfg,
394 }
395 }
396
397 /// Lift one tile's per-tile-decoded boxes to full-frame pixels and append
398 /// them. Idempotent per `placement.index`: a duplicate, out-of-range, or
399 /// foreign placement (one whose `count` disagrees with this accumulator's
400 /// `tiles_total`, i.e. from a different plan/frame) is ignored and its boxes
401 /// dropped, so out-of-order **and** at-least-once delivery both converge to
402 /// the same result. Returns `true` if the tile was newly accepted, `false`
403 /// otherwise.
404 pub fn push_tile(&mut self, tile_boxes: Vec<DetectBox>, placement: &TilePlacement) -> bool {
405 let idx = placement.index;
406 // Runtime guard (not debug-only): reject placements from a different
407 // plan so a mixed-frame mistake can't corrupt fan-in completion.
408 if placement.count != self.tiles_total || idx >= self.tiles_total || self.seen[idx] {
409 return false;
410 }
411 self.seen[idx] = true;
412 self.dets.extend(lift_tile_boxes(tile_boxes, placement));
413 self.received += 1;
414 true
415 }
416
417 /// True once every tile of the frame has been pushed (by distinct index).
418 #[inline]
419 pub fn is_complete(&self) -> bool {
420 self.received >= self.tiles_total
421 }
422
423 /// Tiles still outstanding.
424 #[inline]
425 pub fn remaining(&self) -> usize {
426 self.tiles_total.saturating_sub(self.received)
427 }
428
429 /// Merge all accumulated detections into full-frame **pixel** boxes.
430 #[must_use]
431 pub fn finalize(self) -> Vec<DetectBox> {
432 let span = tracing::trace_span!(
433 "decoder.tiled.merge",
434 tiles = self.tiles_total,
435 boxes_in = self.dets.len(),
436 boxes_out = tracing::field::Empty,
437 );
438 let _s = span.enter();
439 let out = merge_tiled_detections(self.dets, &self.cfg);
440 span.record("boxes_out", out.len());
441 out
442 }
443
444 /// Merge then renormalize to `[0,1]` by `frame_dims` (for the tracker,
445 /// matching the non-tiled normalized-detection contract).
446 ///
447 /// Returns an empty list when `frame_dims` are non-finite or non-positive
448 /// rather than emitting Inf/NaN coordinates.
449 #[must_use]
450 pub fn finalize_normalized(self) -> Vec<DetectBox> {
451 let (fw, fh) = self.frame_dims;
452 if !(fw.is_finite() && fh.is_finite() && fw > 0.0 && fh > 0.0) {
453 return Vec::new();
454 }
455 let inv_w = 1.0 / fw;
456 let inv_h = 1.0 / fh;
457 let mut merged = {
458 let span = tracing::trace_span!(
459 "decoder.tiled.merge",
460 tiles = self.tiles_total,
461 boxes_in = self.dets.len(),
462 boxes_out = tracing::field::Empty,
463 );
464 let _s = span.enter();
465 let out = merge_tiled_detections(self.dets, &self.cfg);
466 span.record("boxes_out", out.len());
467 out
468 };
469 for d in &mut merged {
470 d.bbox.xmin *= inv_w;
471 d.bbox.xmax *= inv_w;
472 d.bbox.ymin *= inv_h;
473 d.bbox.ymax *= inv_h;
474 }
475 merged
476 }
477}
478
479#[cfg(test)]
480mod tests {
481 use super::*;
482
483 fn det(b: [f32; 4], score: f32, label: usize) -> DetectBox {
484 DetectBox {
485 bbox: BoundingBox::new(b[0], b[1], b[2], b[3]),
486 score,
487 label,
488 }
489 }
490
491 // --- lift -------------------------------------------------------------
492
493 fn placement(origin: (f32, f32), crop: (f32, f32)) -> TilePlacement {
494 TilePlacement {
495 index: 0,
496 count: 1,
497 origin,
498 crop_size: crop,
499 letterbox: None,
500 frame_dims: (3840.0, 2160.0),
501 }
502 }
503
504 #[test]
505 fn lift_no_letterbox_matches_origin_plus_scale() {
506 let p = placement((100.0, 200.0), (640.0, 640.0));
507 let lifted = lift_tile_boxes(
508 vec![
509 det([0.0, 0.0, 1.0, 1.0], 0.9, 0),
510 det([0.25, 0.5, 0.75, 1.0], 0.8, 0),
511 ],
512 &p,
513 );
514 assert_eq!(lifted[0].bbox, BoundingBox::new(100.0, 200.0, 740.0, 840.0));
515 assert_eq!(lifted[1].bbox, BoundingBox::new(260.0, 520.0, 580.0, 840.0));
516 }
517
518 #[test]
519 fn lift_with_letterbox_inverts_then_scales() {
520 // A box filling the letterbox content region should, after un-padding,
521 // fill the crop and lift identically to the no-letterbox full-crop box.
522 let mut p = placement((0.0, 0.0), (640.0, 640.0));
523 p.letterbox = Some([0.1, 0.1, 0.9, 0.9]);
524 let lifted = lift_tile_boxes(vec![det([0.1, 0.1, 0.9, 0.9], 0.9, 0)], &p);
525 let b = lifted[0].bbox;
526 assert!((b.xmin - 0.0).abs() < 1e-3);
527 assert!((b.ymin - 0.0).abs() < 1e-3);
528 assert!((b.xmax - 640.0).abs() < 1e-3);
529 assert!((b.ymax - 640.0).abs() < 1e-3);
530 }
531
532 #[test]
533 fn lift_empty_is_empty() {
534 let p = placement((0.0, 0.0), (640.0, 640.0));
535 assert!(lift_tile_boxes(vec![], &p).is_empty());
536 }
537
538 #[test]
539 fn lift_roundtrip_with_letterbox() {
540 // Project a known full-frame box into tile-normalized (letterboxed)
541 // coords, then lift it back and confirm it returns to the original.
542 let p = TilePlacement {
543 index: 0,
544 count: 1,
545 origin: (100.0, 200.0),
546 crop_size: (640.0, 640.0),
547 letterbox: Some([0.1, 0.1, 0.9, 0.9]),
548 frame_dims: (1920.0, 1080.0),
549 };
550 // Full-frame target [228,328,420,520] -> crop-normalized (subtract origin,
551 // /crop) -> letterbox-normalized (scale by lb extent + offset).
552 let crop_n = [
553 (228.0 - 100.0) / 640.0,
554 (328.0 - 200.0) / 640.0,
555 (420.0 - 100.0) / 640.0,
556 (520.0 - 200.0) / 640.0,
557 ];
558 let [lx0, ly0, lx1, ly1] = [0.1, 0.1, 0.9, 0.9];
559 let model_n = det(
560 [
561 lx0 + crop_n[0] * (lx1 - lx0),
562 ly0 + crop_n[1] * (ly1 - ly0),
563 lx0 + crop_n[2] * (lx1 - lx0),
564 ly0 + crop_n[3] * (ly1 - ly0),
565 ],
566 0.9,
567 0,
568 );
569 let lifted = lift_tile_boxes(vec![model_n], &p);
570 let b = lifted[0].bbox;
571 assert!((b.xmin - 228.0).abs() < 1e-2, "xmin {}", b.xmin);
572 assert!((b.ymin - 328.0).abs() < 1e-2, "ymin {}", b.ymin);
573 assert!((b.xmax - 420.0).abs() < 1e-2, "xmax {}", b.xmax);
574 assert!((b.ymax - 520.0).abs() < 1e-2, "ymax {}", b.ymax);
575 }
576
577 #[test]
578 fn lift_letterbox_clamp_fires_and_no_div_by_zero() {
579 // A box outside the letterbox content region clamps to the edge; a
580 // degenerate (zero-span) letterbox axis must not divide by zero.
581 let p = TilePlacement {
582 index: 0,
583 count: 1,
584 origin: (0.0, 0.0),
585 crop_size: (640.0, 640.0),
586 letterbox: Some([0.1, 0.1, 0.9, 0.9]),
587 frame_dims: (640.0, 640.0),
588 };
589 // Box at the very top-left, outside [0.1,0.1] content origin.
590 let lifted = lift_tile_boxes(vec![det([0.0, 0.0, 0.05, 0.05], 0.9, 0)], &p);
591 assert_eq!(lifted[0].bbox.xmin, 0.0);
592 assert_eq!(lifted[0].bbox.ymin, 0.0);
593
594 // Degenerate letterbox (lx0 == lx1): unit scale, finite result.
595 let pd = TilePlacement {
596 letterbox: Some([0.5, 0.1, 0.5, 0.9]),
597 ..p
598 };
599 let out = lift_tile_boxes(vec![det([0.2, 0.2, 0.8, 0.8], 0.9, 0)], &pd);
600 assert!(out[0].bbox.xmin.is_finite() && out[0].bbox.xmax.is_finite());
601 }
602
603 // --- merge: the canonical IOS-vs-IOU case ----------------------------
604
605 #[test]
606 fn ios_merges_fragment_iou_does_not() {
607 // From modelpack tests/test_tiled_merge.py: B fully inside A
608 // (IoS=1.0, IoU=0.167).
609 let a = det([100.0, 100.0, 400.0, 300.0], 0.9, 0);
610 let b = det([350.0, 100.0, 400.0, 300.0], 0.7, 0);
611
612 let ios = merge_tiled_detections(
613 vec![a, b],
614 &MergeConfig {
615 metric: MatchMetric::Ios,
616 threshold: 0.5,
617 ..Default::default()
618 },
619 );
620 assert_eq!(ios.len(), 1);
621 assert_eq!(ios[0].bbox, BoundingBox::new(100.0, 100.0, 400.0, 300.0));
622 assert_eq!(ios[0].score, 0.9);
623
624 let iou = merge_tiled_detections(
625 vec![a, b],
626 &MergeConfig {
627 metric: MatchMetric::Iou,
628 threshold: 0.5,
629 ..Default::default()
630 },
631 );
632 assert_eq!(iou.len(), 2);
633 }
634
635 #[test]
636 fn merge_respects_class_unless_agnostic() {
637 let a = det([100.0, 100.0, 400.0, 300.0], 0.9, 0);
638 let b = det([350.0, 100.0, 400.0, 300.0], 0.7, 1); // different class
639
640 let aware = merge_tiled_detections(vec![a, b], &MergeConfig::default());
641 assert_eq!(aware.len(), 2);
642
643 let agnostic = merge_tiled_detections(
644 vec![a, b],
645 &MergeConfig {
646 class_agnostic: true,
647 ..Default::default()
648 },
649 );
650 assert_eq!(agnostic.len(), 1);
651 // Merged box keeps the base (highest-score) label.
652 assert_eq!(agnostic[0].label, 0);
653 assert_eq!(
654 agnostic[0].bbox,
655 BoundingBox::new(100.0, 100.0, 400.0, 300.0)
656 );
657 }
658
659 #[test]
660 fn merge_enclosing_union_for_partial_overlap() {
661 // Legacy `Union` mode: two boxes overlapping >=0.5 IoS merge to their
662 // enclosing union carrying the group's max score.
663 let a = det([0.0, 0.0, 100.0, 100.0], 0.9, 0);
664 let b = det([50.0, 0.0, 150.0, 100.0], 0.8, 0); // IoS = 0.5
665 let merged = merge_tiled_detections(
666 vec![a, b],
667 &MergeConfig {
668 metric: MatchMetric::Ios,
669 threshold: 0.5,
670 mode: MergeMode::Union,
671 ..Default::default()
672 },
673 );
674 assert_eq!(merged.len(), 1);
675 assert_eq!(merged[0].bbox, BoundingBox::new(0.0, 0.0, 150.0, 100.0));
676 assert_eq!(merged[0].score, 0.9);
677 }
678
679 #[test]
680 fn keep_best_keeps_base_box_unchanged_for_partial_overlap() {
681 // Same geometry as the union test: the lower-scoring box extends
682 // beyond the base and matches at IoS >= 0.5. `KeepBest` (the
683 // default) suppresses it and leaves the base box's coordinates and
684 // score unchanged -- it must not grow into the union.
685 let a = det([0.0, 0.0, 100.0, 100.0], 0.9, 0);
686 let b = det([50.0, 0.0, 150.0, 100.0], 0.8, 0); // IoS = 0.5
687 let merged = merge_tiled_detections(
688 vec![a, b],
689 &MergeConfig {
690 metric: MatchMetric::Ios,
691 threshold: 0.5,
692 mode: MergeMode::KeepBest,
693 ..Default::default()
694 },
695 );
696 assert_eq!(merged.len(), 1);
697 assert_eq!(merged[0].bbox, a.bbox);
698 assert_eq!(merged[0].score, a.score);
699 assert_eq!(merged[0].label, a.label);
700
701 // The default config is keep-best, so it must give the same answer.
702 let by_default = merge_tiled_detections(vec![a, b], &MergeConfig::default());
703 assert_eq!(by_default.len(), 1);
704 assert_eq!(by_default[0].bbox, a.bbox);
705 assert_eq!(by_default[0].score, a.score);
706 }
707
708 #[test]
709 fn keep_best_still_suppresses_every_matched_box_in_the_group() {
710 // Suppression, not a no-op: three same-class boxes chained through
711 // the base collapse to the base alone, and a box matched only
712 // against a suppressed box (not the base) survives, exactly as the
713 // reference's "metric against the ORIGINAL base box" rule dictates.
714 let base = det([0.0, 0.0, 100.0, 100.0], 0.9, 0);
715 let right = det([50.0, 0.0, 150.0, 100.0], 0.8, 0); // IoS 0.5 with base
716 let inside = det([10.0, 10.0, 40.0, 40.0], 0.7, 0); // fully inside base
717 let far = det([140.0, 0.0, 200.0, 100.0], 0.6, 0); // IoS 1/6 with base
718 let merged =
719 merge_tiled_detections(vec![base, right, inside, far], &MergeConfig::default());
720 assert_eq!(merged.len(), 2);
721 assert_eq!(merged[0].bbox, base.bbox);
722 assert_eq!(merged[0].score, 0.9);
723 assert_eq!(merged[1].bbox, far.bbox);
724 assert_eq!(merged[1].score, 0.6);
725 }
726
727 #[test]
728 fn default_merge_mode_is_keep_best() {
729 // The Ocean Cleanup ADIS 4K study (TOP2-836) measured the enclosing
730 // union at about -0.05 AP50 on every frame; keep-best is the default.
731 assert_eq!(MergeConfig::default().mode, MergeMode::KeepBest);
732 assert_eq!(MergeMode::default(), MergeMode::KeepBest);
733 }
734
735 #[test]
736 fn merge_disjoint_boxes_pass_through() {
737 let a = det([0.0, 0.0, 10.0, 10.0], 0.9, 0);
738 let b = det([100.0, 100.0, 110.0, 110.0], 0.8, 0);
739 let merged = merge_tiled_detections(vec![a, b], &MergeConfig::default());
740 assert_eq!(merged.len(), 2);
741 }
742
743 #[test]
744 fn merge_empty_is_empty() {
745 assert!(merge_tiled_detections(vec![], &MergeConfig::default()).is_empty());
746 }
747
748 #[test]
749 fn merge_threshold_boundary_is_inclusive() {
750 // IoS exactly == threshold must merge (>=, not >).
751 let a = det([0.0, 0.0, 100.0, 100.0], 0.9, 0);
752 let b = det([50.0, 0.0, 150.0, 100.0], 0.8, 0); // IoS = 0.5 exactly
753 let merged = merge_tiled_detections(
754 vec![a, b],
755 &MergeConfig {
756 metric: MatchMetric::Ios,
757 threshold: 0.5,
758 ..Default::default()
759 },
760 );
761 assert_eq!(merged.len(), 1);
762 }
763
764 #[test]
765 fn merge_score_threshold_drops_low_groups() {
766 let a = det([0.0, 0.0, 10.0, 10.0], 0.3, 0);
767 let b = det([100.0, 100.0, 110.0, 110.0], 0.8, 0);
768 let merged = merge_tiled_detections(
769 vec![a, b],
770 &MergeConfig {
771 score_threshold: 0.5,
772 ..Default::default()
773 },
774 );
775 assert_eq!(merged.len(), 1);
776 assert_eq!(merged[0].score, 0.8);
777 }
778
779 #[test]
780 fn merge_max_det_caps_highest_scoring() {
781 let dets: Vec<DetectBox> = (0..10)
782 .map(|i| {
783 det(
784 [i as f32 * 50.0, 0.0, i as f32 * 50.0 + 10.0, 10.0],
785 1.0 - i as f32 * 0.05,
786 0,
787 )
788 })
789 .collect();
790 let merged = merge_tiled_detections(
791 dets,
792 &MergeConfig {
793 max_det: 3,
794 ..Default::default()
795 },
796 );
797 assert_eq!(merged.len(), 3);
798 assert!(merged[0].score >= merged[1].score);
799 assert!(merged[1].score >= merged[2].score);
800 }
801
802 #[test]
803 fn merge_max_det_exact_boundary() {
804 // N disjoint boxes; max_det == N keeps all, max_det == N-1 drops one.
805 let make = || -> Vec<DetectBox> {
806 (0..5)
807 .map(|i| {
808 det(
809 [i as f32 * 50.0, 0.0, i as f32 * 50.0 + 10.0, 10.0],
810 1.0 - i as f32 * 0.05,
811 0,
812 )
813 })
814 .collect()
815 };
816 assert_eq!(
817 merge_tiled_detections(
818 make(),
819 &MergeConfig {
820 max_det: 5,
821 ..Default::default()
822 }
823 )
824 .len(),
825 5
826 );
827 assert_eq!(
828 merge_tiled_detections(
829 make(),
830 &MergeConfig {
831 max_det: 4,
832 ..Default::default()
833 }
834 )
835 .len(),
836 4
837 );
838 }
839
840 #[test]
841 fn merge_score_threshold_boundary_is_inclusive() {
842 // A group whose max score == score_threshold is kept (>=).
843 let a = det([0.0, 0.0, 10.0, 10.0], 0.5, 0);
844 let merged = merge_tiled_detections(
845 vec![a],
846 &MergeConfig {
847 score_threshold: 0.5,
848 ..Default::default()
849 },
850 );
851 assert_eq!(merged.len(), 1);
852 }
853
854 // --- accumulator ------------------------------------------------------
855
856 fn empty_placement(index: usize, count: usize) -> TilePlacement {
857 TilePlacement {
858 index,
859 count,
860 origin: (0.0, 0.0),
861 crop_size: (640.0, 640.0),
862 letterbox: None,
863 frame_dims: (640.0, 640.0),
864 }
865 }
866
867 #[test]
868 fn accumulator_fan_in_fence() {
869 let mut acc = TiledFrameAccumulator::new((640.0, 640.0), 3, MergeConfig::default(), 8);
870 assert!(!acc.is_complete());
871 assert_eq!(acc.remaining(), 3);
872 assert!(acc.push_tile(vec![], &empty_placement(0, 3)));
873 assert!(acc.push_tile(vec![], &empty_placement(1, 3)));
874 assert!(!acc.is_complete());
875 assert_eq!(acc.remaining(), 1);
876 assert!(acc.push_tile(vec![], &empty_placement(2, 3)));
877 assert!(acc.is_complete());
878 assert_eq!(acc.remaining(), 0);
879 assert!(acc.finalize().is_empty());
880 }
881
882 #[test]
883 fn accumulator_dedups_and_ignores_overpush() {
884 let mut acc = TiledFrameAccumulator::new((640.0, 640.0), 2, MergeConfig::default(), 8);
885 assert!(acc.push_tile(vec![], &empty_placement(0, 2)));
886 // Duplicate index 0 is ignored (idempotent under at-least-once delivery).
887 assert!(!acc.push_tile(vec![], &empty_placement(0, 2)));
888 assert_eq!(acc.remaining(), 1);
889 assert!(!acc.is_complete());
890 assert!(acc.push_tile(vec![], &empty_placement(1, 2)));
891 assert!(acc.is_complete());
892 // Out-of-range index is ignored, never over-counts.
893 assert!(!acc.push_tile(vec![], &empty_placement(2, 2)));
894 assert_eq!(acc.remaining(), 0);
895 }
896
897 #[test]
898 fn accumulator_rejects_foreign_plan_count() {
899 // A placement from a different plan (count != tiles_total) is rejected
900 // at runtime — not just under debug_assert — so a mixed-frame mistake
901 // can't corrupt fan-in completion.
902 let mut acc = TiledFrameAccumulator::new((640.0, 640.0), 3, MergeConfig::default(), 8);
903 assert!(!acc.push_tile(vec![], &empty_placement(0, 4)));
904 assert_eq!(acc.remaining(), 3);
905 assert!(!acc.is_complete());
906 // The same index from the correct plan is still accepted afterward.
907 assert!(acc.push_tile(vec![], &empty_placement(0, 3)));
908 assert_eq!(acc.remaining(), 2);
909 }
910
911 #[test]
912 fn accumulator_out_of_order_equals_in_order() {
913 let cfg = MergeConfig::default();
914 let frame = (1280.0, 640.0);
915 // Two tiles side by side, a box near the seam in each.
916 let p0 = TilePlacement {
917 index: 0,
918 count: 2,
919 origin: (0.0, 0.0),
920 crop_size: (640.0, 640.0),
921 letterbox: None,
922 frame_dims: frame,
923 };
924 let p1 = TilePlacement {
925 index: 1,
926 count: 2,
927 origin: (640.0, 0.0),
928 crop_size: (640.0, 640.0),
929 letterbox: None,
930 frame_dims: frame,
931 };
932 let t0 = vec![det([0.9, 0.4, 1.0, 0.6], 0.8, 0)];
933 let t1 = vec![det([0.0, 0.4, 0.1, 0.6], 0.9, 0)];
934
935 let mut a = TiledFrameAccumulator::new(frame, 2, cfg, 8);
936 a.push_tile(t0.clone(), &p0);
937 a.push_tile(t1.clone(), &p1);
938 let in_order = a.finalize();
939
940 let mut b = TiledFrameAccumulator::new(frame, 2, cfg, 8);
941 b.push_tile(t1, &p1);
942 b.push_tile(t0, &p0);
943 let out_order = b.finalize();
944
945 assert_eq!(in_order.len(), out_order.len());
946 for (x, y) in in_order.iter().zip(out_order.iter()) {
947 assert_eq!(x.bbox, y.bbox);
948 assert_eq!(x.score, y.score);
949 }
950 }
951
952 #[test]
953 fn finalize_normalized_equals_finalize_over_frame_dims() {
954 let cfg = MergeConfig::default();
955 let frame = (1280.0, 640.0);
956 let p = TilePlacement {
957 index: 0,
958 count: 1,
959 origin: (100.0, 50.0),
960 crop_size: (640.0, 640.0),
961 letterbox: None,
962 frame_dims: frame,
963 };
964 let boxes = vec![det([0.1, 0.1, 0.4, 0.4], 0.9, 0)];
965
966 let mut a = TiledFrameAccumulator::new(frame, 1, cfg, 8);
967 a.push_tile(boxes.clone(), &p);
968 let px = a.finalize();
969
970 let mut b = TiledFrameAccumulator::new(frame, 1, cfg, 8);
971 b.push_tile(boxes, &p);
972 let norm = b.finalize_normalized();
973
974 assert_eq!(px.len(), norm.len());
975 let (fw, fh) = frame;
976 for (p, n) in px.iter().zip(norm.iter()) {
977 assert!((n.bbox.xmin - p.bbox.xmin / fw).abs() < 1e-4);
978 assert!((n.bbox.ymin - p.bbox.ymin / fh).abs() < 1e-4);
979 assert!((n.bbox.xmax - p.bbox.xmax / fw).abs() < 1e-4);
980 assert!((n.bbox.ymax - p.bbox.ymax / fh).abs() < 1e-4);
981 }
982 }
983
984 #[test]
985 fn finalize_normalized_rejects_invalid_frame_dims() {
986 let cfg = MergeConfig::default();
987 let boxes = vec![det([10.0, 10.0, 40.0, 40.0], 0.9, 0)];
988 for frame in [
989 (0.0, 640.0),
990 (1280.0, 0.0),
991 (f32::NAN, 640.0),
992 (1280.0, f32::INFINITY),
993 ] {
994 let p = TilePlacement {
995 index: 0,
996 count: 1,
997 origin: (0.0, 0.0),
998 crop_size: (640.0, 640.0),
999 letterbox: None,
1000 frame_dims: frame,
1001 };
1002 let mut acc = TiledFrameAccumulator::new(frame, 1, cfg, 8);
1003 acc.push_tile(boxes.clone(), &p);
1004 assert!(
1005 acc.finalize_normalized().is_empty(),
1006 "expected empty for frame_dims={frame:?}"
1007 );
1008 }
1009 }
1010
1011 // --- end-to-end: accumulator -> merge -> tracker ----------------------
1012
1013 /// One tile sees the whole object, an overlapping tile sees a contained
1014 /// fragment of it. IOS merges them into a single full-frame detection that
1015 /// the tracker resolves to ONE track; IOU leaves two, yielding TWO tracks
1016 /// (the negative control proving IOS does its job, not that it collapses
1017 /// everything).
1018 #[cfg(feature = "tracker")]
1019 #[test]
1020 fn e2e_ios_one_track_iou_two_tracks() {
1021 use edgefirst_tracker::{ByteTrackBuilder, Tracker};
1022
1023 // Two distinct tiles (indices 0 and 1) of a 2-tile frame, each lifted via
1024 // a whole-frame placement so tile 0 yields the full object and tile 1 a
1025 // contained fragment — simulating overlapping tiles seeing the same object.
1026 let p0 = TilePlacement {
1027 index: 0,
1028 count: 2,
1029 origin: (0.0, 0.0),
1030 crop_size: (640.0, 640.0),
1031 letterbox: None,
1032 frame_dims: (640.0, 640.0),
1033 };
1034 let p1 = TilePlacement { index: 1, ..p0 };
1035 let full = det(
1036 [100.0 / 640.0, 100.0 / 640.0, 400.0 / 640.0, 300.0 / 640.0],
1037 0.9,
1038 0,
1039 );
1040 let frag = det(
1041 [350.0 / 640.0, 100.0 / 640.0, 400.0 / 640.0, 300.0 / 640.0],
1042 0.7,
1043 0,
1044 );
1045
1046 let run = |metric: MatchMetric| -> (usize, usize) {
1047 let cfg = MergeConfig {
1048 metric,
1049 threshold: 0.5,
1050 ..Default::default()
1051 };
1052 let mut acc = TiledFrameAccumulator::new((640.0, 640.0), 2, cfg, 4);
1053 acc.push_tile(vec![full], &p0);
1054 acc.push_tile(vec![frag], &p1);
1055 let merged = acc.finalize_normalized();
1056 let mut tracker = ByteTrackBuilder::new().build::<DetectBox>();
1057 let _ = tracker.update(&merged, 1_000);
1058 (merged.len(), tracker.get_active_tracks().len())
1059 };
1060
1061 let (ios_merged, ios_tracks) = run(MatchMetric::Ios);
1062 assert_eq!(ios_merged, 1, "IOS should merge the fragment");
1063 assert_eq!(ios_tracks, 1, "merged detection yields one track");
1064
1065 let (iou_merged, iou_tracks) = run(MatchMetric::Iou);
1066 assert_eq!(iou_merged, 2, "IOU should leave the fragment separate");
1067 assert_eq!(iou_tracks, 2, "two detections yield two tracks");
1068 }
1069}