Skip to main content

oxidant/api/
onnx.rs

1use log::error;
2
3pub use crate::Detection;
4
5#[flutter_rust_bridge::frb(sync)]
6pub fn detect_persons_rgba(
7    model_path: String,
8    rgba: Vec<u8>,
9    width: u32,
10    height: u32,
11    score_threshold: f32,
12) -> Result<Vec<Detection>, String> {
13    // `person_detection` is feature-gated; this must still compile without ONNX (e.g. WASM).
14    let resolved_model_path = {
15        #[cfg(onnx)]
16        {
17            crate::person_detection::resolve_model_path(Some(&model_path))
18        }
19
20        #[cfg(not(onnx))]
21        {
22            model_path
23        }
24    };
25
26    detect_persons_rgba_impl(&resolved_model_path, &rgba, width, height, score_threshold).map_err(
27        |e| {
28            error!("detect_persons_rgba failed: {e:#}");
29            format!("detect_persons_rgba failed: {e:#}")
30        },
31    )
32}
33
34#[cfg(onnx)]
35fn detect_persons_rgba_impl(
36    model_path: &str,
37    rgba: &[u8],
38    width: u32,
39    height: u32,
40    score_threshold: f32,
41) -> anyhow::Result<Vec<Detection>> {
42    use std::sync::{Mutex, OnceLock};
43
44    use crate::person_detection::PersonDetector;
45
46    struct Cached {
47        model_path: String,
48        detector: PersonDetector,
49    }
50
51    static DETECTOR: OnceLock<Mutex<Option<Cached>>> = OnceLock::new();
52    let mutex = DETECTOR.get_or_init(|| Mutex::new(None));
53
54    let lock_guard = || {
55        mutex
56            .lock()
57            .unwrap_or_else(|poisoned| poisoned.into_inner())
58    };
59
60    // Check whether a reload is needed while holding the lock briefly.
61    let needs_reload = {
62        let guard = lock_guard();
63        guard
64            .as_ref()
65            .map(|c| c.model_path != model_path)
66            .unwrap_or(true)
67    };
68
69    // Load outside the lock so concurrent callers are not blocked for a multi-second load.
70    if needs_reload {
71        let detector = PersonDetector::new(model_path)?;
72        let mut guard = lock_guard();
73        // Re-check: another thread may have loaded the same model meanwhile.
74        let still_needs = guard
75            .as_ref()
76            .map(|c| c.model_path != model_path)
77            .unwrap_or(true);
78        if still_needs {
79            *guard = Some(Cached {
80                model_path: model_path.to_string(),
81                detector,
82            });
83        }
84    }
85
86    let mut guard = lock_guard();
87    let Some(cached) = guard.as_mut() else {
88        return Ok(Vec::new());
89    };
90    cached
91        .detector
92        .infer_persons_rgba(rgba, width, height, score_threshold)
93}
94
95#[cfg(not(onnx))]
96fn detect_persons_rgba_impl(
97    _model_path: &str,
98    _rgba: &[u8],
99    _width: u32,
100    _height: u32,
101    _score_threshold: f32,
102) -> anyhow::Result<Vec<Detection>> {
103    anyhow::bail!("ONNX inference is disabled. Build with --features onnx_tract (or onnxruntime*)")
104}