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