edgefirst_decoder/
modelpack.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! Decode kernels for the Au-Zone ModelPack output format.
5//!
6//! ModelPack emits SSD-style detection (separate box and score tensors) and
7//! semantic segmentation, in flat and per-layer split variants. Every kernel
8//! here is crate-private and reached through [`crate::Decoder`], which picks
9//! the variant from the output shapes at builder time.
10
11use ndarray::{Array2, ArrayView2, ArrayView3};
12use num_traits::{AsPrimitive, Float, PrimInt};
13
14use crate::{
15    byte::{nms_int, postprocess_boxes_quant, quantize_score_threshold},
16    configs::Detection,
17    dequant_detect_box,
18    float::{nms_float, postprocess_boxes_float},
19    BBoxTypeTrait, DecoderError, DetectBox, Quantization, XYWH, XYXY,
20};
21
22/// Configuration for ModelPack split detection decoder. The quantization is
23/// ignored when decoding float models.
24#[derive(Debug, Clone, PartialEq)]
25pub(crate) struct ModelPackDetectionConfig {
26    pub(crate) anchors: Vec<[f32; 2]>,
27    pub(crate) quantization: Option<Quantization>,
28}
29
30impl TryFrom<&Detection> for ModelPackDetectionConfig {
31    type Error = DecoderError;
32
33    fn try_from(value: &Detection) -> Result<Self, DecoderError> {
34        Ok(Self {
35            anchors: value.anchors.clone().ok_or_else(|| {
36                DecoderError::InvalidConfig("ModelPack Split Detection missing anchors".to_string())
37            })?,
38            quantization: value.quantization.map(Quantization::from),
39        })
40    }
41}
42
43/// Decodes ModelPack detection outputs from quantized tensors.
44///
45/// The boxes are expected to be in XYXY format.
46///
47/// Expected shapes of inputs:
48/// - boxes: (num_boxes, 4)
49/// - scores: (num_boxes, num_classes)
50///
51/// # Panics
52/// Panics if shapes don't match the expected dimensions.
53pub(crate) fn decode_modelpack_det<
54    BOX: PrimInt + AsPrimitive<f32> + Send + Sync,
55    SCORE: PrimInt + AsPrimitive<f32> + Send + Sync,
56>(
57    boxes_tensor: (ArrayView2<BOX>, Quantization),
58    scores_tensor: (ArrayView2<SCORE>, Quantization),
59    score_threshold: f32,
60    iou_threshold: f32,
61    max_det: usize,
62    output_boxes: &mut Vec<DetectBox>,
63) where
64    f32: AsPrimitive<SCORE>,
65{
66    impl_modelpack_quant::<XYXY, _, _>(
67        boxes_tensor,
68        scores_tensor,
69        score_threshold,
70        iou_threshold,
71        max_det,
72        output_boxes,
73    )
74}
75
76/// Decodes ModelPack detection outputs from float tensors. The boxes
77/// are expected to be in XYXY format.
78///
79/// Expected shapes of inputs:
80/// - boxes: (num_boxes, 4)
81/// - scores: (num_boxes, num_classes)
82///
83/// # Panics
84/// Panics if shapes don't match the expected dimensions.
85pub(crate) fn decode_modelpack_float<
86    BOX: Float + AsPrimitive<f32> + Send + Sync,
87    SCORE: Float + AsPrimitive<f32> + Send + Sync,
88>(
89    boxes_tensor: ArrayView2<BOX>,
90    scores_tensor: ArrayView2<SCORE>,
91    score_threshold: f32,
92    iou_threshold: f32,
93    max_det: usize,
94    output_boxes: &mut Vec<DetectBox>,
95) where
96    f32: AsPrimitive<SCORE>,
97{
98    impl_modelpack_float::<XYXY, _, _>(
99        boxes_tensor,
100        scores_tensor,
101        score_threshold,
102        iou_threshold,
103        max_det,
104        output_boxes,
105    )
106}
107
108/// Decodes ModelPack split detection outputs from quantized tensors. The boxes
109/// are expected to be in XYWH format.
110///
111/// The `configs` must correspond to the `outputs` in order.
112///
113/// Expected shapes of inputs:
114/// - outputs: (width, height, num_anchors * (5 + num_classes))
115///
116/// # Panics
117/// Panics if shapes don't match the expected dimensions.
118#[cfg(test)]
119pub(crate) fn decode_modelpack_split_quant<D: AsPrimitive<f32>>(
120    outputs: &[ArrayView3<D>],
121    configs: &[ModelPackDetectionConfig],
122    score_threshold: f32,
123    iou_threshold: f32,
124    max_det: usize,
125    output_boxes: &mut Vec<DetectBox>,
126) {
127    impl_modelpack_split_quant::<XYWH, D>(
128        outputs,
129        configs,
130        score_threshold,
131        iou_threshold,
132        max_det,
133        output_boxes,
134    )
135}
136
137/// Decodes ModelPack split detection outputs from float tensors. The boxes
138/// are expected to be in XYWH format.
139///
140/// The `configs` must correspond to the `outputs` in order.
141///
142/// Expected shapes of inputs:
143/// - outputs: (width, height, num_anchors * (5 + num_classes))
144///
145/// # Panics
146/// Panics if shapes don't match the expected dimensions.
147pub(crate) fn decode_modelpack_split_float<D: AsPrimitive<f32>>(
148    outputs: &[ArrayView3<D>],
149    configs: &[ModelPackDetectionConfig],
150    score_threshold: f32,
151    iou_threshold: f32,
152    max_det: usize,
153    output_boxes: &mut Vec<DetectBox>,
154) {
155    impl_modelpack_split_float::<XYWH, D>(
156        outputs,
157        configs,
158        score_threshold,
159        iou_threshold,
160        max_det,
161        output_boxes,
162    );
163}
164/// Implementation of ModelPack detection decoding for quantized tensors.
165///
166/// Expected shapes of inputs:
167/// - boxes: (num_boxes, 4)
168/// - scores: (num_boxes, num_classes)
169///
170/// # Panics
171/// Panics if shapes don't match the expected dimensions.
172pub(crate) fn impl_modelpack_quant<
173    B: BBoxTypeTrait,
174    BOX: PrimInt + AsPrimitive<f32> + Send + Sync,
175    SCORE: PrimInt + AsPrimitive<f32> + Send + Sync,
176>(
177    boxes: (ArrayView2<BOX>, Quantization),
178    scores: (ArrayView2<SCORE>, Quantization),
179    score_threshold: f32,
180    iou_threshold: f32,
181    max_det: usize,
182    output_boxes: &mut Vec<DetectBox>,
183) where
184    f32: AsPrimitive<SCORE>,
185{
186    let (boxes_tensor, quant_boxes) = boxes;
187    let (scores_tensor, quant_scores) = scores;
188    let boxes = {
189        let score_threshold = quantize_score_threshold(score_threshold, quant_scores);
190        postprocess_boxes_quant::<B, _, _>(
191            score_threshold,
192            boxes_tensor,
193            scores_tensor,
194            quant_boxes,
195        )
196    };
197    let boxes = nms_int(iou_threshold, Some(max_det), boxes);
198    output_boxes.clear();
199    for b in boxes.into_iter().take(max_det) {
200        output_boxes.push(dequant_detect_box(&b, quant_scores));
201    }
202}
203
204/// Implementation of ModelPack detection decoding for float tensors.
205///
206/// Expected shapes of inputs:
207/// - boxes: (num_boxes, 4)
208/// - scores: (num_boxes, num_classes)
209///
210/// # Panics
211/// Panics if shapes don't match the expected dimensions.
212pub(crate) fn impl_modelpack_float<
213    B: BBoxTypeTrait,
214    BOX: Float + AsPrimitive<f32> + Send + Sync,
215    SCORE: Float + AsPrimitive<f32> + Send + Sync,
216>(
217    boxes_tensor: ArrayView2<BOX>,
218    scores_tensor: ArrayView2<SCORE>,
219    score_threshold: f32,
220    iou_threshold: f32,
221    max_det: usize,
222    output_boxes: &mut Vec<DetectBox>,
223) where
224    f32: AsPrimitive<SCORE>,
225{
226    let boxes =
227        postprocess_boxes_float::<B, _, _>(score_threshold.as_(), boxes_tensor, scores_tensor);
228    let boxes = nms_float(iou_threshold, Some(max_det), boxes);
229    output_boxes.clear();
230    for b in boxes.into_iter().take(max_det) {
231        output_boxes.push(b);
232    }
233}
234
235/// Implementation of ModelPack split detection decoding for quantized tensors.
236///
237/// Expected shapes of inputs:
238/// - boxes: (num_boxes, 4)
239/// - scores: (num_boxes, num_classes)
240///
241/// # Panics
242/// Panics if shapes don't match the expected dimensions.
243#[cfg(test)]
244pub(crate) fn impl_modelpack_split_quant<B: BBoxTypeTrait, D: AsPrimitive<f32>>(
245    outputs: &[ArrayView3<D>],
246    configs: &[ModelPackDetectionConfig],
247    score_threshold: f32,
248    iou_threshold: f32,
249    max_det: usize,
250    output_boxes: &mut Vec<DetectBox>,
251) {
252    let (boxes_tensor, scores_tensor) = postprocess_modelpack_split_quant(outputs, configs);
253    let boxes = postprocess_boxes_float::<B, _, _>(
254        score_threshold,
255        boxes_tensor.view(),
256        scores_tensor.view(),
257    );
258    let boxes = nms_float(iou_threshold, Some(max_det), boxes);
259    output_boxes.clear();
260    for b in boxes.into_iter().take(max_det) {
261        output_boxes.push(b);
262    }
263}
264
265/// Implementation of ModelPack split detection decoding for float tensors.
266///
267/// The `configs` must correspond to the `outputs` in order.
268///
269/// Expected shapes of inputs:
270/// - outputs: (width, height, num_anchors * (5 + num_classes))
271///
272/// # Panics
273/// Panics if shapes don't match the expected dimensions.
274pub(crate) fn impl_modelpack_split_float<B: BBoxTypeTrait, D: AsPrimitive<f32>>(
275    outputs: &[ArrayView3<D>],
276    configs: &[ModelPackDetectionConfig],
277    score_threshold: f32,
278    iou_threshold: f32,
279    max_det: usize,
280    output_boxes: &mut Vec<DetectBox>,
281) {
282    let (boxes_tensor, scores_tensor) = postprocess_modelpack_split_float(outputs, configs);
283    let boxes = postprocess_boxes_float::<B, _, _>(
284        score_threshold,
285        boxes_tensor.view(),
286        scores_tensor.view(),
287    );
288    let boxes = nms_float(iou_threshold, Some(max_det), boxes);
289    output_boxes.clear();
290    for b in boxes.into_iter().take(max_det) {
291        output_boxes.push(b);
292    }
293}
294
295/// Post processes ModelPack split detection into detection boxes,
296/// filtering out any boxes below the score threshold. Returns the boxes and
297/// scores tensors. Boxes are in XYWH format.
298#[cfg(test)]
299pub(crate) fn postprocess_modelpack_split_quant<T: AsPrimitive<f32>>(
300    outputs: &[ArrayView3<T>],
301    config: &[ModelPackDetectionConfig],
302) -> (Array2<f32>, Array2<f32>) {
303    let mut total_capacity = 0;
304    let mut nc = 0;
305    for (p, detail) in outputs.iter().zip(config) {
306        let shape = p.shape();
307        let na = detail.anchors.len();
308        nc = *shape
309            .last()
310            .expect("Shape must have at least one dimension")
311            / na
312            - 5;
313        total_capacity += shape[0] * shape[1] * na;
314    }
315    let mut bboxes = Vec::with_capacity(total_capacity * 4);
316    let mut bscores = Vec::with_capacity(total_capacity * nc);
317
318    for (p, detail) in outputs.iter().zip(config) {
319        let anchors = &detail.anchors;
320        let na = detail.anchors.len();
321        let shape = p.shape();
322        assert_eq!(
323            shape.iter().product::<usize>(),
324            p.len(),
325            "Shape product doesn't match tensor length"
326        );
327        let p_sigmoid = if let Some(quant) = &detail.quantization {
328            let scaled_zero = -quant.zero_point as f32 * quant.scale;
329            p.mapv(|x| fast_sigmoid_impl(x.as_() * quant.scale + scaled_zero))
330        } else {
331            p.mapv(|x| fast_sigmoid_impl(x.as_()))
332        };
333        let p_sigmoid = p_sigmoid.as_standard_layout();
334
335        // Safe to unwrap since we ensured standard layout above
336        let p = p_sigmoid
337            .as_slice()
338            .expect("Sigmoids are not in standard layout");
339        let height = shape[0];
340        let width = shape[1];
341
342        let div_width = 1.0 / width as f32;
343        let div_height = 1.0 / height as f32;
344
345        let mut grid = Vec::with_capacity(height * width * na * 2);
346        for y in 0..height {
347            for x in 0..width {
348                for _ in 0..na {
349                    grid.push(x as f32 - 0.5);
350                    grid.push(y as f32 - 0.5);
351                }
352            }
353        }
354        for ((p, g), anchor) in p
355            .chunks_exact(nc + 5)
356            .zip(grid.chunks_exact(2))
357            .zip(anchors.iter().cycle())
358        {
359            let (x, y) = (p[0], p[1]);
360            let x = (x * 2.0 + g[0]) * div_width;
361            let y = (y * 2.0 + g[1]) * div_height;
362            let (w, h) = (p[2], p[3]);
363            let w = w * w * 4.0 * anchor[0];
364            let h = h * h * 4.0 * anchor[1];
365
366            bboxes.push(x);
367            bboxes.push(y);
368            bboxes.push(w);
369            bboxes.push(h);
370
371            if nc == 1 {
372                bscores.push(p[4]);
373            } else {
374                let obj = p[4];
375                let probs = p[5..].iter().map(|x| *x * obj);
376                bscores.extend(probs);
377            }
378        }
379    }
380    // Safe to unwrap since we ensured lengths will match above
381
382    debug_assert_eq!(bboxes.len() % 4, 0);
383    debug_assert_eq!(bscores.len() % nc, 0);
384
385    let bboxes = Array2::from_shape_vec((bboxes.len() / 4, 4), bboxes)
386        .expect("Failed to create bboxes array");
387    let bscores = Array2::from_shape_vec((bscores.len() / nc, nc), bscores)
388        .expect("Failed to create bscores array");
389    (bboxes, bscores)
390}
391
392/// Post processes ModelPack split detection into detection boxes,
393/// filtering out any boxes below the score threshold. Returns the boxes and
394/// scores tensors. Boxes are in XYWH format.
395pub(crate) fn postprocess_modelpack_split_float<T: AsPrimitive<f32>>(
396    outputs: &[ArrayView3<T>],
397    config: &[ModelPackDetectionConfig],
398) -> (Array2<f32>, Array2<f32>) {
399    let mut total_capacity = 0;
400    let mut nc = 0;
401    for (p, detail) in outputs.iter().zip(config) {
402        let shape = p.shape();
403        let na = detail.anchors.len();
404        nc = *shape
405            .last()
406            .expect("Shape must have at least one dimension")
407            / na
408            - 5;
409        total_capacity += shape[0] * shape[1] * na;
410    }
411    let mut bboxes = Vec::with_capacity(total_capacity * 4);
412    let mut bscores = Vec::with_capacity(total_capacity * nc);
413
414    for (p, detail) in outputs.iter().zip(config) {
415        let anchors = &detail.anchors;
416        let na = detail.anchors.len();
417        let shape = p.shape();
418        assert_eq!(
419            shape.iter().product::<usize>(),
420            p.len(),
421            "Shape product doesn't match tensor length"
422        );
423        let p_sigmoid = p.mapv(|x| fast_sigmoid_impl(x.as_()));
424        let p_sigmoid = p_sigmoid.as_standard_layout();
425
426        // Safe to unwrap since we ensured standard layout above
427        let p = p_sigmoid
428            .as_slice()
429            .expect("Sigmoids are not in standard layout");
430        let height = shape[0];
431        let width = shape[1];
432
433        let div_width = 1.0 / width as f32;
434        let div_height = 1.0 / height as f32;
435
436        let mut grid = Vec::with_capacity(height * width * na * 2);
437        for y in 0..height {
438            for x in 0..width {
439                for _ in 0..na {
440                    grid.push(x as f32 - 0.5);
441                    grid.push(y as f32 - 0.5);
442                }
443            }
444        }
445        for ((p, g), anchor) in p
446            .chunks_exact(nc + 5)
447            .zip(grid.chunks_exact(2))
448            .zip(anchors.iter().cycle())
449        {
450            let (x, y) = (p[0], p[1]);
451            let x = (x * 2.0 + g[0]) * div_width;
452            let y = (y * 2.0 + g[1]) * div_height;
453            let (w, h) = (p[2], p[3]);
454            let w = w * w * 4.0 * anchor[0];
455            let h = h * h * 4.0 * anchor[1];
456
457            bboxes.push(x);
458            bboxes.push(y);
459            bboxes.push(w);
460            bboxes.push(h);
461
462            if nc == 1 {
463                bscores.push(p[4]);
464            } else {
465                let obj = p[4];
466                let probs = p[5..].iter().map(|x| *x * obj);
467                bscores.extend(probs);
468            }
469        }
470    }
471    // Safe to unwrap since we ensured lengths will match above
472
473    debug_assert_eq!(bboxes.len() % 4, 0);
474    debug_assert_eq!(bscores.len() % nc, 0);
475
476    let bboxes = Array2::from_shape_vec((bboxes.len() / 4, 4), bboxes)
477        .expect("Failed to create bboxes array");
478    let bscores = Array2::from_shape_vec((bscores.len() / nc, nc), bscores)
479        .expect("Failed to create bscores array");
480    (bboxes, bscores)
481}
482
483/// Fast approximation of `eˣ`, vendored bit-for-bit from the (abandoned)
484/// `fast-math` crate by Huon Wilson (MIT/Apache-2.0). Inlining the ~10-line
485/// Schraudolph float-bit-trick lets us drop the unmaintained dependency while
486/// keeping decoder outputs bit-identical to prior releases. Valid for roughly
487/// -88 ≤ x ≤ 88 (kept in range by the guard in `fast_sigmoid_impl`); max
488/// relative error < ~0.011.
489#[inline(always)]
490// Constants are reproduced verbatim from `fast-math` to stay bit-identical;
491// keep the upstream literals rather than rounding them to f32 precision.
492#[allow(clippy::excessive_precision)]
493fn exp_raw(x: f32) -> f32 {
494    const A: f32 = (1u32 << 23) as f32; // 2^SIGNIF for f32
495    const MASK: i32 = 0xff800000u32 as i32;
496    const EXP2_23: f32 = 1.1920929e-7;
497    const C0: f32 = 0.3371894346 * EXP2_23 * EXP2_23;
498    const C1: f32 = 0.657636276 * EXP2_23;
499    const C2: f32 = 1.00172476;
500
501    let a = A * core::f32::consts::LOG2_E;
502    let mul = (a * x) as i32;
503    let floor = mul & MASK;
504    let frac = (mul - floor) as f32;
505    let approx = (C0 * frac + C1) * frac + C2;
506    f32::from_bits(approx.to_bits().wrapping_add(floor as u32))
507}
508
509#[inline(always)]
510fn fast_sigmoid_impl(f: f32) -> f32 {
511    if f.abs() > 80.0 {
512        f.signum() * 0.5 + 0.5
513    } else {
514        // exp_raw is only valid for -88 < x < 88; the guard above keeps us in range
515        1.0 / (1.0 + exp_raw(-f))
516    }
517}
518
519/// Converts ModelPack segmentation into a 2D mask.
520/// The input segmentation is expected to have shape (H, W, num_classes).
521///
522/// The output mask will have shape (H, W), with values `0..num_classes` based
523/// on the argmax across the channels.
524///
525/// # Panics
526/// Panics if the input tensor does not have more than one channel.
527pub(crate) fn modelpack_segmentation_to_mask(segmentation: ArrayView3<u8>) -> Array2<u8> {
528    use argminmax::ArgMinMax;
529    assert!(
530        segmentation.shape()[2] > 1,
531        "Model Instance Segmentation should have shape (H, W, x) where x > 1"
532    );
533    let height = segmentation.shape()[0];
534    let width = segmentation.shape()[1];
535    let channels = segmentation.shape()[2];
536    let segmentation = segmentation.as_standard_layout();
537    // Safe to unwrap since we ensured standard layout above
538    let seg = segmentation
539        .as_slice()
540        .expect("Segmentation is not in standard layout");
541    let argmax = seg
542        .chunks_exact(channels)
543        .map(|x| x.argmax() as u8)
544        .collect::<Vec<_>>();
545
546    Array2::from_shape_vec((height, width), argmax).expect("Failed to create mask array")
547}
548
549#[cfg(test)]
550#[cfg_attr(coverage_nightly, coverage(off))]
551mod modelpack_tests {
552    #![allow(clippy::excessive_precision)]
553    use ndarray::Array3;
554
555    use crate::configs::{DecoderType, DimName};
556
557    use super::*;
558    #[test]
559    fn test_detection_config() {
560        let det = Detection {
561            anchors: Some(vec![[0.1, 0.13], [0.16, 0.30], [0.33, 0.23]]),
562            quantization: Some((0.1, 128).into()),
563            decoder: DecoderType::ModelPack,
564            shape: vec![1, 9, 17, 18],
565            dshape: vec![
566                (DimName::Batch, 1),
567                (DimName::Height, 9),
568                (DimName::Width, 17),
569                (DimName::NumAnchorsXFeatures, 18),
570            ],
571            normalized: Some(true),
572        };
573        let config = ModelPackDetectionConfig::try_from(&det).unwrap();
574        assert_eq!(
575            config,
576            ModelPackDetectionConfig {
577                anchors: vec![[0.1, 0.13], [0.16, 0.30], [0.33, 0.23]],
578                quantization: Some(Quantization::new(0.1, 128)),
579            }
580        );
581
582        let det = Detection {
583            anchors: None,
584            quantization: Some((0.1, 128).into()),
585            decoder: DecoderType::ModelPack,
586            shape: vec![1, 9, 17, 18],
587            dshape: vec![
588                (DimName::Batch, 1),
589                (DimName::Height, 9),
590                (DimName::Width, 17),
591                (DimName::NumAnchorsXFeatures, 18),
592            ],
593            normalized: Some(true),
594        };
595        let result = ModelPackDetectionConfig::try_from(&det);
596        assert!(
597            matches!(result, Err(DecoderError::InvalidConfig(s)) if s == "ModelPack Split Detection missing anchors")
598        );
599    }
600
601    #[test]
602    fn test_fast_sigmoid() {
603        fn full_sigmoid(x: f32) -> f32 {
604            1.0 / (1.0 + (-x).exp())
605        }
606        for i in -2550..=2550 {
607            let x = i as f32 * 0.1;
608            let fast = fast_sigmoid_impl(x);
609            let full = full_sigmoid(x);
610            let diff = (fast - full).abs();
611            assert!(
612                diff < 0.0005,
613                "Fast sigmoid differs from full sigmoid by {} at input {}",
614                diff,
615                x
616            );
617        }
618    }
619
620    #[test]
621    fn test_modelpack_segmentation_to_mask() {
622        let seg = Array3::from_shape_vec(
623            (2, 2, 3),
624            vec![
625                0u8, 10, 5, // pixel (0,0)
626                20, 15, 25, // pixel (0,1)
627                30, 5, 10, // pixel (1,0)
628                0, 0, 0, // pixel (1,1)
629            ],
630        )
631        .unwrap();
632        let mask = modelpack_segmentation_to_mask(seg.view());
633        let expected_mask = Array2::from_shape_vec((2, 2), vec![1u8, 2, 0, 0]).unwrap();
634        assert_eq!(mask, expected_mask);
635    }
636
637    #[test]
638    #[should_panic(
639        expected = "Model Instance Segmentation should have shape (H, W, x) where x > 1"
640    )]
641    fn test_modelpack_segmentation_to_mask_invalid() {
642        let seg = Array3::from_shape_vec((2, 2, 1), vec![0u8, 10, 20, 30]).unwrap();
643        let _ = modelpack_segmentation_to_mask(seg.view());
644    }
645}
edgefirst_decoder/modelpack.rs

edgefirst_decoder/
modelpack.rs
