DarkRoom/core/dr-face/src/classify.rs

//! TRACES: FR-CULL-8a
//! The two small classifiers behind a face's eye state (docs/dev/faces.md §17).
//!
//! **OCEC** — *open closed eyes classification*, Hyodo 2025 — reads one
//! 40×24 eye and answers P(open). **SGC** — *sunglasses classification*,
//! Hyodo 2026 — reads a 48×48 head and answers P(sunglasses); it is shown
//! two framings of each face and the higher answer stands, for the reason
//! [`crate::align::SUNGLASSES_WINDOWS`] gives. Both are
//! depthwise-separable CNNs of a few hundred kilobytes, both MIT with their
//! weights, and both were exported with BatchNorm already folded, which is
//! about the friendliest graph tract can be handed.
//!
//! Neither takes a plain buffer. [`EyeClassifier::classify`] takes an
//! [`EyePatch`] and [`SunglassesClassifier::classify`] a [`HeadViews`], each
//! constructible only by the crop in [`crate::align`] that puts the right
//! pixels in it — the same defence [`crate::embed::Embedder`] makes with
//! [`crate::align::Aligned112`], for the same reason: a classifier handed the
//! wrong region returns a confident probability of nothing. Where the eye
//! box comes from is [`crate::landmarks`]; [`EyeModels::read`] is the whole
//! chain.
//!
//! # The graphs must have a fixed batch
//!
//! Both ship with a dynamic batch dimension, which tract will not analyse.
//! `tools/fix-face-model-shapes.sh` pins it to 1, exactly as it does for the
//! embedder; the shipped files are the pinned ones.
//!
//! # Pre-processing
//!
//! Read off the reference demos rather than assumed: RGB, `x / 255`, NCHW,
//! the crop resized to the input with bilinear interpolation and **without**
//! preserving its aspect. [`crate::align`]'s crops arrive already at the
//! input size in `0..=1`, so there is nothing left to do but lay them out.

use ndarray::Array4;

use crate::align::{
    eye_box, eye_patch, head_views, EyePatch, HeadViews, EYE_PATCH_HEIGHT, EYE_PATCH_WIDTH,
    SUNGLASSES_EDGE,
};
use crate::eyes::{Eye, EyeReading};
use crate::landmarks::{Landmarker, Landmarks};
use crate::{FaceError, Pixels};
use dr_inference_engine::{Form, Model, Role};

/// A loaded OCEC graph.
pub struct EyeClassifier {
    session: Model,
}

/// A loaded SGC graph.
pub struct SunglassesClassifier {
    session: Model,
}

/// Open a single-input, single-output classifier and check it is the shape
/// the crop feeding it will be.
///
/// The check is against the *input*, because that is where these two graphs
/// differ from each other and from everything else in this crate: an SGC file
/// given to the eye classifier would otherwise be resized into by an eye
/// patch, and answer. `expected` names the model in the error.
fn open_classifier(
    bytes: &[u8],
    expected: &'static str,
    (h, w): (usize, usize),
) -> Result<Model, FaceError> {
    let model = dr_inference_engine::open(Role::EyeClassifier, Form::F32, bytes)?;
    let acquired = model.acquire()?;
    let session = acquired.lock();

    let input = session.inputs().first().ok_or(FaceError::WrongModel {
        expected,
        detail: "model has no inputs".into(),
    })?;
    let shape: Option<Vec<i64>> = input.dtype().tensor_shape().map(|s| s.to_vec());
    let want = [1, 3, h as i64, w as i64];
    if shape.as_deref() != Some(&want[..]) {
        return Err(FaceError::WrongModel {
            expected,
            detail: format!(
                "input '{}' is {:?}, expected {:?} (batch pinned to 1)",
                input.name(),
                shape,
                want
            ),
        });
    }
    if session.outputs().len() != 1 {
        return Err(FaceError::WrongModel {
            expected,
            detail: format!("{} outputs, expected one", session.outputs().len()),
        });
    }
    drop(session);
    drop(acquired);
    Ok(model)
}

/// Lay a `h × w` RGB crop out as the `[1, 3, h, w]` tensor both graphs take.
fn to_nchw(pixels: &[f32], h: usize, w: usize) -> Array4<f32> {
    let mut input = Array4::<f32>::zeros((1, 3, h, w));
    for y in 0..h {
        for x in 0..w {
            for c in 0..3 {
                input[[0, c, y, x]] = pixels[(y * w + x) * 3 + c];
            }
        }
    }
    input
}

/// Run a one-number classifier and read its sigmoid back, clamped.
fn run_scalar(model: &Model, input: Array4<f32>, expected: &'static str) -> Result<f32, FaceError> {
    let acquired = model.acquire()?;
    let mut session = acquired.lock();
    let outputs = session
        .run(ort::inputs![
            ort::value::Tensor::from_array(input).map_err(FaceError::Inference)?
        ])
        .map_err(FaceError::Inference)?;
    let (_, data) = outputs[0]
        .try_extract_tensor::<f32>()
        .map_err(FaceError::Inference)?;
    let Some(&p) = data.first() else {
        return Err(FaceError::WrongModel {
            expected,
            detail: "empty output".into(),
        });
    };
    // The graph ends in a sigmoid, so this is a clamp against rounding and
    // nothing more — the reference demo does the same.
    Ok(p.clamp(0.0, 1.0))
}

impl EyeClassifier {
    pub fn from_path(path: impl AsRef<std::path::Path>) -> Result<Self, FaceError> {
        let bytes = std::fs::read(path).map_err(FaceError::ModelRead)?;
        Self::from_bytes(&bytes)
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self, FaceError> {
        Ok(Self {
            session: open_classifier(bytes, "OCEC", (EYE_PATCH_HEIGHT, EYE_PATCH_WIDTH))?,
        })
    }

    /// P(open) for one eye.
    pub fn classify(&mut self, eye: &EyePatch) -> Result<f32, FaceError> {
        let input = to_nchw(eye.pixels(), EYE_PATCH_HEIGHT, EYE_PATCH_WIDTH);
        run_scalar(&self.session, input, "OCEC")
    }
}

impl SunglassesClassifier {
    pub fn from_path(path: impl AsRef<std::path::Path>) -> Result<Self, FaceError> {
        let bytes = std::fs::read(path).map_err(FaceError::ModelRead)?;
        Self::from_bytes(&bytes)
    }

    pub fn from_bytes(bytes: &[u8]) -> Result<Self, FaceError> {
        Ok(Self {
            session: open_classifier(bytes, "SGC", (SUNGLASSES_EDGE, SUNGLASSES_EDGE))?,
        })
    }

    /// P(sunglasses) for one head: the highest answer over its framings.
    pub fn classify(&mut self, head: &HeadViews) -> Result<f32, FaceError> {
        let mut best = 0.0_f32;
        for view in head.views() {
            let input = to_nchw(view, SUNGLASSES_EDGE, SUNGLASSES_EDGE);
            best = best.max(run_scalar(&self.session, input, "SGC")?);
        }
        Ok(best)
    }
}

/// The three models behind a reading, which is how every caller holds them.
///
/// One struct rather than three optional parameters, because a partial
/// reading is not a reading: an eye state with no sunglasses number behind
/// it is exactly the beach-photograph failure [`crate::eyes`] describes, and
/// an eye box without the landmarks is the loose one this module replaced.
/// The models load together or not at all.
pub struct EyeModels {
    pub landmarks: Landmarker,
    pub eyes: EyeClassifier,
    pub sunglasses: SunglassesClassifier,
}

impl EyeModels {
    pub fn from_paths(
        landmarks: impl AsRef<std::path::Path>,
        eyes: impl AsRef<std::path::Path>,
        sunglasses: impl AsRef<std::path::Path>,
    ) -> Result<Self, FaceError> {
        Ok(Self {
            landmarks: Landmarker::from_path(landmarks)?,
            eyes: EyeClassifier::from_path(eyes)?,
            sunglasses: SunglassesClassifier::from_path(sunglasses)?,
        })
    }

    /// Read one face's eyes, and hand back the dense landmarks it read them
    /// from.
    ///
    /// `bbox` is the detector's `(x0, y0, x1, y1)` and `landmarks5` its five
    /// points, both in source pixels; the buffer is the one the aligned
    /// crop was taken from, so an eye is read from the same pixels the
    /// embedder saw the face in. `None` where nothing could be cut — a
    /// degenerate box or landmarks — which the caller stores as "not read".
    ///
    /// The landmarks come back because they cost a model run the caller will
    /// not want to pay twice: stored beside the reading, a later pass over
    /// faces — head pose, expression — has them without the original.
    pub fn read(
        &mut self,
        px: Pixels<'_>,
        width: usize,
        height: usize,
        bbox: (f32, f32, f32, f32),
        landmarks5: &[(f32, f32); 5],
    ) -> Result<Option<(EyeReading, Landmarks)>, FaceError> {
        let Some(lm) = self.landmarks.landmarks(px, width, height, bbox)? else {
            return Ok(None);
        };
        let Some(head) = head_views(px, width, height, landmarks5) else {
            return Ok(None);
        };
        let mut eye = |contour: &[(f32, f32)]| -> Result<Eye, FaceError> {
            // A hidden eye's contour can collapse to no width. Its numbers
            // are then zero — no pixels, no sharpness — which is what the
            // rule in `crate::eyes` reads as "not readable".
            let Some(b) = eye_box(contour) else {
                return Ok(Eye {
                    open: 0.0,
                    px: 0.0,
                    sharpness: 0.0,
                });
            };
            let Some(patch) = eye_patch(px, width, height, b) else {
                return Ok(Eye {
                    open: 0.0,
                    px: 0.0,
                    sharpness: 0.0,
                });
            };
            Ok(Eye {
                open: self.eyes.classify(&patch)?,
                px: patch.source_px(),
                sharpness: patch.sharpness(),
            })
        };
        let right = eye(&lm.right_eye())?;
        let left = eye(&lm.left_eye())?;
        let reading = EyeReading {
            right,
            left,
            sunglasses: self.sunglasses.classify(&head)?,
        };
        Ok(Some((reading, lm)))
    }
}