Skip to main content

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}