media_pp\elements\sink/
ort_detector.rs

1use std::{path::Path, sync::Arc};
2
3use crate::pp_log::{PpLog, pp_error, pp_info};
4use ffmpeg_next as ffmpeg;
5use ndarray::{Array4, Axis, s};
6use ort::{inputs, session::Session, value::TensorRef};
7use thiserror::Error as ThisError;
8
9use crate::{
10    buffer::MediaBuffer,
11    control::ControlMsg,
12    element::{Element, ElementType, Sink, element_pp_log},
13    error::Result,
14};
15
16/// One detected object, in the pixel space of the frame [`OrtDetector`]
17/// was handed — see its own doc comment for why no further rescaling is
18/// needed to place this on top of that same frame.
19#[derive(Debug, Clone, Copy)]
20pub struct Detection {
21    /// Index into whatever label set the model was trained on — see
22    /// [`COCO_CLASS_LABELS`] for stock Ultralytics YOLOv8/v11 weights.
23    pub class_id: usize,
24    pub score: f32,
25    /// Top-left corner (not center — already converted from the model's
26    /// own center/width/height encoding).
27    pub x: f32,
28    pub y: f32,
29    pub width: f32,
30    pub height: f32,
31}
32
33/// Convenience label table for the 80 COCO classes stock Ultralytics
34/// YOLOv8/v11 weights are trained on. Meaningless for a custom-trained
35/// model with a different class set — index [`Detection::class_id`] into
36/// your own labels in that case instead.
37#[rustfmt::skip]
38pub const COCO_CLASS_LABELS: [&str; 80] = [
39    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
40    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
41    "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
42    "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
43    "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
44    "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet",
45    "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
46    "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
47];
48
49/// Errors specific to `OrtDetector`. Converts into the crate-wide
50/// `Error` via `?` (see [`crate::error::Error`]).
51#[derive(Debug, ThisError)]
52pub enum OrtDetectorError {
53    #[error("onnxruntime error: {0}")]
54    Ort(#[from] ort::Error),
55
56    #[error(
57        "OrtDetector only accepts RGB24 Video frames, got {0:?}; \
58         link it straight after a Scaler configured with Pixel::RGB24"
59    )]
60    UnsupportedFormat(ffmpeg::format::Pixel),
61
62    #[error(
63        "OrtDetector only accepts decoded Video frames, got a {0}; \
64         link it straight after a Scaler"
65    )]
66    UnsupportedBuffer(&'static str),
67}
68
69/// Terminal sink that runs a YOLOv8/v11-style ONNX object-detection model
70/// (an Ultralytics export: one image input, one `[1, 4 + num_classes,
71/// num_boxes]` output, box coordinates as center/width/height) on every
72/// incoming frame via `ort`, then hands the decoded, NMS-filtered
73/// detections to a plain closure — same "bring your own closure" shape as
74/// [`crate::elements::AppSink`], except the closure gets structured
75/// [`Detection`]s instead of a raw [`MediaBuffer`].
76///
77/// Expects every frame's pixel dimensions to already match the model's own
78/// input resolution (e.g. 640x640 for stock YOLOv8/11 weights) and its
79/// format to be `Pixel::RGB24` — put a [`crate::elements::Scaler`]
80/// configured that way directly upstream. Because of that, a detection's
81/// box coordinates need no rescaling back to some "original" resolution:
82/// they come straight out of the model in the exact same pixel space as
83/// the frame handed to the closure.
84///
85/// Input/output tensors are bound by position, not by name (`images` /
86/// `output0` aren't assumed) — whatever the export happens to call its
87/// single input and single output, this binds to index `0` of each.
88///
89/// NMS is per-class (a box only suppresses another box of the *same*
90/// `class_id`), matching Ultralytics' own default (non-agnostic) NMS.
91pub struct OrtDetector<F> {
92    pp_log: PpLog,
93    name: Arc<str>,
94    session: Session,
95    conf_threshold: f32,
96    iou_threshold: f32,
97    on_detections: F,
98}
99
100impl<F> OrtDetector<F>
101where
102    F: FnMut(&ffmpeg::frame::Video, &[Detection]) -> Result<()> + Send + 'static,
103{
104    /// `conf_threshold` drops candidate boxes below that class score before
105    /// NMS ever sees them; `iou_threshold` is how much two same-class boxes
106    /// may overlap before the lower-scoring one is suppressed as a
107    /// duplicate of the other.
108    pub fn new(
109        name: impl Into<String>,
110        model_path: impl AsRef<Path>,
111        conf_threshold: f32,
112        iou_threshold: f32,
113        on_detections: F,
114    ) -> Result<Self> {
115        let model_path_display = model_path.as_ref().display().to_string();
116        let session = Session::builder()
117            .map_err(OrtDetectorError::from)?
118            .commit_from_file(model_path)
119            .map_err(OrtDetectorError::from)?;
120        let name: Arc<str> = name.into().into();
121        let pp_log = element_pp_log(ElementType::OrtDetector, &name, None);
122        pp_info!(
123            pp_log: &pp_log,
124            "model loaded: path={model_path_display}, conf_threshold={conf_threshold}, iou_threshold={iou_threshold}"
125        );
126        Ok(Self {
127            name,
128            pp_log,
129            session,
130            conf_threshold,
131            iou_threshold,
132            on_detections,
133        })
134    }
135
136    /// Builds the `[1, 3, height, width]` normalized input tensor from
137    /// `frame`'s packed RGB24 bytes (skipping over `stride`'s per-row
138    /// padding, which is usually wider than `width * 3`), runs inference,
139    /// then decodes + NMS-filters the raw output into [`Detection`]s.
140    fn detect(&mut self, frame: &ffmpeg::frame::Video) -> Result<Vec<Detection>> {
141        let width = frame.width() as usize;
142        let height = frame.height() as usize;
143        let stride = frame.stride(0);
144        let data = frame.data(0);
145
146        let mut input = Array4::<f32>::zeros((1, 3, height, width));
147        for y in 0..height {
148            let row = &data[y * stride..y * stride + width * 3];
149            for x in 0..width {
150                let pixel = &row[x * 3..x * 3 + 3];
151                input[[0, 0, y, x]] = pixel[0] as f32 / 255.0;
152                input[[0, 1, y, x]] = pixel[1] as f32 / 255.0;
153                input[[0, 2, y, x]] = pixel[2] as f32 / 255.0;
154            }
155        }
156
157        let outputs = self
158            .session
159            .run(inputs![
160                TensorRef::from_array_view(&input).map_err(OrtDetectorError::from)?
161            ])
162            .map_err(OrtDetectorError::from)?;
163        // `[1, 4 + num_classes, num_boxes]` -> transpose -> `[num_boxes, 4 +
164        // num_classes, 1]` -> drop the now-trailing batch axis -> `[num_boxes,
165        // 4 + num_classes]`, one row per candidate box.
166        let output = outputs[0]
167            .try_extract_array::<f32>()
168            .map_err(OrtDetectorError::from)?
169            .t()
170            .into_owned();
171        let output = output.slice(s![.., .., 0]);
172
173        let mut candidates = Vec::new();
174        for row in output.axis_iter(Axis(0)) {
175            let (class_id, score) = row
176                .iter()
177                // first 4 columns are the box, not a class score
178                .skip(4)
179                .enumerate()
180                .map(|(index, value)| (index, *value))
181                .reduce(|best, next| if next.1 > best.1 { next } else { best })
182                .expect("model output has at least one class column");
183            if score < self.conf_threshold {
184                continue;
185            }
186            let (cx, cy, w, h) = (row[0usize], row[1usize], row[2usize], row[3usize]);
187            candidates.push(Detection {
188                class_id,
189                score,
190                x: cx - w / 2.0,
191                y: cy - h / 2.0,
192                width: w,
193                height: h,
194            });
195        }
196
197        Ok(non_max_suppression(candidates, self.iou_threshold))
198    }
199}
200
201fn iou(a: &Detection, b: &Detection) -> f32 {
202    let (ax2, ay2) = (a.x + a.width, a.y + a.height);
203    let (bx2, by2) = (b.x + b.width, b.y + b.height);
204    let overlap_w = (ax2.min(bx2) - a.x.max(b.x)).max(0.0);
205    let overlap_h = (ay2.min(by2) - a.y.max(b.y)).max(0.0);
206    let intersection = overlap_w * overlap_h;
207    let union = a.width * a.height + b.width * b.height - intersection;
208    if union <= 0.0 {
209        0.0
210    } else {
211        intersection / union
212    }
213}
214
215/// Highest score first, then greedily keeps each box that doesn't overlap
216/// (past `iou_threshold`) an already-kept box of the same `class_id`.
217fn non_max_suppression(mut candidates: Vec<Detection>, iou_threshold: f32) -> Vec<Detection> {
218    candidates.sort_by(|a, b| b.score.total_cmp(&a.score));
219
220    let mut kept: Vec<Detection> = Vec::with_capacity(candidates.len());
221    'candidates: for candidate in candidates {
222        for already_kept in &kept {
223            if already_kept.class_id == candidate.class_id
224                && iou(already_kept, &candidate) > iou_threshold
225            {
226                continue 'candidates;
227            }
228        }
229        kept.push(candidate);
230    }
231    kept
232}
233
234impl<F> Element for OrtDetector<F>
235where
236    F: FnMut(&ffmpeg::frame::Video, &[Detection]) -> Result<()> + Send + 'static,
237{
238    fn name(&self) -> Arc<str> {
239        self.name.clone()
240    }
241
242    fn element_type(&self) -> ElementType {
243        ElementType::OrtDetector
244    }
245
246    fn pp_log(&self) -> &PpLog {
247        &self.pp_log
248    }
249
250    fn pp_log_mut(&mut self) -> &mut PpLog {
251        &mut self.pp_log
252    }
253}
254
255impl<F> Sink for OrtDetector<F>
256where
257    F: FnMut(&ffmpeg::frame::Video, &[Detection]) -> Result<()> + Send + 'static,
258{
259    fn consume(&mut self, buf: MediaBuffer) -> Result<()> {
260        match buf {
261            MediaBuffer::Video(frame) => {
262                if frame.format() != ffmpeg::format::Pixel::RGB24 {
263                    pp_error!(self, "unsupported pixel format: {:?}", frame.format());
264                    return Err(OrtDetectorError::UnsupportedFormat(frame.format()).into());
265                }
266                let detections = self
267                    .detect(&frame)
268                    .inspect_err(|error| pp_error!(self, "detect failed: {error}"))?;
269                (self.on_detections)(&frame, &detections)
270            }
271            MediaBuffer::Eos => Ok(()),
272            MediaBuffer::Packet(_) => {
273                pp_error!(self, "unsupported buffer: Packet");
274                Err(OrtDetectorError::UnsupportedBuffer("Packet").into())
275            }
276            MediaBuffer::Audio(_) => {
277                pp_error!(self, "unsupported buffer: Audio");
278                Err(OrtDetectorError::UnsupportedBuffer("Audio").into())
279            }
280        }
281    }
282
283    fn control(&mut self, _msg: ControlMsg) -> Result<()> {
284        // Terminal, same as AppSink/D3d12Renderer: nothing buffered or
285        // downstream to flush/forward for any ControlMsg.
286        Ok(())
287    }
288}
media_pp\elements\sink/ort_detector.rs

media_pp\elements\sink/
ort_detector.rs
