298 lines
10 KiB
Rust
298 lines
10 KiB
Rust
//! 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::{computed_keys, draw_simple_text, ComputedValue, PipelineContext, ProcessingStep, StepOutcome};
|
|
use async_trait::async_trait;
|
|
use image::{DynamicImage, GenericImageView, Rgb, RgbImage};
|
|
|
|
/// 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.require_image("head pose estimation") {
|
|
Ok(img) => img,
|
|
Err(e) => return StepOutcome::Error { ctx, error: e },
|
|
};
|
|
|
|
// 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);
|
|
}
|
|
};
|
|
|
|
// Extract a tighter face crop if we have the face rectangle
|
|
// DMHead works better with tight face crops centered on the face
|
|
let face_image: DynamicImage = if let Some(face_rect) = ctx
|
|
.get_computed(computed_keys::FACE_RECT)
|
|
.and_then(|v| v.as_face_rect())
|
|
{
|
|
// Use the face rectangle to extract a tighter crop
|
|
let (img_w, img_h) = image.dimensions();
|
|
let x = (face_rect.x1 as u32).min(img_w.saturating_sub(1));
|
|
let y = (face_rect.y1 as u32).min(img_h.saturating_sub(1));
|
|
let w = ((face_rect.x2 - face_rect.x1) as u32).min(img_w - x);
|
|
let h = ((face_rect.y2 - face_rect.y1) as u32).min(img_h - y);
|
|
|
|
if w > 10 && h > 10 {
|
|
// Add a small margin around the face for better model performance
|
|
let margin = (w.max(h) / 4).min(20);
|
|
let x = x.saturating_sub(margin);
|
|
let y = y.saturating_sub(margin);
|
|
let w = (w + margin * 2).min(img_w - x);
|
|
let h = (h + margin * 2).min(img_h - y);
|
|
|
|
image.crop_imm(x, y, w, h)
|
|
} else {
|
|
// Face rect too small, use full image
|
|
image.clone()
|
|
}
|
|
} else {
|
|
// No face rect available, use full image
|
|
image.clone()
|
|
};
|
|
|
|
// Run inference on the face crop
|
|
let pose = match model.estimate(&face_image) {
|
|
Ok(p) => p,
|
|
Err(e) => {
|
|
return StepOutcome::Error {
|
|
ctx,
|
|
error: format!("Head pose estimation failed: {}", e),
|
|
};
|
|
}
|
|
};
|
|
|
|
// Store pose in computed values
|
|
ctx.set_computed(computed_keys::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 {
|
|
ctx,
|
|
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 {
|
|
ctx,
|
|
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 {
|
|
ctx,
|
|
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, _config: &Config) -> Option<DynamicImage> {
|
|
// Get head pose from computed values
|
|
let pose = ctx
|
|
.get_computed(computed_keys::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 a center crosshair
|
|
let cx = width / 2;
|
|
let cy = height / 2;
|
|
let crosshair_size = 20u32;
|
|
|
|
// Horizontal line
|
|
for x in cx.saturating_sub(crosshair_size)..=(cx + crosshair_size).min(width - 1) {
|
|
debug_img.put_pixel(x, cy, Rgb([0, 255, 0]));
|
|
}
|
|
// Vertical line
|
|
for y in cy.saturating_sub(crosshair_size)..=(cy + crosshair_size).min(height - 1) {
|
|
debug_img.put_pixel(cx, y, Rgb([0, 255, 0]));
|
|
}
|
|
|
|
// Draw pose direction arrow from center
|
|
// Yaw rotates left/right, pitch rotates up/down
|
|
let arrow_len = 40.0_f32;
|
|
let yaw_rad = pose.yaw.to_radians();
|
|
let pitch_rad = pose.pitch.to_radians();
|
|
|
|
// Arrow endpoint based on yaw and pitch
|
|
let dx = (yaw_rad.sin() * arrow_len) as i32;
|
|
let dy = (-pitch_rad.sin() * arrow_len) as i32; // Negative because y increases downward
|
|
|
|
let ex = (cx as i32 + dx).clamp(0, width as i32 - 1) as u32;
|
|
let ey = (cy as i32 + dy).clamp(0, height as i32 - 1) as u32;
|
|
|
|
// Draw arrow line using Bresenham's algorithm
|
|
draw_line(&mut debug_img, cx as i32, cy as i32, ex as i32, ey as i32, Rgb([255, 0, 0]));
|
|
|
|
// Draw roll indicator as a tilted line through center
|
|
let roll_rad = pose.roll.to_radians();
|
|
let roll_len = 30.0_f32;
|
|
let rx1 = (cx as f32 - roll_rad.cos() * roll_len) as u32;
|
|
let ry1 = (cy as f32 - roll_rad.sin() * roll_len) as u32;
|
|
let rx2 = (cx as f32 + roll_rad.cos() * roll_len) as u32;
|
|
let ry2 = (cy as f32 + roll_rad.sin() * roll_len) as u32;
|
|
draw_line(&mut debug_img, rx1 as i32, ry1 as i32, rx2 as i32, ry2 as i32, Rgb([0, 255, 255]));
|
|
|
|
// Draw text background bar at bottom for pose values
|
|
let bar_height = 20u32;
|
|
for y in height.saturating_sub(bar_height)..height {
|
|
for x in 0..width {
|
|
debug_img.put_pixel(x, y, Rgb([0, 0, 0]));
|
|
}
|
|
}
|
|
|
|
// Draw simple text representation of values using block characters
|
|
// Format: Y:-20 P:+29 R:-21
|
|
let text = format!(
|
|
"Y:{:+.0} P:{:+.0} R:{:+.0}",
|
|
pose.yaw, pose.pitch, pose.roll
|
|
);
|
|
draw_simple_text(&mut debug_img, 5, height - bar_height + 4, &text, Rgb([255, 255, 255]));
|
|
|
|
Some(DynamicImage::ImageRgb8(debug_img))
|
|
}
|
|
}
|
|
|
|
/// Draw a line using Bresenham's algorithm.
|
|
fn draw_line(img: &mut RgbImage, x0: i32, y0: i32, x1: i32, y1: i32, color: Rgb<u8>) {
|
|
let (width, height) = (img.width() as i32, img.height() as i32);
|
|
|
|
let dx = (x1 - x0).abs();
|
|
let dy = -(y1 - y0).abs();
|
|
let sx = if x0 < x1 { 1 } else { -1 };
|
|
let sy = if y0 < y1 { 1 } else { -1 };
|
|
let mut err = dx + dy;
|
|
|
|
let mut x = x0;
|
|
let mut y = y0;
|
|
|
|
loop {
|
|
if x >= 0 && x < width && y >= 0 && y < height {
|
|
img.put_pixel(x as u32, y as u32, color);
|
|
}
|
|
|
|
if x == x1 && y == y1 {
|
|
break;
|
|
}
|
|
|
|
let e2 = 2 * err;
|
|
if e2 >= dy {
|
|
err += dy;
|
|
x += sx;
|
|
}
|
|
if e2 <= dx {
|
|
err += dx;
|
|
y += sy;
|
|
}
|
|
}
|
|
}
|
|
|
|
#[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 { error, .. } => {
|
|
assert!(error.contains("No image"));
|
|
}
|
|
other => panic!("Expected Error, got {:?}", other),
|
|
}
|
|
}
|
|
}
|