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

kataglyphis_rustprojecttemplate/api/
onnx.rs
