vision_rs/detect/
mod.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//! Object detection interface.
19//!
20//! # Example
21//!
22//! ```rust,no_run
23//! use vision_rs::detect::{DetectorConfig, ObjectDetector, Yolo26DetectorConfig};
24//! use vision_rs::models::yolo::yolo26::Yolo26Variant;
25//!
26//! # #[tokio::main]
27//! # async fn main() -> anyhow::Result<()> {
28//! let config = DetectorConfig::Yolo26(Yolo26DetectorConfig::new(
29//!     Yolo26Variant::N,
30//!     "weights/yolo26n.bin",
31//!     vec!["car".into(), "truck".into(), "person".into()],
32//! ));
33//!
34//! let detector = ObjectDetector::new(config)?;
35//! let image_bytes = std::fs::read("frame.jpg")?;
36//! let detections = detector.detect(&image_bytes).await?;
37//!
38//! for d in &detections {
39//!     println!("{} {:.2} {:?}", d.class, d.confidence, d.bbox);
40//! }
41//! # Ok(())
42//! # }
43//! ```
44
45use std::path::PathBuf;
46
47use crate::models::yolo::yolo26::Yolo26Variant;
48
49// ── Detection result ──────────────────────────────────────────────────────────
50
51/// A single detected object.
52#[derive(Clone, Debug)]
53pub struct Detection {
54    /// Bounding box as `[cx, cy, w, h]` normalised to `[0, 1]`.
55    pub bbox: [f32; 4],
56    /// Class label resolved from the detector's `class_names` list.
57    pub class: String,
58    /// Confidence score in `[0, 1]`.
59    pub confidence: f32,
60}
61
62// ── YOLO26 config ─────────────────────────────────────────────────────────────
63
64/// Configuration for a YOLO26-backed [`ObjectDetector`].
65#[derive(Clone, Debug)]
66pub struct Yolo26DetectorConfig {
67    /// Which YOLO26 model size to use (N/S/M/L/X).
68    pub variant: Yolo26Variant,
69    /// Path to the model weights file.
70    pub weights: PathBuf,
71    /// Class labels, indexed by the model's output class index.
72    pub class_names: Vec<String>,
73    /// Minimum confidence to retain a detection. Default: `0.25`.
74    pub conf_threshold: f32,
75    /// IoU threshold for non-maximum suppression. Default: `0.45`.
76    pub nms_iou_threshold: f32,
77    /// Square input resolution the model was trained at (e.g. `640`).
78    pub img_size: usize,
79}
80
81impl Yolo26DetectorConfig {
82    /// Creates a config with default `conf_threshold` (`0.25`), `nms_iou_threshold`
83    /// (`0.45`), and `img_size` (`640`).
84    pub fn new(
85        variant: Yolo26Variant,
86        weights: impl Into<PathBuf>,
87        class_names: Vec<String>,
88    ) -> Self {
89        Self {
90            variant,
91            weights: weights.into(),
92            class_names,
93            conf_threshold: 0.25,
94            nms_iou_threshold: 0.45,
95            img_size: 640,
96        }
97    }
98}
99
100// ── Top-level config enum ─────────────────────────────────────────────────────
101
102/// Detector configuration, tagged by model family.
103#[derive(Clone, Debug)]
104pub enum DetectorConfig {
105    /// Use a YOLO26 model.
106    Yolo26(Yolo26DetectorConfig),
107}
108
109// ── ObjectDetector ────────────────────────────────────────────────────────────
110
111/// A model-agnostic object detector, dispatching to the configured model family.
112pub struct ObjectDetector {
113    inner: DetectorInner,
114}
115
116impl ObjectDetector {
117    /// Build a detector from the given config.
118    pub fn new(config: DetectorConfig) -> anyhow::Result<Self> {
119        let inner = match config {
120            DetectorConfig::Yolo26(cfg) => DetectorInner::Yolo26(Yolo26Detector { config: cfg }),
121        };
122        Ok(Self { inner })
123    }
124
125    /// Run inference on raw JPEG or PNG `image_bytes`.
126    ///
127    /// Returns every detection that clears the config's `conf_threshold`,
128    /// after non-maximum suppression.
129    pub async fn detect(&self, image_bytes: &[u8]) -> anyhow::Result<Vec<Detection>> {
130        match &self.inner {
131            DetectorInner::Yolo26(d) => d.detect(image_bytes).await,
132        }
133    }
134}
135
136// ── Private per-model implementations ────────────────────────────────────────
137
138enum DetectorInner {
139    Yolo26(Yolo26Detector),
140}
141
142struct Yolo26Detector {
143    config: Yolo26DetectorConfig,
144}
145
146impl Yolo26Detector {
147    async fn detect(&self, _image_bytes: &[u8]) -> anyhow::Result<Vec<Detection>> {
148        let _ = &self.config;
149        todo!("YOLO26 inference")
150    }
151}
vision_rs/detect/mod.rs

vision_rs/detect/
mod.rs
