384 lines
12 KiB
Rust
384 lines
12 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};
|
|
|
|
/// 3D point for cube vertices
|
|
#[derive(Clone, Copy)]
|
|
struct Point3D {
|
|
x: f32,
|
|
y: f32,
|
|
z: f32,
|
|
}
|
|
|
|
impl Point3D {
|
|
fn new(x: f32, y: f32, z: f32) -> Self {
|
|
Self { x, y, z }
|
|
}
|
|
|
|
/// Rotate around Y axis (yaw)
|
|
fn rotate_y(self, angle_deg: f32) -> Self {
|
|
let rad = angle_deg.to_radians();
|
|
let cos = rad.cos();
|
|
let sin = rad.sin();
|
|
Self {
|
|
x: self.x * cos + self.z * sin,
|
|
y: self.y,
|
|
z: -self.x * sin + self.z * cos,
|
|
}
|
|
}
|
|
|
|
/// Rotate around X axis (pitch)
|
|
fn rotate_x(self, angle_deg: f32) -> Self {
|
|
let rad = angle_deg.to_radians();
|
|
let cos = rad.cos();
|
|
let sin = rad.sin();
|
|
Self {
|
|
x: self.x,
|
|
y: self.y * cos - self.z * sin,
|
|
z: self.y * sin + self.z * cos,
|
|
}
|
|
}
|
|
|
|
/// Rotate around Z axis (roll)
|
|
fn rotate_z(self, angle_deg: f32) -> Self {
|
|
let rad = angle_deg.to_radians();
|
|
let cos = rad.cos();
|
|
let sin = rad.sin();
|
|
Self {
|
|
x: self.x * cos - self.y * sin,
|
|
y: self.x * sin + self.y * cos,
|
|
z: self.z,
|
|
}
|
|
}
|
|
|
|
/// Project 3D point to 2D using perspective projection
|
|
fn project(self, cx: f32, cy: f32, focal_length: f32) -> (i32, i32) {
|
|
let scale = focal_length / (focal_length + self.z);
|
|
let x2d = cx + self.x * scale;
|
|
let y2d = cy + self.y * scale;
|
|
(x2d as i32, y2d as i32)
|
|
}
|
|
}
|
|
|
|
/// Draw 3D pose axes (RGB = XYZ) from face center, rotated by head pose.
|
|
/// This is the standard visualization for head pose estimation.
|
|
fn draw_pose_axes(
|
|
img: &mut RgbImage,
|
|
cx: f32,
|
|
cy: f32,
|
|
axis_length: f32,
|
|
yaw: f32,
|
|
pitch: f32,
|
|
roll: f32,
|
|
) {
|
|
// Define axis endpoints (origin at 0,0,0)
|
|
// X axis (red) - points right
|
|
// Y axis (green) - points down (image coordinates)
|
|
// Z axis (blue) - points out of screen (towards camera)
|
|
let axes = [
|
|
(Point3D::new(axis_length, 0.0, 0.0), Rgb([255, 0, 0])), // X - red
|
|
(Point3D::new(0.0, axis_length, 0.0), Rgb([0, 255, 0])), // Y - green
|
|
(Point3D::new(0.0, 0.0, -axis_length), Rgb([0, 0, 255])), // Z - blue (negative = towards camera)
|
|
];
|
|
|
|
// Negate roll and pitch to convert from model convention to image coordinates
|
|
// (image Y-axis points down, model assumes Y-axis points up)
|
|
let focal_length = axis_length * 2.0;
|
|
let origin = Point3D::new(0.0, 0.0, 0.0)
|
|
.rotate_y(yaw)
|
|
.rotate_x(-pitch)
|
|
.rotate_z(-roll)
|
|
.project(cx, cy, focal_length);
|
|
|
|
for (endpoint, color) in axes {
|
|
let rotated = endpoint
|
|
.rotate_y(yaw)
|
|
.rotate_x(-pitch)
|
|
.rotate_z(-roll)
|
|
.project(cx, cy, focal_length);
|
|
|
|
draw_line(img, origin.0, origin.1, rotated.0, rotated.1, color);
|
|
}
|
|
}
|
|
|
|
/// Draw an axis-aligned bounding box (no rotation applied).
|
|
fn draw_axis_aligned_rect(img: &mut RgbImage, x1: i32, y1: i32, x2: i32, y2: i32, color: Rgb<u8>) {
|
|
draw_line(img, x1, y1, x2, y1, color); // top
|
|
draw_line(img, x2, y1, x2, y2, color); // right
|
|
draw_line(img, x2, y2, x1, y2, color); // bottom
|
|
draw_line(img, x1, y2, x1, y1, color); // left
|
|
}
|
|
|
|
/// 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;
|
|
}
|
|
}
|
|
}
|
|
|
|
/// 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 {
|
|
let image: &DynamicImage = 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 square face crop if we have the face rectangle
|
|
// Square crop prevents aspect ratio distortion when DMHead resizes to 224x224
|
|
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_with_margin = x.saturating_sub(margin);
|
|
let y_with_margin = y.saturating_sub(margin);
|
|
let w_with_margin = (w + margin * 2).min(img_w - x_with_margin);
|
|
let h_with_margin = (h + margin * 2).min(img_h - y_with_margin);
|
|
|
|
// Make the crop square by using the larger dimension
|
|
let size = w_with_margin.max(h_with_margin);
|
|
|
|
// Center the square crop around the face
|
|
let center_x = x_with_margin + w_with_margin / 2;
|
|
let center_y = y_with_margin + h_with_margin / 2;
|
|
|
|
let square_x = center_x.saturating_sub(size / 2);
|
|
let square_y = center_y.saturating_sub(size / 2);
|
|
|
|
// Ensure the square crop doesn't go out of bounds
|
|
let final_x = square_x.min(img_w.saturating_sub(size));
|
|
let final_y = square_y.min(img_h.saturating_sub(size));
|
|
let final_size = size.min(img_w - final_x).min(img_h - final_y);
|
|
|
|
image.crop_imm(final_x, final_y, final_size, final_size)
|
|
} 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;
|
|
|
|
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
|
|
)),
|
|
};
|
|
}
|
|
|
|
// Note: Roll is not checked for pass/fail - only used for visualization
|
|
// Roll (head tilt) is less important for timelapse alignment since
|
|
// the alignment step can handle rotated faces
|
|
|
|
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();
|
|
|
|
// Get face rectangle if available
|
|
let face_rect = ctx
|
|
.get_computed(computed_keys::FACE_RECT)
|
|
.and_then(|v| v.as_face_rect());
|
|
|
|
// Draw axis-aligned face bounding box and 3D pose axes
|
|
if let Some(rect) = face_rect {
|
|
let x1 = rect.x1 as i32;
|
|
let y1 = rect.y1 as i32;
|
|
let x2 = rect.x2 as i32;
|
|
let y2 = rect.y2 as i32;
|
|
|
|
// Draw axis-aligned bounding box in cyan
|
|
draw_axis_aligned_rect(&mut debug_img, x1, y1, x2, y2, Rgb([0, 255, 255]));
|
|
|
|
// Draw 3D pose axes from face center
|
|
let cx = (x1 + x2) as f32 / 2.0;
|
|
let cy = (y1 + y2) as f32 / 2.0;
|
|
let axis_length = ((x2 - x1).max(y2 - y1) as f32) * 0.6;
|
|
draw_pose_axes(
|
|
&mut debug_img,
|
|
cx,
|
|
cy,
|
|
axis_length,
|
|
pose.yaw,
|
|
pose.pitch,
|
|
pose.roll,
|
|
);
|
|
}
|
|
|
|
// 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
|
|
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))
|
|
}
|
|
}
|
|
|
|
#[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_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),
|
|
}
|
|
}
|
|
}
|