From 7c6c37868711dd4c6d157392be39bf14afa2beb3 Mon Sep 17 00:00:00 2001 From: Arnaud_Cayrol Date: Mon, 2 Feb 2026 20:11:31 +0100 Subject: [PATCH] working on debug visualisation --- src/face_processing/crop.rs | 30 +++- src/job/mod.rs | 43 ++---- src/main.rs | 8 +- src/models/dlib_landmarks.rs | 53 +++++-- src/models/dmhead.rs | 20 +-- src/pipeline/mod.rs | 97 ++++++++---- src/pipeline/steps/brightness.rs | 116 ++++++++++++++- src/pipeline/steps/crop.rs | 13 +- src/pipeline/steps/decode.rs | 1 + src/pipeline/steps/face_resolution.rs | 3 +- src/pipeline/steps/head_pose.rs | 207 +++++++++++++++++++++----- src/pipeline/steps/landmarks.rs | 20 ++- src/pipeline/traits.rs | 49 +++++- src/video/ffmpeg.rs | 4 +- 14 files changed, 517 insertions(+), 147 deletions(-) diff --git a/src/face_processing/crop.rs b/src/face_processing/crop.rs index 9489fee..47c823f 100644 --- a/src/face_processing/crop.rs +++ b/src/face_processing/crop.rs @@ -3,12 +3,23 @@ //! Extracts and resizes face regions from images using bounding box data. use crate::error::{Error, Result}; +use crate::face_processing::types::BoundingBox; use crate::immich_api::FaceData; use image::imageops::FilterType; use image::{DynamicImage, GenericImageView}; +/// Result of cropping a face from an image. +pub struct CropResult { + /// The cropped image at full resolution. + pub cropped: DynamicImage, + /// The cropped image resized to output size. + pub resized: DynamicImage, + /// The face bounding box in crop coordinates. + pub face_rect: BoundingBox, +} + /// Crop and resize the face from an image using bounding box. -/// Returns (cropped_full_res, resized_final) for intermediate saving. +/// Returns CropResult containing cropped images and face rectangle in crop coordinates. /// /// This is a simplified version that just uses the bounding box. /// A full implementation would use facial landmarks for alignment. @@ -16,7 +27,7 @@ pub fn crop_face_with_intermediate( img: &DynamicImage, face_data: &FaceData, output_size: u32, -) -> Result<(DynamicImage, DynamicImage)> { +) -> Result { let (img_width, img_height) = img.dimensions(); // Scale bounding box from metadata dimensions to actual image dimensions. @@ -70,5 +81,18 @@ pub fn crop_face_with_intermediate( // Resize to output size let resized = cropped.resize_exact(output_size, output_size, FilterType::Lanczos3); - Ok((cropped, resized)) + // Calculate the face bounding box in crop coordinates + // These are the original face coordinates relative to the crop origin + let face_rect = BoundingBox { + x1: x1.saturating_sub(crop_x1) as f32, + y1: y1.saturating_sub(crop_y1) as f32, + x2: x2.saturating_sub(crop_x1).min(actual_crop_size) as f32, + y2: y2.saturating_sub(crop_y1).min(actual_crop_size) as f32, + }; + + Ok(CropResult { + cropped, + resized, + face_rect, + }) } diff --git a/src/job/mod.rs b/src/job/mod.rs index 3bc63e1..682bdb4 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -43,22 +43,14 @@ struct OutputDirs { debug: Option, } -/// Debug output directories for visualizing each processing stage. -/// Each folder contains images with debug overlays showing what happened at that step. -/// -/// Note: `landmarks` and `alignment` fields are placeholders for future features -/// (facial landmark detection and face alignment). They are intentionally unused -/// until those features are implemented. -#[derive(Debug, Clone, Default)] +/// Debug output base directory for visualizing processing stages. +/// Subdirectories are created on-demand by the pipeline for each step: +/// - `{step_id}/passed/` - Images that passed the step +/// - `{step_id}/failed/` - Images that failed/skipped at the step +#[derive(Debug, Clone)] struct DebugDirs { - /// Original image with bounding box and crop region overlayed - crop: Option, - /// Face with landmark points drawn (future: facial landmark detection) - #[allow(dead_code)] - landmarks: Option, - /// Before/after alignment visualization (future: face alignment) - #[allow(dead_code)] - alignment: Option, + /// Base directory for debug output (e.g., `output/PersonName/debug`) + base: PathBuf, } /// Run the complete processing pipeline. @@ -137,19 +129,8 @@ async fn run_job_inner( // Create debug directories if enabled let debug = if config.processing.output.keep_intermediates { let debug_base = person_dir.join("debug"); - - let crop_dir = debug_base.join("crop"); - tokio::fs::create_dir_all(&crop_dir).await?; - - // Future directories (created on-demand when those features are implemented): - // - debug_base.join("landmarks") - face with landmark points - // - debug_base.join("alignment") - before/after alignment - - Some(DebugDirs { - crop: Some(crop_dir), - landmarks: None, - alignment: None, - }) + tokio::fs::create_dir_all(&debug_base).await?; + Some(DebugDirs { base: debug_base }) } else { None }; @@ -499,11 +480,7 @@ async fn process_single_asset( ctx = ctx.with_bytes(image_bytes); // Determine debug directory - let debug_dir = if config.processing.output.keep_intermediates { - output_dirs.debug.as_ref().and_then(|d| d.crop.as_ref().map(|p| p.parent().unwrap().to_path_buf())) - } else { - None - }; + let debug_dir = output_dirs.debug.as_ref().map(|d| d.base.clone()); // Execute the pipeline let result = pipeline.execute( diff --git a/src/main.rs b/src/main.rs index 623cccf..2ee11d3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,7 +2,7 @@ //! //! Web server for creating selfie timelapses from Immich. -use immich_timelapse::{config::Config, web}; +use immich_timelapse::{config::Config, models::DlibLandmarks, web}; use std::net::SocketAddr; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; @@ -35,6 +35,12 @@ async fn main() -> anyhow::Result<()> { Err(e) => tracing::warn!("FFmpeg not available: {} - video compilation will fail", e), } + // Pre-load ML models to avoid loading during processing + match DlibLandmarks::init() { + Ok(_) => tracing::info!("Dlib landmarks model loaded"), + Err(e) => tracing::warn!("Dlib landmarks model not available: {} - landmark detection will be skipped", e), + } + // Create application state let state = web::AppState::new(config); diff --git a/src/models/dlib_landmarks.rs b/src/models/dlib_landmarks.rs index 6baf20a..fb11fb1 100644 --- a/src/models/dlib_landmarks.rs +++ b/src/models/dlib_landmarks.rs @@ -51,12 +51,25 @@ impl DlibLandmarks { .map_err(|e| Error::Model(e.to_string())) } + /// Eagerly initialize the model at startup. + /// + /// Call this during server startup to load the model before processing begins. + /// This ensures any model loading messages appear during startup rather than + /// during image processing. + pub fn init() -> Result<()> { + Self::global()?; + Ok(()) + } + /// 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) + /// * `face_rect` - Optional face bounding box (x1, y1, x2, y2) in image coordinates. + /// If provided, uses this rectangle for landmark detection. + /// If None, attempts to detect the face or uses the whole image. /// /// # Returns /// Landmarks struct containing the 68 facial landmark points, or an error @@ -66,6 +79,7 @@ impl DlibLandmarks { width: usize, height: usize, pixels: &[u8], + face_rect: Option<(i64, i64, i64, i64)>, ) -> Result { // Create image matrix for dlib let matrix = unsafe { ImageMatrix::new(width, height, pixels.as_ptr()) }; @@ -80,23 +94,30 @@ impl DlibLandmarks { .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() + // Determine the face rectangle to use + let rect = if let Some((x1, y1, x2, y2)) = face_rect { + // Use the provided face rectangle + Rectangle { + left: x1.max(0), + top: y1.max(0), + right: x2.min(width as i64), + bottom: y2.min(height as i64), + } } else { - face_rect + // No face rect provided - try to detect or use whole image + let faces = detector.face_locations(&matrix); + if !faces.is_empty() { + faces[0].clone() + } else { + // Fallback: use whole image with small margin + let margin = 5; + Rectangle { + left: margin, + top: margin, + right: (width as i64) - margin, + bottom: (height as i64) - margin, + } + } }; // Detect landmarks diff --git a/src/models/dmhead.rs b/src/models/dmhead.rs index 4d50e1a..b87bf94 100644 --- a/src/models/dmhead.rs +++ b/src/models/dmhead.rs @@ -95,10 +95,11 @@ impl DMHeadModel { 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 + // No normalization - DMHead expects raw [0, 255] pixel values as floats + // See: https://github.com/PINTO0309/DMHead/blob/main/demo_video.py + input_data[[0, 0, y as usize, x as usize]] = pixel[0] as f32; // R + input_data[[0, 1, y as usize, x as usize]] = pixel[1] as f32; // G + input_data[[0, 2, y as usize, x as usize]] = pixel[2] as f32; // B } } @@ -157,10 +158,11 @@ impl DMHeadModel { #[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 + fn test_input_format() { + // DMHead expects raw [0, 255] pixel values as floats, no normalization + // See: https://github.com/PINTO0309/DMHead/blob/main/demo_video.py + assert_eq!(0_u8 as f32, 0.0); + assert_eq!(255_u8 as f32, 255.0); + assert_eq!(128_u8 as f32, 128.0); } } diff --git a/src/pipeline/mod.rs b/src/pipeline/mod.rs index f9ea1ec..cfed085 100644 --- a/src/pipeline/mod.rs +++ b/src/pipeline/mod.rs @@ -38,13 +38,15 @@ pub enum PipelineResult { /// Timestamp for file naming. timestamp: String, /// Debug images if keep_intermediates was enabled. - debug_images: Vec<(String, DynamicImage)>, + debug_images: Vec<(String, DebugImage)>, }, /// Image was skipped. Skipped { asset_id: String, reason: String, detail: Option, + /// Debug images generated up to (and including) the failing step. + debug_images: Vec<(String, DebugImage)>, }, /// Processing error. Error { @@ -145,17 +147,24 @@ impl Pipeline { StepOutcome::Continue(new_ctx) => { ctx = new_ctx; - // Generate debug visualization if enabled + // Generate debug visualization if enabled (step passed) if config.processing.output.keep_intermediates { if let Some(debug_img) = step.debug_visualize(&ctx) { - ctx.add_debug_image(step.id(), debug_img); + ctx.add_debug_image(step.id(), debug_img, true); } } } - StepOutcome::Skip { reason, detail } => { + StepOutcome::Skip { mut ctx, reason, detail } => { // Update skip stats with the step ID as reason skip_stats.increment(&reason); + // Generate debug visualization for the failing step if enabled + if config.processing.output.keep_intermediates { + if let Some(debug_img) = step.debug_visualize(&ctx) { + ctx.add_debug_image(step.id(), debug_img, false); + } + } + tracing::debug!( "Asset {} skipped at step '{}': {} ({})", asset_id, @@ -164,10 +173,18 @@ impl Pipeline { detail.as_deref().unwrap_or("no detail") ); + // Collect and save debug images before returning + let debug_images: Vec<(String, DebugImage)> = ctx.debug_images.into_iter().collect(); + + if let Some(debug_base) = debug_dir { + save_debug_images(debug_base, &debug_images, ×tamp, &asset_id).await; + } + return PipelineResult::Skipped { asset_id, reason, detail, + debug_images, }; } StepOutcome::Error(error) => { @@ -194,34 +211,12 @@ impl Pipeline { } }; + // Collect debug images + let debug_images: Vec<(String, DebugImage)> = ctx.debug_images.into_iter().collect(); + // Save debug images if enabled - let debug_images: Vec<(String, DynamicImage)> = ctx.debug_images.into_iter().collect(); - if let Some(debug_base) = debug_dir { - for (step_id, debug_img) in &debug_images { - let step_dir = debug_base.join(step_id); - if let Err(e) = tokio::fs::create_dir_all(&step_dir).await { - tracing::warn!("Failed to create debug dir {}: {}", step_dir.display(), e); - continue; - } - - let filename = format!("{}_{}.jpg", timestamp, asset_id) - .chars() - .map(|c| if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' { c } else { '_' }) - .collect::(); - - let debug_path = step_dir.join(&filename); - let mut buffer = Cursor::new(Vec::new()); - - if let Err(e) = debug_img.write_to(&mut buffer, ImageFormat::Jpeg) { - tracing::warn!("Failed to encode debug image: {}", e); - continue; - } - - if let Err(e) = tokio::fs::write(&debug_path, buffer.into_inner()).await { - tracing::warn!("Failed to save debug image: {}", e); - } - } + save_debug_images(debug_base, &debug_images, ×tamp, &asset_id).await; } PipelineResult::Success { @@ -233,6 +228,48 @@ impl Pipeline { } } +/// Save debug images to disk with passed/failed subdirectories. +/// +/// Creates directory structure: +/// ```text +/// debug/{step_id}/passed/{filename}.jpg +/// debug/{step_id}/failed/{filename}.jpg +/// ``` +async fn save_debug_images( + debug_base: &PathBuf, + debug_images: &[(String, DebugImage)], + timestamp: &str, + asset_id: &str, +) { + for (step_id, debug_img) in debug_images { + // Determine subdirectory based on pass/fail status + let status_dir = if debug_img.passed { "passed" } else { "failed" }; + let step_dir = debug_base.join(step_id).join(status_dir); + + if let Err(e) = tokio::fs::create_dir_all(&step_dir).await { + tracing::warn!("Failed to create debug dir {}: {}", step_dir.display(), e); + continue; + } + + let filename = format!("{}_{}.jpg", timestamp, asset_id) + .chars() + .map(|c| if c.is_alphanumeric() || c == '-' || c == '_' || c == '.' { c } else { '_' }) + .collect::(); + + let debug_path = step_dir.join(&filename); + let mut buffer = Cursor::new(Vec::new()); + + if let Err(e) = debug_img.image.write_to(&mut buffer, ImageFormat::Jpeg) { + tracing::warn!("Failed to encode debug image: {}", e); + continue; + } + + if let Err(e) = tokio::fs::write(&debug_path, buffer.into_inner()).await { + tracing::warn!("Failed to save debug image: {}", e); + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/pipeline/steps/brightness.rs b/src/pipeline/steps/brightness.rs index 670b8b1..da6ac9f 100644 --- a/src/pipeline/steps/brightness.rs +++ b/src/pipeline/steps/brightness.rs @@ -5,7 +5,7 @@ use crate::config::Config; use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome}; use async_trait::async_trait; -use image::DynamicImage; +use image::{DynamicImage, Rgb, RgbImage}; /// Validates image brightness and skips images outside acceptable range. /// @@ -76,6 +76,7 @@ impl ProcessingStep for BrightnessStep { if brightness < step_config.min_brightness { return StepOutcome::Skip { + ctx, reason: "too_dark".to_string(), detail: Some(format!( "{:.2} (min: {:.2})", @@ -86,6 +87,7 @@ impl ProcessingStep for BrightnessStep { if brightness > step_config.max_brightness { return StepOutcome::Skip { + ctx, reason: "too_bright".to_string(), detail: Some(format!( "{:.2} (max: {:.2})", @@ -96,6 +98,118 @@ impl ProcessingStep for BrightnessStep { StepOutcome::Continue(ctx) } + + fn debug_visualize(&self, ctx: &PipelineContext) -> Option { + // Get brightness from computed values + let brightness = ctx + .get_computed("brightness") + .and_then(|v| v.as_float())?; + + // 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 horizontal brightness bar at the bottom + let bar_height = 20u32; + let bar_y = height.saturating_sub(bar_height); + let bar_width = (width as f32 * 0.8) as u32; + let bar_x = (width - bar_width) / 2; + + // Draw background (dark gray) + for y in bar_y..height { + for x in 0..width { + debug_img.put_pixel(x, y, Rgb([40, 40, 40])); + } + } + + // Draw bar outline (white) + let outline_y = bar_y + 4; + let outline_height = bar_height - 8; + for x in bar_x..bar_x + bar_width { + debug_img.put_pixel(x, outline_y, Rgb([200, 200, 200])); + debug_img.put_pixel(x, outline_y + outline_height - 1, Rgb([200, 200, 200])); + } + for y in outline_y..outline_y + outline_height { + debug_img.put_pixel(bar_x, y, Rgb([200, 200, 200])); + debug_img.put_pixel(bar_x + bar_width - 1, y, Rgb([200, 200, 200])); + } + + // Fill the bar based on brightness value + let fill_width = ((bar_width - 4) as f32 * brightness.clamp(0.0, 1.0)) as u32; + let fill_color = brightness_to_color(brightness); + for y in (outline_y + 2)..(outline_y + outline_height - 2) { + for x in (bar_x + 2)..(bar_x + 2 + fill_width) { + if x < width { + debug_img.put_pixel(x, y, fill_color); + } + } + } + + // Draw brightness text value + let text = format!("B:{:.2}", brightness); + draw_simple_text(&mut debug_img, 5, bar_y + 6, &text, Rgb([255, 255, 255])); + + Some(DynamicImage::ImageRgb8(debug_img)) + } +} + +/// Convert brightness value to a color (red for dark/bright, green for good) +fn brightness_to_color(brightness: f32) -> Rgb { + // Very dark or very bright = red, middle range = green + if brightness < 0.15 || brightness > 0.85 { + Rgb([255, 80, 80]) // Red + } else if brightness < 0.25 || brightness > 0.75 { + Rgb([255, 200, 80]) // Yellow/orange + } else { + Rgb([80, 255, 80]) // Green + } +} + +/// Draw simple text using a basic 5x7 pixel font. +fn draw_simple_text(img: &mut RgbImage, x: u32, y: u32, text: &str, color: Rgb) { + let (width, height) = (img.width(), img.height()); + let mut cursor_x = x; + + for ch in text.chars() { + let pattern = get_char_pattern(ch); + for (row_idx, row) in pattern.iter().enumerate() { + for col in 0..5 { + if (row >> (4 - col)) & 1 == 1 { + let px = cursor_x + col; + let py = y + row_idx as u32; + if px < width && py < height { + img.put_pixel(px, py, color); + } + } + } + } + cursor_x += 6; // 5 pixels wide + 1 pixel spacing + } +} + +/// Get a 5x7 pixel pattern for a character. +fn get_char_pattern(ch: char) -> [u8; 7] { + match ch { + '0' => [0b01110, 0b10001, 0b10011, 0b10101, 0b11001, 0b10001, 0b01110], + '1' => [0b00100, 0b01100, 0b00100, 0b00100, 0b00100, 0b00100, 0b01110], + '2' => [0b01110, 0b10001, 0b00001, 0b00110, 0b01000, 0b10000, 0b11111], + '3' => [0b01110, 0b10001, 0b00001, 0b00110, 0b00001, 0b10001, 0b01110], + '4' => [0b00010, 0b00110, 0b01010, 0b10010, 0b11111, 0b00010, 0b00010], + '5' => [0b11111, 0b10000, 0b11110, 0b00001, 0b00001, 0b10001, 0b01110], + '6' => [0b00110, 0b01000, 0b10000, 0b11110, 0b10001, 0b10001, 0b01110], + '7' => [0b11111, 0b00001, 0b00010, 0b00100, 0b01000, 0b01000, 0b01000], + '8' => [0b01110, 0b10001, 0b10001, 0b01110, 0b10001, 0b10001, 0b01110], + '9' => [0b01110, 0b10001, 0b10001, 0b01111, 0b00001, 0b00010, 0b01100], + 'B' => [0b11110, 0b10001, 0b10001, 0b11110, 0b10001, 0b10001, 0b11110], + ':' => [0b00000, 0b00100, 0b00000, 0b00000, 0b00100, 0b00000, 0b00000], + '.' => [0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00100], + ' ' => [0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000], + _ => [0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000], + } } #[cfg(test)] diff --git a/src/pipeline/steps/crop.rs b/src/pipeline/steps/crop.rs index 8a7ba45..18dc431 100644 --- a/src/pipeline/steps/crop.rs +++ b/src/pipeline/steps/crop.rs @@ -5,7 +5,7 @@ use crate::config::Config; use crate::face_processing::crop_face_with_intermediate; use crate::face_processing::debug::draw_crop_debug; -use crate::pipeline::{PipelineContext, ProcessingStep, StepOutcome}; +use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome}; use async_trait::async_trait; use image::DynamicImage; @@ -33,15 +33,18 @@ impl ProcessingStep for CropFaceStep { } }; - // Crop returns (full_res_crop, resized). We want the full_res_crop here - // and let the resize step handle the final sizing. + // Crop returns CropResult with cropped images and face rectangle in crop coordinates. + // We use the full_res_crop here and let the resize step handle the final sizing. match crop_face_with_intermediate(image, &ctx.face_data, config.processing.output.size) { - Ok((cropped_full, _resized)) => { + Ok(crop_result) => { // Use the full-resolution cropped image; resize step will handle final size - ctx.image = Some(cropped_full); + ctx.image = Some(crop_result.cropped); + // Store the face rectangle in crop coordinates for later steps + ctx.set_computed("face_rect", ComputedValue::FaceRect(crop_result.face_rect)); StepOutcome::Continue(ctx) } Err(e) => StepOutcome::Skip { + ctx, reason: "crop_failed".to_string(), detail: Some(e.to_string()), }, diff --git a/src/pipeline/steps/decode.rs b/src/pipeline/steps/decode.rs index 9d46c41..6c73f5c 100644 --- a/src/pipeline/steps/decode.rs +++ b/src/pipeline/steps/decode.rs @@ -40,6 +40,7 @@ impl ProcessingStep for DecodeImageStep { StepOutcome::Continue(ctx) } Err(e) => StepOutcome::Skip { + ctx, reason: "decode_failed".to_string(), detail: Some(e.to_string()), }, diff --git a/src/pipeline/steps/face_resolution.rs b/src/pipeline/steps/face_resolution.rs index 78181cc..d86db31 100644 --- a/src/pipeline/steps/face_resolution.rs +++ b/src/pipeline/steps/face_resolution.rs @@ -44,6 +44,7 @@ impl ProcessingStep for FaceResolutionStep { if face_size < threshold { return StepOutcome::Skip { + ctx, reason: "face_too_small".to_string(), detail: Some(format!("{}px (threshold: {}px)", face_size, threshold)), }; @@ -77,7 +78,7 @@ mod tests { let config = Config::default(); // threshold is 80px match step.execute(ctx, &config).await { - StepOutcome::Skip { reason, detail } => { + StepOutcome::Skip { reason, detail, .. } => { assert_eq!(reason, "face_too_small"); assert!(detail.unwrap().contains("50px")); } diff --git a/src/pipeline/steps/head_pose.rs b/src/pipeline/steps/head_pose.rs index 1546c26..865d1d9 100644 --- a/src/pipeline/steps/head_pose.rs +++ b/src/pipeline/steps/head_pose.rs @@ -7,7 +7,7 @@ use crate::config::Config; use crate::models::DMHeadModel; use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome}; use async_trait::async_trait; -use image::{DynamicImage, Rgb}; +use image::{DynamicImage, GenericImageView, Rgb, RgbImage}; /// Estimates head pose and filters non-frontal faces. /// @@ -50,8 +50,39 @@ impl ProcessingStep for HeadPoseStep { } }; - // Run inference - let pose = match model.estimate(image) { + // 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("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(format!("Head pose estimation failed: {}", e)); @@ -73,6 +104,7 @@ impl ProcessingStep for HeadPoseStep { 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}°", @@ -83,6 +115,7 @@ impl ProcessingStep for HeadPoseStep { 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}°", @@ -93,6 +126,7 @@ impl ProcessingStep for HeadPoseStep { 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}°", @@ -118,45 +152,144 @@ impl ProcessingStep for HeadPoseStep { // 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 a center crosshair + let cx = width / 2; + let cy = height / 2; + let crosshair_size = 20u32; - // 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); + // 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) { + 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; + } + } +} + +/// Draw simple text using a basic 5x7 pixel font. +/// Only supports basic ASCII characters needed for pose display. +fn draw_simple_text(img: &mut RgbImage, x: u32, y: u32, text: &str, color: Rgb) { + let (width, height) = (img.width(), img.height()); + let mut cursor_x = x; + + for ch in text.chars() { + let pattern = get_char_pattern(ch); + for (row_idx, row) in pattern.iter().enumerate() { + for col in 0..5 { + if (row >> (4 - col)) & 1 == 1 { + let px = cursor_x + col; + let py = y + row_idx as u32; + if px < width && py < height { + img.put_pixel(px, py, color); + } } } } + cursor_x += 6; // 5 pixels wide + 1 pixel spacing + } +} - // 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)) +/// Get a 5x7 pixel pattern for a character. +fn get_char_pattern(ch: char) -> [u8; 7] { + match ch { + '0' => [0b01110, 0b10001, 0b10011, 0b10101, 0b11001, 0b10001, 0b01110], + '1' => [0b00100, 0b01100, 0b00100, 0b00100, 0b00100, 0b00100, 0b01110], + '2' => [0b01110, 0b10001, 0b00001, 0b00110, 0b01000, 0b10000, 0b11111], + '3' => [0b01110, 0b10001, 0b00001, 0b00110, 0b00001, 0b10001, 0b01110], + '4' => [0b00010, 0b00110, 0b01010, 0b10010, 0b11111, 0b00010, 0b00010], + '5' => [0b11111, 0b10000, 0b11110, 0b00001, 0b00001, 0b10001, 0b01110], + '6' => [0b00110, 0b01000, 0b10000, 0b11110, 0b10001, 0b10001, 0b01110], + '7' => [0b11111, 0b00001, 0b00010, 0b00100, 0b01000, 0b01000, 0b01000], + '8' => [0b01110, 0b10001, 0b10001, 0b01110, 0b10001, 0b10001, 0b01110], + '9' => [0b01110, 0b10001, 0b10001, 0b01111, 0b00001, 0b00010, 0b01100], + 'Y' => [0b10001, 0b10001, 0b01010, 0b00100, 0b00100, 0b00100, 0b00100], + 'P' => [0b11110, 0b10001, 0b10001, 0b11110, 0b10000, 0b10000, 0b10000], + 'R' => [0b11110, 0b10001, 0b10001, 0b11110, 0b10100, 0b10010, 0b10001], + ':' => [0b00000, 0b00100, 0b00000, 0b00000, 0b00100, 0b00000, 0b00000], + '+' => [0b00000, 0b00100, 0b00100, 0b11111, 0b00100, 0b00100, 0b00000], + '-' => [0b00000, 0b00000, 0b00000, 0b11111, 0b00000, 0b00000, 0b00000], + ' ' => [0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000], + '.' => [0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00100], + _ => [0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000, 0b00000], } } diff --git a/src/pipeline/steps/landmarks.rs b/src/pipeline/steps/landmarks.rs index 22ab4f5..984859f 100644 --- a/src/pipeline/steps/landmarks.rs +++ b/src/pipeline/steps/landmarks.rs @@ -54,6 +54,7 @@ impl ProcessingStep for LandmarksStep { // If model isn't available, skip this step with a warning tracing::warn!("Dlib landmarks model not available: {}", e); return StepOutcome::Skip { + ctx, reason: "landmarks_failed".to_string(), detail: Some(e.to_string()), }; @@ -65,9 +66,15 @@ impl ProcessingStep for LandmarksStep { let (width, height) = (rgb.width() as usize, rgb.height() as usize); let pixels = rgb.into_raw(); + // Get the face rectangle if available + let face_rect: Option<(i64, i64, i64, i64)> = ctx + .get_computed("face_rect") + .and_then(|v| v.as_face_rect()) + .map(|r| (r.x1 as i64, r.y1 as i64, r.x2 as i64, r.y2 as i64)); + // Run dlib operations in a blocking thread to avoid dropping in async context let landmarks_result = task::spawn_blocking(move || -> Result { - dlib.detect_landmarks(width, height, &pixels) + dlib.detect_landmarks(width, height, &pixels, face_rect) .map_err(|e| e.to_string()) }) .await; @@ -76,6 +83,7 @@ impl ProcessingStep for LandmarksStep { Ok(Ok(l)) => l, Ok(Err(e)) => { return StepOutcome::Skip { + ctx, reason: "landmarks_failed".to_string(), detail: Some(e), }; @@ -98,6 +106,7 @@ impl ProcessingStep for LandmarksStep { let min_ear = config.processing.eye_filter.min_ear; if avg_ear < min_ear { return StepOutcome::Skip { + ctx, reason: "eyes_closed".to_string(), detail: Some(format!("EAR {:.3} below threshold {:.3}", avg_ear, min_ear)), }; @@ -171,15 +180,20 @@ fn draw_cross(img: &mut RgbImage, x: u32, y: u32, color: Rgb) { let (width, height) = (img.width(), img.height()); let size = 2; + // Check base coordinates are in bounds + if x >= width || y >= height { + return; + } + for dx in 0..=size * 2 { let px = (x as i32 + dx as i32 - size as i32) as u32; - if px < width { + if px < width && y < height { 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 { + if x < width && py < height { img.put_pixel(x, py, color); } } diff --git a/src/pipeline/traits.rs b/src/pipeline/traits.rs index a47e498..96f4846 100644 --- a/src/pipeline/traits.rs +++ b/src/pipeline/traits.rs @@ -4,7 +4,7 @@ //! extensible image processing pipelines. use crate::config::Config; -use crate::face_processing::types::{HeadPose, Landmarks}; +use crate::face_processing::types::{BoundingBox, HeadPose, Landmarks}; use crate::immich_api::FaceData; use async_trait::async_trait; use bytes::Bytes; @@ -17,8 +17,12 @@ use std::fmt; pub enum StepOutcome { /// Continue to the next step with the updated context. Continue(PipelineContext), - /// Skip this image with the given reason. - Skip { reason: String, detail: Option }, + /// Skip this image with the given reason. Context is returned for debug visualization. + Skip { + ctx: PipelineContext, + reason: String, + detail: Option, + }, /// An error occurred during processing. Error(String), } @@ -38,6 +42,8 @@ pub enum ComputedValue { HeadPose(HeadPose), /// Facial landmarks (68 points). Landmarks(Box), + /// Face bounding box in current image coordinates. + FaceRect(BoundingBox), } impl ComputedValue { @@ -88,6 +94,30 @@ impl ComputedValue { _ => None, } } + + /// Get as BoundingBox if this is a FaceRect variant. + pub fn as_face_rect(&self) -> Option<&BoundingBox> { + match self { + ComputedValue::FaceRect(v) => Some(v), + _ => None, + } + } +} + +/// A debug image with metadata about whether the step passed or failed. +#[derive(Debug, Clone)] +pub struct DebugImage { + /// The debug visualization image. + pub image: DynamicImage, + /// Whether the step that generated this image passed (true) or failed/skipped (false). + pub passed: bool, +} + +impl DebugImage { + /// Create a new debug image. + pub fn new(image: DynamicImage, passed: bool) -> Self { + Self { image, passed } + } } /// Context passed through the pipeline, carrying data between steps. @@ -105,8 +135,8 @@ pub struct PipelineContext { pub face_data: FaceData, /// Values computed by previous steps (e.g., brightness, landmarks). pub computed: HashMap, - /// Debug images generated by steps (step_id -> debug image). - pub debug_images: HashMap, + /// Debug images generated by steps (step_id -> debug image with pass/fail status). + pub debug_images: HashMap, } impl PipelineContext { @@ -146,8 +176,13 @@ impl PipelineContext { } /// Add a debug image for a step. - pub fn add_debug_image(&mut self, step_id: impl Into, image: DynamicImage) { - self.debug_images.insert(step_id.into(), image); + /// + /// # Arguments + /// * `step_id` - The identifier of the step that generated this image + /// * `image` - The debug visualization image + /// * `passed` - Whether the step passed (true) or failed/skipped (false) + pub fn add_debug_image(&mut self, step_id: impl Into, image: DynamicImage, passed: bool) { + self.debug_images.insert(step_id.into(), DebugImage::new(image, passed)); } } diff --git a/src/video/ffmpeg.rs b/src/video/ffmpeg.rs index 7fd12d0..0f1ac7d 100644 --- a/src/video/ffmpeg.rs +++ b/src/video/ffmpeg.rs @@ -128,7 +128,9 @@ where if line.starts_with("frame=") { if let Some(frame_str) = line.strip_prefix("frame=") { if let Ok(frame) = frame_str.trim().parse::() { - progress_callback(frame, image_count); + // Clamp to image_count - ffmpeg can report higher values internally + let clamped_frame = frame.min(image_count); + progress_callback(clamped_frame, image_count); } } }