face-detector/src/faceembed/facenet.rs

pub mod mnn;
pub mod ort;

use crate::errors::*;
use error_stack::ResultExt;
use ndarray::{Array1, Array2, ArrayView3, ArrayView4};
use ndarray_math::{CosineSimilarity, EuclideanDistance};

/// Configuration for face embedding processing
#[derive(Debug, Clone, PartialEq)]
pub struct FaceEmbeddingConfig {
    /// Input image width expected by the model
    pub input_width: usize,
    /// Input image height expected by the model
    pub input_height: usize,
    /// Whether to normalize embeddings to unit vectors
    pub normalize: bool,
}

impl FaceEmbeddingConfig {
    pub fn with_input_size(mut self, width: usize, height: usize) -> Self {
        self.input_width = width;
        self.input_height = height;
        self
    }

    pub fn with_normalization(mut self, normalize: bool) -> Self {
        self.normalize = normalize;
        self
    }
}

impl Default for FaceEmbeddingConfig {
    fn default() -> Self {
        Self {
            input_width: 320,
            input_height: 320,
            normalize: false,
        }
    }
}

/// Represents a face embedding vector
#[derive(Debug, Clone, PartialEq)]
pub struct FaceEmbedding {
    /// The embedding vector
    pub vector: Array1<f32>,
    /// Optional confidence score for the embedding quality
    pub confidence: Option<f32>,
}

impl FaceEmbedding {
    pub fn new(vector: Array1<f32>) -> Self {
        Self {
            vector,
            confidence: None,
        }
    }

    pub fn with_confidence(mut self, confidence: f32) -> Self {
        self.confidence = Some(confidence);
        self
    }

    /// Calculate cosine similarity with another embedding
    pub fn cosine_similarity(&self, other: &FaceEmbedding) -> f32 {
        self.vector.cosine_similarity(&other.vector).unwrap_or(0.0)
    }

    /// Calculate Euclidean distance with another embedding
    pub fn euclidean_distance(&self, other: &FaceEmbedding) -> f32 {
        self.vector
            .euclidean_distance(other.vector.view())
            .unwrap_or(f32::INFINITY)
    }

    /// Normalize the embedding vector to unit length
    pub fn normalize(&mut self) {
        let norm = self.vector.mapv(|x| x * x).sum().sqrt();
        if norm > 0.0 {
            self.vector.mapv_inplace(|x| x / norm);
        }
    }

    /// Get the dimensionality of the embedding
    pub fn dimension(&self) -> usize {
        self.vector.len()
    }
}

/// Raw model outputs that can be converted to embeddings
pub trait IntoEmbeddings {
    fn into_embeddings(self, config: &FaceEmbeddingConfig) -> Result<Vec<FaceEmbedding>>;
}

impl IntoEmbeddings for Array2<f32> {
    fn into_embeddings(self, config: &FaceEmbeddingConfig) -> Result<Vec<FaceEmbedding>> {
        let mut embeddings = Vec::new();

        for row in self.rows() {
            let mut vector = row.to_owned();

            if config.normalize {
                let norm = vector.mapv(|x| x * x).sum().sqrt();
                if norm > 0.0 {
                    vector.mapv_inplace(|x| x / norm);
                }
            }

            embeddings.push(FaceEmbedding::new(vector));
        }

        Ok(embeddings)
    }
}

/// Common trait for face embedding backends
pub trait FaceNetEmbedder {
    /// Generate embeddings for a batch of face images
    fn run_model(&mut self, faces: ArrayView4<u8>) -> Result<Array2<f32>>;

    /// Generate embeddings with full pipeline including postprocessing
    fn generate_embeddings(
        &mut self,
        faces: ArrayView4<u8>,
        config: FaceEmbeddingConfig,
    ) -> Result<Vec<FaceEmbedding>> {
        let raw_output = self
            .run_model(faces)
            .change_context(Error)
            .attach_printable("Failed to generate embeddings")?;

        raw_output
            .into_embeddings(&config)
            .attach_printable("Failed to process embeddings")
    }

    /// Generate a single embedding from a single face image
    fn generate_embedding(
        &mut self,
        face: ArrayView3<u8>,
        config: FaceEmbeddingConfig,
    ) -> Result<FaceEmbedding> {
        // Add batch dimension
        let face_batch = face.insert_axis(ndarray::Axis(0));
        let embeddings = self.generate_embeddings(face_batch.view(), config)?;

        embeddings
            .into_iter()
            .next()
            .ok_or(Error)
            .attach_printable("No embedding generated for input face")
    }
}

/// Utility functions for embedding processing
pub mod utils {
    use super::*;

    /// Compute pairwise cosine similarities between two sets of embeddings
    pub fn pairwise_cosine_similarities(
        embeddings1: &[FaceEmbedding],
        embeddings2: &[FaceEmbedding],
    ) -> Array2<f32> {
        let n1 = embeddings1.len();
        let n2 = embeddings2.len();
        let mut similarities = Array2::zeros((n1, n2));

        for (i, emb1) in embeddings1.iter().enumerate() {
            for (j, emb2) in embeddings2.iter().enumerate() {
                similarities[(i, j)] = emb1.cosine_similarity(emb2);
            }
        }

        similarities
    }

    /// Find the best matching embedding from a gallery for each query
    pub fn find_best_matches(
        queries: &[FaceEmbedding],
        gallery: &[FaceEmbedding],
    ) -> Vec<(usize, f32)> {
        let similarities = pairwise_cosine_similarities(queries, gallery);
        let mut best_matches = Vec::new();

        for i in 0..queries.len() {
            let row = similarities.row(i);
            let (best_idx, best_score) = row
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .unwrap();
            best_matches.push((best_idx, *best_score));
        }

        best_matches
    }

    /// Filter embeddings by minimum quality threshold
    pub fn filter_by_confidence(
        embeddings: Vec<FaceEmbedding>,
        min_confidence: f32,
    ) -> Vec<FaceEmbedding> {
        embeddings
            .into_iter()
            .filter(|emb| emb.confidence.map_or(true, |conf| conf >= min_confidence))
            .collect()
    }
}