vision_rs/models/yolo/loss/
anchor.rs

1/*
2 * Copyright 2026 Teenygrad
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17
18//! Anchor grid generation for YOLO detection heads.
19//!
20//! YOLO26 uses three FPN levels with strides [8, 16, 32].  For each level,
21//! anchor centres are placed at the centre of every grid cell.
22
23/// Flattened anchor descriptors for all FPN levels.
24///
25/// Anchors are ordered level-by-level (stride 8 first, then 16, then 32),
26/// and within each level in row-major order (y, x).
27pub struct AnchorGrid {
28    /// Anchor centre x in input-image pixel coordinates. Shape: `[A]`.
29    pub cx: Vec<f32>,
30    /// Anchor centre y in input-image pixel coordinates. Shape: `[A]`.
31    pub cy: Vec<f32>,
32    /// Stride (scale factor from feature to input). Shape: `[A]`.
33    pub strides: Vec<f32>,
34    /// Total number of anchors A.
35    pub n_anchors: usize,
36}
37
38impl AnchorGrid {
39    /// Build the anchor grid for YOLO26 given the input image dimensions.
40    ///
41    /// `img_h` and `img_w` must be divisible by 32.
42    pub fn yolo26(img_h: usize, img_w: usize) -> Self {
43        Self::new(img_h, img_w, &[8, 16, 32])
44    }
45
46    /// Build the anchor grid for arbitrary strides.
47    pub fn new(img_h: usize, img_w: usize, strides: &[usize]) -> Self {
48        let total: usize = strides.iter().map(|&s| (img_h / s) * (img_w / s)).sum();
49        let mut cx      = Vec::with_capacity(total);
50        let mut cy      = Vec::with_capacity(total);
51        let mut stride_v = Vec::with_capacity(total);
52
53        for &s in strides {
54            let fh = img_h / s;
55            let fw = img_w / s;
56            for gy in 0..fh {
57                for gx in 0..fw {
58                    cx.push((gx as f32 + 0.5) * s as f32);
59                    cy.push((gy as f32 + 0.5) * s as f32);
60                    stride_v.push(s as f32);
61                }
62            }
63        }
64
65        AnchorGrid { n_anchors: total, cx, cy, strides: stride_v }
66    }
67
68    /// Decode raw LTRB predictions into absolute XYWH boxes.
69    ///
70    /// Input `ltrb`: `[4, A]` channels-first — (l, t, r, b) distances.
71    /// Output: `[4, A]` channels-first — (cx, cy, w, h) in pixel coords.
72    pub fn decode_ltrb_to_xywh(&self, ltrb: &[f32]) -> Vec<f32> {
73        let a = self.n_anchors;
74        assert_eq!(ltrb.len(), 4 * a);
75        let mut out = vec![0.0f32; 4 * a];
76        for i in 0..a {
77            let l = ltrb[i];
78            let t = ltrb[a + i];
79            let r = ltrb[2 * a + i];
80            let b = ltrb[3 * a + i];
81            let s = self.strides[i];
82            out[i]           = self.cx[i] + s * (r - l) * 0.5; // cx
83            out[a + i]       = self.cy[i] + s * (b - t) * 0.5; // cy
84            out[2 * a + i]   = s * (l + r);                     // w
85            out[3 * a + i]   = s * (t + b);                     // h
86        }
87        out
88    }
89
90    /// Backward through the LTRB → XYWH decode.
91    ///
92    /// `d_xywh`: `[4, A]` gradient w.r.t. decoded (cx, cy, w, h).
93    /// Returns: `[4, A]` gradient w.r.t. raw (l, t, r, b) predictions.
94    pub fn decode_backward(&self, d_xywh: &[f32]) -> Vec<f32> {
95        let a = self.n_anchors;
96        assert_eq!(d_xywh.len(), 4 * a);
97        let mut d_ltrb = vec![0.0f32; 4 * a];
98        for i in 0..a {
99            let d_cx = d_xywh[i];
100            let d_cy = d_xywh[a + i];
101            let d_w  = d_xywh[2 * a + i];
102            let d_h  = d_xywh[3 * a + i];
103            let s = self.strides[i];
104            d_ltrb[i]         = d_cx * (-s * 0.5) + d_w * s; // d_l
105            d_ltrb[a + i]     = d_cy * (-s * 0.5) + d_h * s; // d_t
106            d_ltrb[2 * a + i] = d_cx * ( s * 0.5) + d_w * s; // d_r
107            d_ltrb[3 * a + i] = d_cy * ( s * 0.5) + d_h * s; // d_b
108        }
109        d_ltrb
110    }
111}
vision_rs/models/yolo/loss/anchor.rs

vision_rs/models/yolo/loss/
anchor.rs
