vision_rs/models/yolo/yolo26/blocks/
detect.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//! YOLO26 Detect head.
19//!
20//! Implements `ultralytics.nn.modules.head.Detect(nc, ch)` with `reg_max=1`.
21//!
22//! Architecture per scale:
23//!   `cv2[i]`: Conv(c_in,c2,3,1) → Conv(c2,c2,3,1) → Conv2d(c2,4*reg_max,1,bias=True)
24//!   `cv3[i]`: DWConv(c_in,3,1) → Conv(c_in,c3,1,1) → DWConv(c3,3,1) → Conv(c3,c3,1,1)
25//!           → Conv2d(c3,nc,1,bias=True)
26
27use teeny_core::{dtype::Float, graph::{Op, SymTensor}, name_scope::name_scope};
28
29use super::conv::{conv, conv_plain, dwconv};
30
31// ── Output types ──────────────────────────────────────────────────────────────
32
33/// Output of the YOLO26 Detect head (training mode).
34pub struct DetectOutput {
35    /// Raw box predictions, concatenated across all FPN scales: (B, 4·A).
36    /// A = H0·W0 + H1·W1 + H2·W2 total anchors.  reg_max = 1.
37    pub boxes: SymTensor,
38    /// Raw class logits, concatenated across all FPN scales: (B, nc·A).
39    pub scores: SymTensor,
40}
41
42/// Combined output of both YOLO26 detect heads for dual-assignment training.
43///
44/// `one2many` (cv2/cv3) receives TAL assignment with top_k=10 and dominates
45/// early training; `one2one` (one2one_cv2/cv3) receives TAL with top_k=1 and
46/// is used exclusively at inference time.
47pub struct DualDetectOutput {
48    /// Output of the `cv2`/`cv3` (dense, TAL top_k=10) head.
49    pub one2many: DetectOutput,
50    /// Output of the `one2one_cv2`/`one2one_cv3` (TAL top_k=1) head.
51    pub one2one: DetectOutput,
52}
53
54// ── Graph helper ──────────────────────────────────────────────────────────────
55
56/// Concatenate (B, C, H_i, W_i) tensors along the flattened spatial dimension.
57///
58/// Treats each input as (B, C·H_i·W_i) and concatenates along that flat
59/// channel/spatial axis, producing (B, sum(C·H_i·W_i), 1, 1) in the graph IR.
60/// This matches the `torch.cat([t.view(B, C, -1) for t in ...], dim=-1)` pattern
61/// used in the ultralytics Detect head forward pass.
62fn channel_cat_flat(tensors: Vec<SymTensor>) -> SymTensor {
63    let c_total: usize = tensors.iter()
64        .map(|t| {
65            t.shape[1].unwrap_or(0)
66                * t.shape[2].unwrap_or(1)
67                * t.shape[3].unwrap_or(1)
68        })
69        .sum();
70    let first = &tensors[0];
71    let shape = vec![first.shape[0], Some(c_total), Some(1), Some(1)];
72    let inputs: Vec<usize> = tensors.iter().map(|t| t.node_id).collect();
73    let node_id = first.graph.borrow_mut().add_node(
74        Op::ChannelCat { c_total },
75        inputs,
76        first.dtype,
77        shape.clone(),
78    );
79    SymTensor { node_id, graph: first.graph.clone(), dtype: first.dtype, shape }
80}
81
82// ── Detect head ───────────────────────────────────────────────────────────────
83
84/// Which weight namespace the detect head should bind to.
85///
86/// `OneToMany` maps to `cv2` / `cv3` (the training head used for dense
87/// anchor assignment).  `OneToOne` maps to `one2one_cv2` / `one2one_cv3`
88/// (the inference head used by ultralytics in eval mode — better mAP).
89pub enum DetectHead {
90    /// The dense training head (`cv2`/`cv3`).
91    OneToMany,
92    /// The inference head (`one2one_cv2`/`one2one_cv3`).
93    OneToOne,
94}
95
96/// YOLO26 Detect head: 3 FPN feature maps → raw boxes + class logits.
97///
98/// Parameters:
99///   - `nc`   — number of classes (e.g. 80 for COCO)
100///   - `ch`   — channel counts for each FPN input, e.g. `[256, 512, 512]`
101///   - `head` — which weight namespace to use (`OneToMany` for training,
102///     `OneToOne` for inference — the latter loads `one2one_cv2/cv3`
103///     weights and matches ultralytics eval-mode mAP)
104///
105/// Derived widths (matching ultralytics defaults, reg_max = 1):
106///   - `c2 = max(16, ch[0]/4, reg_max*4)` — box branch hidden width (16 for n-size)
107///   - `c3 = max(ch[0], min(nc, 100))`    — cls branch hidden width (80 for n-size)
108pub fn detect<D: Float + 'static>(
109    nc: usize,
110    ch: &[usize],
111    head: DetectHead,
112) -> impl Fn(Vec<SymTensor>) -> DetectOutput + use<D> {
113    let reg_max = 1usize;
114    let c2 = [16usize, ch[0] / 4, reg_max * 4].into_iter().max().unwrap();
115    let c3 = ch[0].max(nc.min(100));
116    let (cv2_prefix, cv3_prefix) = match head {
117        DetectHead::OneToMany => ("cv2", "cv3"),
118        DetectHead::OneToOne  => ("one2one_cv2", "one2one_cv3"),
119    };
120
121    // cv2[i]: Conv(c_in→c2,3) → Conv(c2→c2,3) → Conv2d(c2→4,1,bias)
122    let cv2: Vec<Box<dyn Fn(SymTensor) -> SymTensor>> = ch
123        .iter()
124        .map(|&c_in| {
125            let l1 = conv::<D>(c_in, c2, 3, 1);
126            let l2 = conv::<D>(c2, c2, 3, 1);
127            let l3 = conv_plain::<D>(c2, 4 * reg_max, 1, 1);
128            Box::new(move |x: SymTensor| {
129                let x = { let _g = name_scope("0"); l1(x) };
130                let x = { let _g = name_scope("1"); l2(x) };
131                { let _g = name_scope("2"); l3(x) }
132            }) as Box<dyn Fn(SymTensor) -> SymTensor>
133        })
134        .collect();
135
136    // cv3[i]: Sequential(DWConv, Conv) → Sequential(DWConv, Conv) → nn.Conv2d
137    // Ultralytics naming: cv3[i][0][0]=DWConv, cv3[i][0][1]=Conv,
138    //                     cv3[i][1][0]=DWConv, cv3[i][1][1]=Conv, cv3[i][2]=plain conv
139    let cv3: Vec<Box<dyn Fn(SymTensor) -> SymTensor>> = ch
140        .iter()
141        .map(|&c_in| {
142            let dw1 = dwconv::<D>(c_in, 3, 1);
143            let pw1 = conv::<D>(c_in, c3, 1, 1);
144            let dw2 = dwconv::<D>(c3, 3, 1);
145            let pw2 = conv::<D>(c3, c3, 1, 1);
146            let out = conv_plain::<D>(c3, nc, 1, 1);
147            Box::new(move |x: SymTensor| {
148                let x = { let _g = name_scope("0.0"); dw1(x) };
149                let x = { let _g = name_scope("0.1"); pw1(x) };
150                let x = { let _g = name_scope("1.0"); dw2(x) };
151                let x = { let _g = name_scope("1.1"); pw2(x) };
152                { let _g = name_scope("2"); out(x) }
153            }) as Box<dyn Fn(SymTensor) -> SymTensor>
154        })
155        .collect();
156
157    move |feats: Vec<SymTensor>| {
158        let box_tensors: Vec<SymTensor> = feats.iter().enumerate()
159            .zip(cv2.iter())
160            .map(|((i, x), f)| { let _g = name_scope(format!("{cv2_prefix}.{i}")); f(x.clone()) })
161            .collect();
162
163        let cls_tensors: Vec<SymTensor> = feats.iter().enumerate()
164            .zip(cv3.iter())
165            .map(|((i, x), f)| { let _g = name_scope(format!("{cv3_prefix}.{i}")); f(x.clone()) })
166            .collect();
167
168        let boxes  = channel_cat_flat(box_tensors);
169        let scores = channel_cat_flat(cls_tensors);
170
171        DetectOutput { boxes, scores }
172    }
173}
vision_rs/models/yolo/yolo26/blocks/detect.rs

vision_rs/models/yolo/yolo26/blocks/
detect.rs
