wip face landmark detection and alignment

This commit is contained in:
Arnaud_Cayrol 2026-02-01 14:15:39 +01:00
parent 291989cf4a
commit 41175c1c1b
17 changed files with 1450 additions and 20 deletions

View file

@ -28,6 +28,11 @@ image = { version = "0.25", features = ["jpeg", "png", "webp"] }
imageproc = "0.25"
kamadak-exif = "0.5"
# ML models
ort = "2.0.0-rc.11"
dlib-face-recognition = { version = "0.3", features = ["embed-all"] }
ndarray = "0.16"
# Error handling
thiserror = "1"
anyhow = "1"

View file

@ -23,14 +23,24 @@
min_brightness: 0.1,
max_brightness: 0.95,
},
head_pose: {
enabled: true,
max_yaw: 35.0,
max_pitch: 35.0,
max_roll: 25.0,
},
eye_filter: {
enabled: false,
min_ear: 0.2,
},
output: {
size: 512,
keep_intermediates: false,
},
alignment: {
enabled: false,
left_eye_pos: [0.35, 0.4],
right_eye_pos: [0.65, 0.4],
enabled: true,
eye_y_position: 0.35,
inter_eye_distance: 0.30,
},
},
video: {
@ -92,14 +102,24 @@
min_brightness: 0.1,
max_brightness: 0.95,
},
head_pose: {
enabled: true,
max_yaw: 35.0,
max_pitch: 35.0,
max_roll: 25.0,
},
eye_filter: {
enabled: false,
min_ear: 0.2,
},
output: {
size: 512,
keep_intermediates: false,
},
alignment: {
enabled: false,
left_eye_pos: [0.35, 0.4],
right_eye_pos: [0.65, 0.4],
enabled: true,
eye_y_position: 0.35,
inter_eye_distance: 0.30,
},
},
video: {
@ -248,6 +268,147 @@
</div>
{/if}
</div>
<!-- Head Pose Filter Section -->
<div class="setting-section">
<div class="section-header">
<span class="section-title">Head Pose Filter</span>
<input
type="checkbox"
bind:checked={config.processing.head_pose.enabled}
/>
</div>
{#if config.processing.head_pose.enabled}
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Max Yaw</span>
<span class="setting-hint">Maximum left/right turn angle</span>
</label>
<div class="setting-control">
<input
type="range"
bind:value={config.processing.head_pose.max_yaw}
min="5"
max="90"
step="5"
/>
<span class="value">{config.processing.head_pose.max_yaw.toFixed(0)}°</span>
</div>
</div>
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Max Pitch</span>
<span class="setting-hint">Maximum up/down tilt angle</span>
</label>
<div class="setting-control">
<input
type="range"
bind:value={config.processing.head_pose.max_pitch}
min="5"
max="90"
step="5"
/>
<span class="value">{config.processing.head_pose.max_pitch.toFixed(0)}°</span>
</div>
</div>
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Max Roll</span>
<span class="setting-hint">Maximum head tilt angle</span>
</label>
<div class="setting-control">
<input
type="range"
bind:value={config.processing.head_pose.max_roll}
min="5"
max="90"
step="5"
/>
<span class="value">{config.processing.head_pose.max_roll.toFixed(0)}°</span>
</div>
</div>
{/if}
</div>
<!-- Eye Filter Section -->
<div class="setting-section">
<div class="section-header">
<span class="section-title">Eye Filter (Blink Detection)</span>
<input
type="checkbox"
bind:checked={config.processing.eye_filter.enabled}
/>
</div>
{#if config.processing.eye_filter.enabled}
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Min EAR</span>
<span class="setting-hint">Eye Aspect Ratio threshold (lower = more closed)</span>
</label>
<div class="setting-control">
<input
type="range"
bind:value={config.processing.eye_filter.min_ear}
min="0.1"
max="0.4"
step="0.02"
/>
<span class="value">{config.processing.eye_filter.min_ear.toFixed(2)}</span>
</div>
</div>
{/if}
</div>
<!-- Alignment Section -->
<div class="setting-section">
<div class="section-header">
<span class="section-title">Face Alignment</span>
<input
type="checkbox"
bind:checked={config.processing.alignment.enabled}
/>
</div>
{#if config.processing.alignment.enabled}
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Eye Y Position</span>
<span class="setting-hint">Vertical position of eyes (% from top)</span>
</label>
<div class="setting-control">
<input
type="range"
bind:value={config.processing.alignment.eye_y_position}
min="0.2"
max="0.5"
step="0.01"
/>
<span class="value">{(config.processing.alignment.eye_y_position * 100).toFixed(0)}%</span>
</div>
</div>
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Inter-eye Distance</span>
<span class="setting-hint">Distance between eyes (% of width)</span>
</label>
<div class="setting-control">
<input
type="range"
bind:value={config.processing.alignment.inter_eye_distance}
min="0.2"
max="0.5"
step="0.01"
/>
<span class="value">{(config.processing.alignment.inter_eye_distance * 100).toFixed(0)}%</span>
</div>
</div>
{/if}
</div>
</fieldset>
{:else if activeTab === 'output'}
<fieldset disabled={disabled || saving}>

View file

@ -9,6 +9,7 @@ export default defineConfig({
'/api': {
target: 'http://localhost:5000',
changeOrigin: true,
ws: true, // Enable WebSocket proxying
},
},
},

View file

@ -205,29 +205,144 @@ impl OutputConfig {
}
}
/// Face alignment configuration (for future landmark-based alignment).
/// Head pose estimation configuration.
///
/// Uses DMHead ONNX model to estimate head pose angles and filter
/// non-front-facing faces.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HeadPoseConfig {
/// Whether head pose filtering is enabled.
pub enabled: bool,
/// Maximum allowed yaw angle (left/right turn) in degrees.
pub max_yaw: f32,
/// Maximum allowed pitch angle (up/down tilt) in degrees.
pub max_pitch: f32,
/// Maximum allowed roll angle (head tilt) in degrees.
pub max_roll: f32,
}
impl Default for HeadPoseConfig {
fn default() -> Self {
Self {
enabled: true,
max_yaw: 35.0,
max_pitch: 35.0,
max_roll: 25.0,
}
}
}
impl HeadPoseConfig {
/// Validate the configuration values.
pub fn validate(&self) -> Result<()> {
if self.enabled {
if self.max_yaw < 0.0 || self.max_yaw > 90.0 {
return Err(Error::Config(
"Head pose max_yaw must be between 0 and 90 degrees".to_string(),
));
}
if self.max_pitch < 0.0 || self.max_pitch > 90.0 {
return Err(Error::Config(
"Head pose max_pitch must be between 0 and 90 degrees".to_string(),
));
}
if self.max_roll < 0.0 || self.max_roll > 90.0 {
return Err(Error::Config(
"Head pose max_roll must be between 0 and 90 degrees".to_string(),
));
}
}
Ok(())
}
}
/// Eye filter configuration for blink detection.
///
/// Uses Eye Aspect Ratio (EAR) computed from facial landmarks to detect
/// closed eyes.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EyeFilterConfig {
/// Whether eye filtering is enabled.
pub enabled: bool,
/// Minimum Eye Aspect Ratio (EAR) threshold.
/// Eyes with EAR below this are considered closed.
pub min_ear: f32,
}
impl Default for EyeFilterConfig {
fn default() -> Self {
Self {
enabled: false,
min_ear: 0.2,
}
}
}
impl EyeFilterConfig {
/// Validate the configuration values.
pub fn validate(&self) -> Result<()> {
if self.enabled {
if self.min_ear < 0.0 || self.min_ear > 0.5 {
return Err(Error::Config(
"Eye filter min_ear must be between 0.0 and 0.5".to_string(),
));
}
}
Ok(())
}
}
/// Face alignment configuration for landmark-based alignment.
///
/// Aligns faces based on eye positions detected from facial landmarks,
/// ensuring consistent eye placement across all images for smoother timelapses.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AlignmentConfig {
/// Whether face alignment is enabled.
pub enabled: bool,
/// Target position for left eye as (x%, y%) of output image.
pub left_eye_pos: (f32, f32),
/// Target Y position for eyes as percentage from top (0.0-1.0).
/// Default 0.35 places eyes at 35% from the top.
pub eye_y_position: f32,
/// Target position for right eye as (x%, y%) of output image.
pub right_eye_pos: (f32, f32),
/// Target inter-eye distance as percentage of output width (0.0-1.0).
/// Default 0.3 makes the distance between eye centers 30% of image width.
pub inter_eye_distance: f32,
}
impl Default for AlignmentConfig {
fn default() -> Self {
Self {
enabled: false,
left_eye_pos: (0.35, 0.4),
right_eye_pos: (0.65, 0.4),
enabled: true,
eye_y_position: 0.35,
inter_eye_distance: 0.30,
}
}
}
impl AlignmentConfig {
/// Validate the configuration values.
pub fn validate(&self) -> Result<()> {
if self.enabled {
if self.eye_y_position < 0.2 || self.eye_y_position > 0.5 {
return Err(Error::Config(
"Alignment eye_y_position must be between 0.2 and 0.5".to_string(),
));
}
if self.inter_eye_distance < 0.2 || self.inter_eye_distance > 0.5 {
return Err(Error::Config(
"Alignment inter_eye_distance must be between 0.2 and 0.5".to_string(),
));
}
}
Ok(())
}
}
// ============================================================================
// Main Processing Configuration
// ============================================================================
@ -250,11 +365,19 @@ pub struct ProcessingConfig {
#[serde(default)]
pub brightness: BrightnessConfig,
/// Head pose estimation settings.
#[serde(default)]
pub head_pose: HeadPoseConfig,
/// Eye filter settings (blink detection).
#[serde(default)]
pub eye_filter: EyeFilterConfig,
/// Output image settings.
#[serde(default)]
pub output: OutputConfig,
/// Face alignment settings (requires landmarks - future feature).
/// Face alignment settings (landmark-based).
#[serde(default)]
pub alignment: AlignmentConfig,
}
@ -265,6 +388,8 @@ impl Default for ProcessingConfig {
max_workers: num_cpus(),
face_resolution: FaceResolutionConfig::default(),
brightness: BrightnessConfig::default(),
head_pose: HeadPoseConfig::default(),
eye_filter: EyeFilterConfig::default(),
output: OutputConfig::default(),
alignment: AlignmentConfig::default(),
}
@ -276,7 +401,10 @@ impl ProcessingConfig {
pub fn validate(&self) -> Result<()> {
self.face_resolution.validate()?;
self.brightness.validate()?;
self.head_pose.validate()?;
self.eye_filter.validate()?;
self.output.validate()?;
self.alignment.validate()?;
Ok(())
}
}

View file

@ -32,6 +32,12 @@ pub enum Error {
#[error("FFmpeg error: {0}")]
FFmpeg(String),
#[error("ML model error: {0}")]
Model(String),
#[error("Landmark detection error: {0}")]
LandmarkDetection(String),
#[error("I/O error: {0}")]
Io(#[from] std::io::Error),

View file

@ -6,7 +6,7 @@
mod crop;
pub mod debug;
mod orientation;
mod types;
pub mod types;
pub use crop::*;
pub use orientation::load_image_with_orientation;

View file

@ -101,6 +101,64 @@ impl Landmarks {
pub fn right_mouth(&self) -> Point {
self.points[54]
}
/// Calculate Eye Aspect Ratio (EAR) for blink detection.
///
/// EAR is computed as:
/// EAR = (||p2-p6|| + ||p3-p5||) / (2 * ||p1-p4||)
///
/// Where p1-p6 are the 6 eye landmark points.
/// For left eye: points 36-41
/// For right eye: points 42-47
pub fn eye_aspect_ratio(&self) -> EyeAspectRatio {
let left_ear = Self::compute_ear(&self.points[36..42]);
let right_ear = Self::compute_ear(&self.points[42..48]);
EyeAspectRatio {
left: left_ear,
right: right_ear,
}
}
/// Compute EAR for a single eye given 6 landmark points.
fn compute_ear(eye: &[Point]) -> f32 {
if eye.len() != 6 {
return 0.0;
}
// Vertical distances
let v1 = Self::distance(&eye[1], &eye[5]); // p2-p6
let v2 = Self::distance(&eye[2], &eye[4]); // p3-p5
// Horizontal distance
let h = Self::distance(&eye[0], &eye[3]); // p1-p4
if h == 0.0 {
return 0.0;
}
(v1 + v2) / (2.0 * h)
}
/// Euclidean distance between two points.
fn distance(p1: &Point, p2: &Point) -> f32 {
let dx = p2.x - p1.x;
let dy = p2.y - p1.y;
(dx * dx + dy * dy).sqrt()
}
/// Get the angle (in radians) to rotate the face so eyes are horizontal.
pub fn eye_rotation_angle(&self) -> f32 {
let left_eye = self.left_eye_center();
let right_eye = self.right_eye_center();
let dy = right_eye.y - left_eye.y;
let dx = right_eye.x - left_eye.x;
dy.atan2(dx)
}
/// Get the distance between eye centers.
pub fn inter_eye_distance(&self) -> f32 {
Self::distance(&self.left_eye_center(), &self.right_eye_center())
}
}
/// Head pose angles.

View file

@ -10,6 +10,7 @@ pub mod error;
pub mod face_processing;
pub mod immich_api;
pub mod job;
pub mod models;
pub mod pipeline;
pub mod utils;
pub mod video;

View file

@ -0,0 +1,127 @@
//! Dlib landmark predictor wrapper.
//!
//! Wraps the dlib-face-recognition crate's LandmarkPredictor and FaceDetector
//! in a thread-safe singleton to avoid reloading the model for every image.
use crate::error::{Error, Result};
use crate::face_processing::types::{Landmarks, Point};
use dlib_face_recognition::{
FaceDetector, FaceDetectorTrait, ImageMatrix, LandmarkPredictor, LandmarkPredictorTrait,
Rectangle,
};
use std::sync::{Mutex, OnceLock};
/// Global landmark predictor instance.
/// Loaded lazily on first use.
static LANDMARK_PREDICTOR: OnceLock<Result<DlibLandmarks>> = OnceLock::new();
/// Thread-safe wrapper for dlib's FaceDetector and LandmarkPredictor.
///
/// The dlib types are not thread-safe, so we wrap them in a Mutex.
/// The model is loaded once and reused for all subsequent calls.
pub struct DlibLandmarks {
detector: Mutex<FaceDetector>,
predictor: Mutex<LandmarkPredictor>,
}
// Safety: The Mutex ensures thread-safe access to the inner types
unsafe impl Send for DlibLandmarks {}
unsafe impl Sync for DlibLandmarks {}
impl DlibLandmarks {
/// Load the landmark predictor model.
///
/// This will download/check the model file (shape_predictor_68_face_landmarks.dat).
fn load() -> Result<Self> {
let detector = FaceDetector::default();
let predictor = LandmarkPredictor::default()
.map_err(|e| Error::Model(format!("Failed to load landmark predictor: {}", e)))?;
Ok(Self {
detector: Mutex::new(detector),
predictor: Mutex::new(predictor),
})
}
/// Get or initialize the global landmark predictor instance.
pub fn global() -> Result<&'static DlibLandmarks> {
LANDMARK_PREDICTOR
.get_or_init(DlibLandmarks::load)
.as_ref()
.map_err(|e| Error::Model(e.to_string()))
}
/// Detect 68 facial landmarks from a cropped face image.
///
/// # Arguments
/// * `width` - Image width
/// * `height` - Image height
/// * `pixels` - Raw RGB pixel data (width * height * 3 bytes)
///
/// # Returns
/// Landmarks struct containing the 68 facial landmark points, or an error
/// if landmarks could not be detected.
pub fn detect_landmarks(
&self,
width: usize,
height: usize,
pixels: &[u8],
) -> Result<Landmarks> {
// Create image matrix for dlib
let matrix = unsafe { ImageMatrix::new(width, height, pixels.as_ptr()) };
// Lock detector and predictor
let detector = self
.detector
.lock()
.map_err(|e| Error::Model(format!("Failed to lock detector: {}", e)))?;
let predictor = self
.predictor
.lock()
.map_err(|e| Error::Model(format!("Failed to lock predictor: {}", e)))?;
// Since we have a cropped face, create a rectangle covering the whole image
let margin = 5;
let face_rect = Rectangle {
left: margin,
top: margin,
right: (width as i64) - margin,
bottom: (height as i64) - margin,
};
// Try to detect face in the cropped image first
let faces = detector.face_locations(&matrix);
// Use detected face if found, otherwise use the whole-image rectangle
let rect = if !faces.is_empty() {
faces[0].clone()
} else {
face_rect
};
// Detect landmarks
let landmarks_raw = predictor.face_landmarks(&matrix, &rect);
// Convert dlib landmarks to our Landmarks type
let points: Vec<Point> = landmarks_raw
.iter()
.map(|p| Point::new(p.x() as f32, p.y() as f32))
.collect();
Landmarks::new(points).ok_or_else(|| {
Error::Model("Could not detect 68 facial landmarks".to_string())
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_global_initialization() {
// Just test that we can get the global instance
// (this will fail if model file is not present, which is expected in CI)
let _ = DlibLandmarks::global();
}
}

166
src/models/dmhead.rs Normal file
View file

@ -0,0 +1,166 @@
//! DMHead head pose estimation model wrapper.
//!
//! DMHead is a lightweight head pose estimation model that predicts yaw, pitch,
//! and roll angles from a cropped face image.
//!
//! Model source: https://github.com/PINTO0309/DMHead
//! Input: 224x224 RGB image, normalized to [-1, 1]
//! Output: [yaw, pitch, roll] in degrees
use crate::error::{Error, Result};
use crate::face_processing::types::HeadPose;
use image::DynamicImage;
use ndarray::Array4;
use ort::session::{builder::GraphOptimizationLevel, Session};
use std::path::Path;
use std::sync::{Mutex, OnceLock};
/// Global model instance for DMHead.
/// Loaded lazily on first use.
static DMHEAD_MODEL: OnceLock<Result<DMHeadModel>> = OnceLock::new();
/// DMHead ONNX model for head pose estimation.
pub struct DMHeadModel {
session: Mutex<Session>,
}
impl DMHeadModel {
/// Model input size (width and height).
pub const INPUT_SIZE: u32 = 224;
/// Load the DMHead model from the given path.
pub fn load(model_path: impl AsRef<Path>) -> Result<Self> {
let session = Session::builder()
.map_err(|e| Error::Model(format!("Failed to create ONNX session builder: {}", e)))?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|e| Error::Model(format!("Failed to set optimization level: {}", e)))?
.commit_from_file(model_path.as_ref())
.map_err(|e| {
Error::Model(format!(
"Failed to load DMHead model from {}: {}",
model_path.as_ref().display(),
e
))
})?;
Ok(Self {
session: Mutex::new(session),
})
}
/// Get or initialize the global DMHead model instance.
///
/// The model is loaded from `models/dmhead_nomask_Nx3x224x224.onnx` relative
/// to the current working directory.
pub fn global() -> Result<&'static DMHeadModel> {
DMHEAD_MODEL
.get_or_init(|| {
let model_path = Path::new("models/dmhead_nomask_Nx3x224x224.onnx");
if !model_path.exists() {
return Err(Error::Model(format!(
"DMHead model not found at {}. \
Download from: https://github.com/PINTO0309/DMHead/releases",
model_path.display()
)));
}
DMHeadModel::load(model_path)
})
.as_ref()
.map_err(|e| Error::Model(e.to_string()))
}
/// Estimate head pose from a face image.
///
/// The image should be a cropped face. It will be resized to 224x224 if necessary.
///
/// Returns (yaw, pitch, roll) in degrees where:
/// - Yaw: left/right rotation (-90 to +90, positive = looking right)
/// - Pitch: up/down rotation (-90 to +90, positive = looking up)
/// - Roll: head tilt (-90 to +90, positive = tilting right)
pub fn estimate(&self, image: &DynamicImage) -> Result<HeadPose> {
// Resize image to model input size
let resized = image.resize_exact(
Self::INPUT_SIZE,
Self::INPUT_SIZE,
image::imageops::FilterType::Triangle,
);
// Convert to RGB and normalize to [-1, 1]
let rgb = resized.to_rgb8();
let (width, height) = rgb.dimensions();
// Create input tensor: [1, 3, 224, 224] in NCHW format
let mut input_data = Array4::<f32>::zeros((1, 3, height as usize, width as usize));
for y in 0..height {
for x in 0..width {
let pixel = rgb.get_pixel(x, y);
// Normalize from [0, 255] to [0, 1]
input_data[[0, 0, y as usize, x as usize]] = pixel[0] as f32 / 255.0; // R
input_data[[0, 1, y as usize, x as usize]] = pixel[1] as f32 / 255.0; // G
input_data[[0, 2, y as usize, x as usize]] = pixel[2] as f32 / 255.0; // B
}
}
// Flatten the array for ort input (ort 2.0 requires owned data)
let shape = [1_usize, 3, height as usize, width as usize];
let (input_vec, _offset) = input_data.into_raw_vec_and_offset();
// Create input tensor
let input_tensor = ort::value::Tensor::from_array((shape, input_vec))
.map_err(|e| Error::Model(format!("Failed to create input tensor: {}", e)))?;
// Run inference
let mut session = self
.session
.lock()
.map_err(|e| Error::Model(format!("Failed to lock session: {}", e)))?;
// Get the input name from the model (don't assume it's "input")
let input_name = session
.inputs()
.first()
.map(|i| i.name().to_string())
.unwrap_or_else(|| "input".to_string());
let outputs = session
.run(ort::inputs![input_name => input_tensor])
.map_err(|e| Error::Model(format!("DMHead inference failed: {}", e)))?;
// Extract output: [1, 3] containing [yaw, pitch, roll]
// Get the first output (model may use different output names)
let output_value = outputs
.iter()
.next()
.map(|(_, v)| v)
.ok_or_else(|| Error::Model("DMHead model returned no outputs".to_string()))?;
let (_shape, output_data) = output_value
.try_extract_tensor::<f32>()
.map_err(|e| Error::Model(format!("Failed to extract output tensor: {}", e)))?;
if output_data.len() < 3 {
return Err(Error::Model(format!(
"Expected 3 output values, got {}",
output_data.len()
)));
}
let yaw = output_data[0];
let pitch = output_data[1];
let roll = output_data[2];
Ok(HeadPose { yaw, pitch, roll })
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_input_normalization() {
// Test that normalization is correct
assert_eq!((0.0_f32 / 127.5) - 1.0, -1.0); // Black -> -1
assert_eq!((255.0_f32 / 127.5) - 1.0, 1.0); // White -> 1 (approx)
assert!((127.0_f32 / 127.5 - 1.0).abs() < 0.01); // Mid-gray -> ~0
}
}

10
src/models/mod.rs Normal file
View file

@ -0,0 +1,10 @@
//! ML model wrappers for face processing.
//!
//! This module provides wrappers for machine learning models used in the
//! face processing pipeline.
pub mod dlib_landmarks;
pub mod dmhead;
pub use dlib_landmarks::DlibLandmarks;
pub use dmhead::DMHeadModel;

View file

@ -83,6 +83,16 @@ impl Pipeline {
}
/// Create the default processing pipeline with standard steps.
///
/// Pipeline order:
/// 1. FaceResolutionStep - Validate face size from Immich metadata
/// 2. DecodeImageStep - Load and orient the image
/// 3. BrightnessStep - Filter by luminance
/// 4. CropFaceStep - Extract face region with padding
/// 5. HeadPoseStep - Filter non-frontal faces (DMHead)
/// 6. LandmarksStep - Detect 68 facial landmarks (dlib)
/// 7. AlignmentStep - Align face based on eye positions
/// 8. ResizeStep - Final resize to output size
pub fn with_default_steps() -> Self {
use steps::*;
@ -91,6 +101,9 @@ impl Pipeline {
pipeline.add_step(Box::new(DecodeImageStep));
pipeline.add_step(Box::new(BrightnessStep));
pipeline.add_step(Box::new(CropFaceStep));
pipeline.add_step(Box::new(HeadPoseStep));
pipeline.add_step(Box::new(LandmarksStep));
pipeline.add_step(Box::new(AlignmentStep));
pipeline.add_step(Box::new(ResizeStep));
pipeline
}
@ -233,6 +246,9 @@ mod tests {
assert!(ids.contains(&"decode"));
assert!(ids.contains(&"brightness"));
assert!(ids.contains(&"crop"));
assert!(ids.contains(&"head_pose"));
assert!(ids.contains(&"landmarks"));
assert!(ids.contains(&"alignment"));
assert!(ids.contains(&"resize"));
}
}

View file

@ -0,0 +1,275 @@
//! Eye-based face alignment step.
//!
//! Aligns faces based on eye positions to ensure consistent eye placement
//! across all images in the timelapse.
use crate::config::Config;
use crate::face_processing::types::Point;
use crate::pipeline::{PipelineContext, ProcessingStep, StepOutcome};
use async_trait::async_trait;
use image::{DynamicImage, GenericImageView, Rgb, RgbImage};
use imageproc::geometric_transformations::{rotate_about_center, Interpolation};
/// Aligns faces based on eye positions.
///
/// This step:
/// 1. Retrieves landmarks from ctx.computed["landmarks"]
/// 2. Calculates rotation angle from eye positions
/// 3. Applies affine transformation to align eyes horizontally
/// 4. Scales and crops to position eyes at configured positions
pub struct AlignmentStep;
#[async_trait]
impl ProcessingStep for AlignmentStep {
fn id(&self) -> &'static str {
"alignment"
}
fn name(&self) -> &'static str {
"Alignment"
}
async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome {
// Skip if alignment is disabled
if !config.processing.alignment.enabled {
return StepOutcome::Continue(ctx);
}
// Get landmarks from previous step
let landmarks: crate::face_processing::types::Landmarks = match ctx
.get_computed("landmarks")
.and_then(|v| v.as_landmarks())
{
Some(l) => l.clone(),
None => {
// If landmarks aren't available, skip alignment but continue
tracing::warn!("Landmarks not available, skipping alignment");
return StepOutcome::Continue(ctx);
}
};
let image = match ctx.image.take() {
Some(img) => img,
None => {
return StepOutcome::Error("No image available for alignment".to_string());
}
};
let (width, height) = image.dimensions();
let output_size = config.processing.output.size;
// Get eye centers
let left_eye = landmarks.left_eye_center();
let right_eye = landmarks.right_eye_center();
// Calculate rotation angle to make eyes horizontal
let angle = landmarks.eye_rotation_angle();
// Calculate current inter-eye distance
let current_eye_dist = landmarks.inter_eye_distance();
// Target inter-eye distance based on config (as fraction of output width)
let target_eye_dist = output_size as f32 * config.processing.alignment.inter_eye_distance;
// Calculate scale factor
let scale = target_eye_dist / current_eye_dist;
// Target eye positions
let target_eye_y = output_size as f32 * config.processing.alignment.eye_y_position;
let target_left_eye_x = (output_size as f32 - target_eye_dist) / 2.0;
let _target_right_eye_x = target_left_eye_x + target_eye_dist;
// Eye center (midpoint between eyes)
let eye_center = Point::new(
(left_eye.x + right_eye.x) / 2.0,
(left_eye.y + right_eye.y) / 2.0,
);
// First, rotate the image to make eyes horizontal
let rgb = image.to_rgb8();
let rotated = rotate_about_center(
&rgb,
-angle, // Negative because we want to counter-rotate
Interpolation::Bilinear,
Rgb([0, 0, 0]), // Black background for rotated areas
);
// After rotation, the eye center moves. Calculate new position.
// For small angles, we can approximate that the center stays roughly the same
// For more accuracy, we'd need to transform the point through the rotation
// Calculate the new eye center after rotation
let cos_a = angle.cos();
let sin_a = angle.sin();
let cx = width as f32 / 2.0;
let cy = height as f32 / 2.0;
// Rotate eye_center around image center
let dx = eye_center.x - cx;
let dy = eye_center.y - cy;
let rotated_eye_center = Point::new(
cx + dx * cos_a + dy * sin_a,
cy - dx * sin_a + dy * cos_a,
);
// Now calculate crop region to achieve the desired scale and positioning
// We want the eye center at (output_size/2, target_eye_y)
let target_center_x = output_size as f32 / 2.0;
let _target_center_y = target_eye_y;
// Calculate crop region in the rotated image
// The crop should be (output_size / scale) pixels, centered appropriately
let crop_size = (output_size as f32 / scale) as u32;
// Crop center in source image (accounting for where we want eyes to end up)
let crop_center_x = rotated_eye_center.x - (target_center_x - output_size as f32 / 2.0) / scale;
let crop_center_y = rotated_eye_center.y + (target_eye_y - output_size as f32 / 2.0) / scale;
// Calculate crop bounds
let crop_x = (crop_center_x - crop_size as f32 / 2.0).max(0.0) as u32;
let crop_y = (crop_center_y - crop_size as f32 / 2.0).max(0.0) as u32;
// Clamp to image bounds
let (rot_width, rot_height) = (rotated.width(), rotated.height());
let crop_x = crop_x.min(rot_width.saturating_sub(crop_size));
let crop_y = crop_y.min(rot_height.saturating_sub(crop_size));
let actual_crop_size = crop_size.min(rot_width - crop_x).min(rot_height - crop_y);
// Crop and resize
let rotated_dyn = DynamicImage::ImageRgb8(rotated);
let cropped = rotated_dyn.crop_imm(crop_x, crop_y, actual_crop_size, actual_crop_size);
let aligned = cropped.resize_exact(
output_size,
output_size,
image::imageops::FilterType::Lanczos3,
);
ctx.image = Some(aligned);
tracing::trace!(
"Aligned: rotation={:.2}deg, scale={:.2}, crop={}x{} at ({},{})",
angle.to_degrees(),
scale,
actual_crop_size,
actual_crop_size,
crop_x,
crop_y
);
StepOutcome::Continue(ctx)
}
fn debug_visualize(&self, ctx: &PipelineContext) -> Option<DynamicImage> {
// Get landmarks for visualization
let landmarks: &crate::face_processing::types::Landmarks =
ctx.get_computed("landmarks").and_then(|v| v.as_landmarks())?;
let image = ctx.image.as_ref()?;
let mut debug_img = image.to_rgb8();
let (width, height) = (debug_img.width(), debug_img.height());
// Draw target eye positions
let output_size = width; // Assuming square output
let eye_y = (output_size as f32 * 0.35) as u32; // Default eye_y_position
// Draw horizontal line at target eye Y position
for x in 0..width {
if eye_y < height {
debug_img.put_pixel(x, eye_y, Rgb([0, 255, 0]));
}
}
// Draw vertical lines at target eye X positions (assuming 0.3 inter_eye_distance)
let inter_eye = (output_size as f32 * 0.3) as u32;
let left_x = (width - inter_eye) / 2;
let right_x = left_x + inter_eye;
for y in 0..height {
if left_x < width {
debug_img.put_pixel(left_x, y, Rgb([0, 255, 0]));
}
if right_x < width {
debug_img.put_pixel(right_x, y, Rgb([0, 255, 0]));
}
}
// Draw actual eye positions
let left_eye = landmarks.left_eye_center();
let right_eye = landmarks.right_eye_center();
draw_marker(&mut debug_img, left_eye.x as u32, left_eye.y as u32, Rgb([255, 0, 0]));
draw_marker(&mut debug_img, right_eye.x as u32, right_eye.y as u32, Rgb([255, 0, 0]));
Some(DynamicImage::ImageRgb8(debug_img))
}
}
/// Draw a marker (small filled square) at the given position.
fn draw_marker(img: &mut RgbImage, x: u32, y: u32, color: Rgb<u8>) {
let (width, height) = (img.width(), img.height());
let size = 3;
for dy in 0..=size * 2 {
for dx in 0..=size * 2 {
let px = (x as i32 + dx as i32 - size as i32) as u32;
let py = (y as i32 + dy as i32 - size as i32) as u32;
if px < width && py < height {
img.put_pixel(px, py, color);
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::immich_api::FaceData;
fn make_test_ctx() -> PipelineContext {
let face_data = FaceData {
bounding_box_x1: 0.0,
bounding_box_y1: 0.0,
bounding_box_x2: 100.0,
bounding_box_y2: 100.0,
image_width: 100,
image_height: 100,
};
PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data)
}
#[tokio::test]
async fn test_disabled_skips_alignment() {
let step = AlignmentStep;
let ctx = make_test_ctx();
let mut config = Config::default();
config.processing.alignment.enabled = false;
// Create a dummy image
let img = DynamicImage::ImageRgb8(RgbImage::new(100, 100));
let ctx = ctx.with_image(img);
match step.execute(ctx, &config).await {
StepOutcome::Continue(new_ctx) => {
assert!(new_ctx.image.is_some());
}
other => panic!("Expected Continue when disabled, got {:?}", other),
}
}
#[tokio::test]
async fn test_no_landmarks_continues() {
let step = AlignmentStep;
let ctx = make_test_ctx();
let mut config = Config::default();
config.processing.alignment.enabled = true;
// Create a dummy image but no landmarks
let img = DynamicImage::ImageRgb8(RgbImage::new(100, 100));
let ctx = ctx.with_image(img);
match step.execute(ctx, &config).await {
StepOutcome::Continue(_) => {} // Expected - continues without alignment
other => panic!("Expected Continue without landmarks, got {:?}", other),
}
}
}

View file

@ -0,0 +1,212 @@
//! Head pose estimation step.
//!
//! Uses the DMHead ONNX model to estimate head pose (yaw, pitch, roll) and
//! filter out non-front-facing faces.
use crate::config::Config;
use crate::models::DMHeadModel;
use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome};
use async_trait::async_trait;
use image::{DynamicImage, Rgb};
/// Estimates head pose and filters non-frontal faces.
///
/// This step:
/// 1. Runs DMHead inference on the cropped face image
/// 2. Stores the HeadPose result in ctx.computed["head_pose"]
/// 3. Skips if any angle exceeds configured thresholds
pub struct HeadPoseStep;
#[async_trait]
impl ProcessingStep for HeadPoseStep {
fn id(&self) -> &'static str {
"head_pose"
}
fn name(&self) -> &'static str {
"Head Pose"
}
async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome {
// Skip if head pose filtering is disabled
if !config.processing.head_pose.enabled {
return StepOutcome::Continue(ctx);
}
let image = match &ctx.image {
Some(img) => img,
None => {
return StepOutcome::Error("No image available for head pose estimation".to_string());
}
};
// Load the DMHead model
let model = match DMHeadModel::global() {
Ok(m) => m,
Err(e) => {
// If model isn't available, skip this step with a warning
tracing::warn!("DMHead model not available, skipping head pose check: {}", e);
return StepOutcome::Continue(ctx);
}
};
// Run inference
let pose = match model.estimate(image) {
Ok(p) => p,
Err(e) => {
return StepOutcome::Error(format!("Head pose estimation failed: {}", e));
}
};
// Store pose in computed values
ctx.set_computed("head_pose", ComputedValue::HeadPose(pose));
// Check against thresholds
let head_pose_config = &config.processing.head_pose;
tracing::debug!(
"Head pose detected: yaw={:.1}°, pitch={:.1}°, roll={:.1}°",
pose.yaw,
pose.pitch,
pose.roll
);
if pose.yaw.abs() > head_pose_config.max_yaw {
return StepOutcome::Skip {
reason: "head_turned".to_string(),
detail: Some(format!(
"Yaw {:.1}° exceeds threshold {:.1}°",
pose.yaw, head_pose_config.max_yaw
)),
};
}
if pose.pitch.abs() > head_pose_config.max_pitch {
return StepOutcome::Skip {
reason: "head_turned".to_string(),
detail: Some(format!(
"Pitch {:.1}° exceeds threshold {:.1}°",
pose.pitch, head_pose_config.max_pitch
)),
};
}
if pose.roll.abs() > head_pose_config.max_roll {
return StepOutcome::Skip {
reason: "head_turned".to_string(),
detail: Some(format!(
"Roll {:.1}° exceeds threshold {:.1}°",
pose.roll, head_pose_config.max_roll
)),
};
}
StepOutcome::Continue(ctx)
}
fn debug_visualize(&self, ctx: &PipelineContext) -> Option<DynamicImage> {
// Get head pose from computed values
let pose = ctx
.get_computed("head_pose")
.and_then(|v| v.as_head_pose())?;
// Get the current image to draw on
let image = ctx.image.as_ref()?;
let rgb = image.to_rgb8();
let (width, height) = (rgb.width(), rgb.height());
// Create a copy for visualization
let mut debug_img = rgb.clone();
// Draw pose info as text overlay
// For simplicity, we'll draw colored bars indicating pose angles
// Green = within range, Red = out of range
// Draw yaw indicator (horizontal bar at top)
let yaw_pos = ((pose.yaw / 90.0 + 1.0) / 2.0 * width as f32) as u32;
let yaw_pos = yaw_pos.min(width - 1);
for x in 0..width {
let color = if x == yaw_pos {
Rgb([255, 255, 0]) // Yellow marker
} else if x == width / 2 {
Rgb([0, 255, 0]) // Green center
} else {
Rgb([50, 50, 50]) // Dark background
};
for y in 0..5 {
if y < height {
debug_img.put_pixel(x, y, color);
}
}
}
// Draw pitch indicator (vertical bar on left)
let pitch_pos = ((pose.pitch / 90.0 + 1.0) / 2.0 * height as f32) as u32;
let pitch_pos = pitch_pos.min(height - 1);
for y in 0..height {
let color = if y == pitch_pos {
Rgb([255, 255, 0]) // Yellow marker
} else if y == height / 2 {
Rgb([0, 255, 0]) // Green center
} else {
Rgb([50, 50, 50]) // Dark background
};
for x in 0..5 {
debug_img.put_pixel(x, y, color);
}
}
Some(DynamicImage::ImageRgb8(debug_img))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::immich_api::FaceData;
use image::RgbImage;
fn make_test_ctx() -> PipelineContext {
let face_data = FaceData {
bounding_box_x1: 0.0,
bounding_box_y1: 0.0,
bounding_box_x2: 100.0,
bounding_box_y2: 100.0,
image_width: 100,
image_height: 100,
};
PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data)
}
#[tokio::test]
async fn test_disabled_skips_check() {
let step = HeadPoseStep;
let ctx = make_test_ctx();
let mut config = Config::default();
config.processing.head_pose.enabled = false;
// Create a dummy image
let img = DynamicImage::ImageRgb8(RgbImage::new(100, 100));
let ctx = ctx.with_image(img);
match step.execute(ctx, &config).await {
StepOutcome::Continue(_) => {} // Expected
other => panic!("Expected Continue when disabled, got {:?}", other),
}
}
#[tokio::test]
async fn test_no_image_error() {
let step = HeadPoseStep;
let ctx = make_test_ctx();
let mut config = Config::default();
config.processing.head_pose.enabled = true;
match step.execute(ctx, &config).await {
StepOutcome::Error(msg) => {
assert!(msg.contains("No image"));
}
other => panic!("Expected Error, got {:?}", other),
}
}
}

View file

@ -0,0 +1,237 @@
//! Facial landmark detection step.
//!
//! Uses dlib to detect 68 facial landmarks for alignment and eye filtering.
use crate::config::Config;
use crate::face_processing::types::Landmarks;
use crate::models::DlibLandmarks;
use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome};
use async_trait::async_trait;
use image::{DynamicImage, Rgb, RgbImage};
use tokio::task;
/// Detects facial landmarks and optionally filters based on eye aspect ratio.
///
/// This step:
/// 1. Uses dlib to detect faces and 68 landmarks
/// 2. Stores Landmarks in ctx.computed["landmarks"]
/// 3. Computes EAR and stores in ctx.computed["ear"]
/// 4. Optionally skips if EAR is below threshold (eyes closed)
pub struct LandmarksStep;
#[async_trait]
impl ProcessingStep for LandmarksStep {
fn id(&self) -> &'static str {
"landmarks"
}
fn name(&self) -> &'static str {
"Landmarks"
}
async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome {
// We always need landmarks if alignment is enabled, even if eye filter is disabled
let need_landmarks =
config.processing.alignment.enabled || config.processing.eye_filter.enabled;
if !need_landmarks {
return StepOutcome::Continue(ctx);
}
let image = match &ctx.image {
Some(img) => img,
None => {
return StepOutcome::Error(
"No image available for landmark detection".to_string(),
);
}
};
// Get the global landmark predictor (loaded once, reused for all images)
let dlib = match DlibLandmarks::global() {
Ok(d) => d,
Err(e) => {
// If model isn't available, skip this step with a warning
tracing::warn!("Dlib landmarks model not available: {}", e);
return StepOutcome::Skip {
reason: "landmarks_failed".to_string(),
detail: Some(e.to_string()),
};
}
};
// Convert to RGB for dlib
let rgb = image.to_rgb8();
let (width, height) = (rgb.width() as usize, rgb.height() as usize);
let pixels = rgb.into_raw();
// Run dlib operations in a blocking thread to avoid dropping in async context
let landmarks_result = task::spawn_blocking(move || -> Result<Landmarks, String> {
dlib.detect_landmarks(width, height, &pixels)
.map_err(|e| e.to_string())
})
.await;
let landmarks = match landmarks_result {
Ok(Ok(l)) => l,
Ok(Err(e)) => {
return StepOutcome::Skip {
reason: "landmarks_failed".to_string(),
detail: Some(e),
};
}
Err(e) => {
return StepOutcome::Error(format!("Landmark detection task failed: {}", e));
}
};
// Compute and store EAR
let ear = landmarks.eye_aspect_ratio();
let avg_ear = (ear.left + ear.right) / 2.0;
ctx.set_computed("ear", ComputedValue::Float(avg_ear));
// Store landmarks
ctx.set_computed("landmarks", ComputedValue::Landmarks(Box::new(landmarks)));
// Check eye filter if enabled
if config.processing.eye_filter.enabled {
let min_ear = config.processing.eye_filter.min_ear;
if avg_ear < min_ear {
return StepOutcome::Skip {
reason: "eyes_closed".to_string(),
detail: Some(format!("EAR {:.3} below threshold {:.3}", avg_ear, min_ear)),
};
}
}
tracing::trace!(
"Landmarks detected: EAR left={:.3}, right={:.3}, avg={:.3}",
ear.left,
ear.right,
avg_ear
);
StepOutcome::Continue(ctx)
}
fn debug_visualize(&self, ctx: &PipelineContext) -> Option<DynamicImage> {
// Get landmarks from computed values
let landmarks: &Landmarks = ctx
.get_computed("landmarks")
.and_then(|v| v.as_landmarks())?;
// Get the current image to draw on
let image = ctx.image.as_ref()?;
let mut debug_img = image.to_rgb8();
// Draw all 68 landmark points
let points = landmarks.points();
for (i, point) in points.iter().enumerate() {
let x = point.x as u32;
let y = point.y as u32;
// Color-code different facial regions
let color = match i {
0..=16 => Rgb([255, 0, 0]), // Jaw (red)
17..=21 => Rgb([0, 255, 0]), // Left eyebrow (green)
22..=26 => Rgb([0, 255, 0]), // Right eyebrow (green)
27..=35 => Rgb([0, 0, 255]), // Nose (blue)
36..=41 => Rgb([255, 255, 0]), // Left eye (yellow)
42..=47 => Rgb([255, 255, 0]), // Right eye (yellow)
48..=67 => Rgb([255, 0, 255]), // Mouth (magenta)
_ => Rgb([255, 255, 255]), // Other (white)
};
// Draw a small cross at each point
draw_cross(&mut debug_img, x, y, color);
}
// Draw eye centers
let left_eye = landmarks.left_eye_center();
let right_eye = landmarks.right_eye_center();
draw_cross(
&mut debug_img,
left_eye.x as u32,
left_eye.y as u32,
Rgb([0, 255, 255]),
);
draw_cross(
&mut debug_img,
right_eye.x as u32,
right_eye.y as u32,
Rgb([0, 255, 255]),
);
Some(DynamicImage::ImageRgb8(debug_img))
}
}
/// Draw a small cross at the given position.
fn draw_cross(img: &mut RgbImage, x: u32, y: u32, color: Rgb<u8>) {
let (width, height) = (img.width(), img.height());
let size = 2;
for dx in 0..=size * 2 {
let px = (x as i32 + dx as i32 - size as i32) as u32;
if px < width {
img.put_pixel(px, y, color);
}
}
for dy in 0..=size * 2 {
let py = (y as i32 + dy as i32 - size as i32) as u32;
if py < height {
img.put_pixel(x, py, color);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::immich_api::FaceData;
fn make_test_ctx() -> PipelineContext {
let face_data = FaceData {
bounding_box_x1: 0.0,
bounding_box_y1: 0.0,
bounding_box_x2: 100.0,
bounding_box_y2: 100.0,
image_width: 100,
image_height: 100,
};
PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data)
}
#[tokio::test]
async fn test_disabled_skips_check() {
let step = LandmarksStep;
let ctx = make_test_ctx();
let mut config = Config::default();
config.processing.alignment.enabled = false;
config.processing.eye_filter.enabled = false;
// Create a dummy image
let img = DynamicImage::ImageRgb8(RgbImage::new(100, 100));
let ctx = ctx.with_image(img);
match step.execute(ctx, &config).await {
StepOutcome::Continue(_) => {} // Expected
other => panic!("Expected Continue when disabled, got {:?}", other),
}
}
#[tokio::test]
async fn test_no_image_error() {
let step = LandmarksStep;
let ctx = make_test_ctx();
let mut config = Config::default();
config.processing.alignment.enabled = true;
match step.execute(ctx, &config).await {
StepOutcome::Error(msg) => {
assert!(msg.contains("No image"));
}
other => panic!("Expected Error, got {:?}", other),
}
}
}

View file

@ -3,14 +3,20 @@
//! Each step implements the `ProcessingStep` trait and performs a specific
//! operation in the image processing pipeline.
mod face_resolution;
mod decode;
mod alignment;
mod brightness;
mod crop;
mod decode;
mod face_resolution;
mod head_pose;
mod landmarks;
mod resize;
pub use face_resolution::FaceResolutionStep;
pub use decode::DecodeImageStep;
pub use alignment::AlignmentStep;
pub use brightness::BrightnessStep;
pub use crop::CropFaceStep;
pub use decode::DecodeImageStep;
pub use face_resolution::FaceResolutionStep;
pub use head_pose::HeadPoseStep;
pub use landmarks::LandmarksStep;
pub use resize::ResizeStep;

View file

@ -4,6 +4,7 @@
//! extensible image processing pipelines.
use crate::config::Config;
use crate::face_processing::types::{HeadPose, Landmarks};
use crate::immich_api::FaceData;
use async_trait::async_trait;
use bytes::Bytes;
@ -33,6 +34,10 @@ pub enum ComputedValue {
Bool(bool),
/// A string value.
String(String),
/// Head pose estimation result (yaw, pitch, roll).
HeadPose(HeadPose),
/// Facial landmarks (68 points).
Landmarks(Box<Landmarks>),
}
impl ComputedValue {
@ -67,6 +72,22 @@ impl ComputedValue {
_ => None,
}
}
/// Get as HeadPose if this is a HeadPose variant.
pub fn as_head_pose(&self) -> Option<&HeadPose> {
match self {
ComputedValue::HeadPose(v) => Some(v),
_ => None,
}
}
/// Get as Landmarks if this is a Landmarks variant.
pub fn as_landmarks(&self) -> Option<&Landmarks> {
match self {
ComputedValue::Landmarks(v) => Some(v),
_ => None,
}
}
}
/// Context passed through the pipeline, carrying data between steps.