insight-face-rs 1.4.0

Face detection and recognition library using ONNX Runtime (Rust)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

use crate::Result;
use ort::{
    ep,
    session::{Session, SessionOutputs},
    value::{PrimitiveTensorElementType, Tensor},
};
use std::{fmt::Debug, path::Path};

pub(crate) struct OnnxModel {
    session: Session,
}

impl OnnxModel {
    pub(crate) fn new(model_path: impl AsRef<Path>) -> Result<Self> {
        let session = Session::builder()?
            .with_execution_providers([ep::CUDA::default().build(), ep::CPU::default().build()])?
            .commit_from_file(model_path)?;
        Ok(Self { session })
    }

    pub(crate) fn run<T>(&mut self, input: Tensor<T>) -> Result<SessionOutputs<'_>>
    where
        T: PrimitiveTensorElementType + Clone + 'static + Debug,
    {
        let outputs = self.session.run(ort::inputs![input.into_dyn()])?;
        Ok(outputs)
    }
}