feat: Added stuff

src/faceembed/facenet.rs

@@ -2,11 +2,57 @@ use crate::errors::*;
 use mnn_bridge::ndarray::*;
 use ndarray::{Array1, Array2, ArrayView3, ArrayView4};
 use std::path::Path;
+mod mnn_impl;
+mod ort_impl;
 
 #[derive(Debug)]
 pub struct EmbeddingGenerator {
     handle: mnn_sync::SessionHandle,
 }
+pub struct EmbeddingGeneratorBuilder {
+    schedule_config: Option<mnn::ScheduleConfig>,
+    backend_config: Option<mnn::BackendConfig>,
+    model: mnn::Interpreter,
+}
+
+impl EmbeddingGeneratorBuilder {
+    pub fn new(model: impl AsRef<[u8]>) -> Result<Self> {
+        Ok(Self {
+            schedule_config: None,
+            backend_config: None,
+            model: mnn::Interpreter::from_bytes(model.as_ref())
+                .map_err(|e| e.into_inner())
+                .change_context(Error)
+                .attach_printable("Failed to load model from bytes")?,
+        })
+    }
+
+    pub fn with_forward_type(mut self, forward_type: mnn::ForwardType) -> Self {
+        self.schedule_config
+            .get_or_insert_default()
+            .set_type(forward_type);
+        self
+    }
+
+    pub fn with_schedule_config(mut self, config: mnn::ScheduleConfig) -> Self {
+        self.schedule_config = Some(config);
+        self
+    }
+
+    pub fn with_backend_config(mut self, config: mnn::BackendConfig) -> Self {
+        self.backend_config = Some(config);
+        self
+    }
+
+    pub fn build(self) -> Result<EmbeddingGenerator> {
+        let model = self.model;
+        let sc = self.schedule_config.unwrap_or_default();
+        let handle = mnn_sync::SessionHandle::new(model, sc)
+            .change_context(Error)
+            .attach_printable("Failed to create session handle")?;
+        Ok(EmbeddingGenerator { handle })
+    }
+}
 
 impl EmbeddingGenerator {
     const INPUT_NAME: &'static str = "serving_default_input_6:0";
@@ -18,6 +64,11 @@ impl EmbeddingGenerator {
         Self::new_from_bytes(&model)
     }
 
+    pub fn builder<T: AsRef<[u8]>>()
+    -> fn(T) -> std::result::Result<EmbeddingGeneratorBuilder, Report<Error>> {
+        EmbeddingGeneratorBuilder::new
+    }
+
     pub fn new_from_bytes(model: &[u8]) -> Result<Self> {
         tracing::info!("Loading face embedding model from bytes");
         let mut model = mnn::Interpreter::from_bytes(model)
@@ -57,16 +108,24 @@ impl EmbeddingGenerator {
                 tracing::trace!("Image Tensor shape: {:?}", tensor.shape());
                 let (intptr, session) = sr.both_mut();
                 tracing::trace!("Copying input tensor to host");
-                unsafe {
+                let needs_resize = unsafe {
                     let mut input = intptr.input_unresized::<f32>(session, Self::INPUT_NAME)?;
                     tracing::trace!("Input shape: {:?}", input.shape());
                     if *input.shape() != shape {
                         tracing::trace!("Resizing input tensor to shape: {:?}", shape);
                         // input.resize(shape);
                         intptr.resize_tensor(input.view_mut(), shape);
+                        true
+                    } else {
+                        false
                     }
+                };
+                if needs_resize {
+                    tracing::trace!("Resized input tensor to shape: {:?}", shape);
+                    let now = std::time::Instant::now();
+                    intptr.resize_session(session);
+                    tracing::trace!("Session resized in {:?}", now.elapsed());
                 }
-                intptr.resize_session(session);
                 let mut input = intptr.input::<f32>(session, Self::INPUT_NAME)?;
                 tracing::trace!("Input shape: {:?}", input.shape());
                 input.copy_from_host_tensor(tensor.view())?;