wip face landmark detection and alignment
This commit is contained in:
parent
291989cf4a
commit
41175c1c1b
17 changed files with 1450 additions and 20 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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}>
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ export default defineConfig({
|
|||
'/api': {
|
||||
target: 'http://localhost:5000',
|
||||
changeOrigin: true,
|
||||
ws: true, // Enable WebSocket proxying
|
||||
},
|
||||
},
|
||||
},
|
||||
|
|
|
|||
146
src/config.rs
146
src/config.rs
|
|
@ -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(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
127
src/models/dlib_landmarks.rs
Normal file
127
src/models/dlib_landmarks.rs
Normal 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
166
src/models/dmhead.rs
Normal 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
10
src/models/mod.rs
Normal 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;
|
||||
|
|
@ -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"));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
275
src/pipeline/steps/alignment.rs
Normal file
275
src/pipeline/steps/alignment.rs
Normal 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),
|
||||
}
|
||||
}
|
||||
}
|
||||
212
src/pipeline/steps/head_pose.rs
Normal file
212
src/pipeline/steps/head_pose.rs
Normal 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),
|
||||
}
|
||||
}
|
||||
}
|
||||
237
src/pipeline/steps/landmarks.rs
Normal file
237
src/pipeline/steps/landmarks.rs
Normal 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),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Reference in a new issue