From 41d50167ebf150160d07109555cd57e058a3d1a9 Mon Sep 17 00:00:00 2001 From: Arnaud_Cayrol Date: Sat, 31 Jan 2026 14:51:40 +0100 Subject: [PATCH] Changed the face processing to generic tasks with traits. No need to hardcode them anymore. --- Cargo.toml | 2 + .../src/lib/components/SettingsPanel.svelte | 105 ++++++++ src/config.rs | 30 +++ src/job/mod.rs | 233 ++++++++--------- src/lib.rs | 1 + src/main.rs | 8 + src/pipeline/mod.rs | 238 ++++++++++++++++++ src/pipeline/steps/brightness.rs | 227 +++++++++++++++++ src/pipeline/steps/crop.rs | 144 +++++++++++ src/pipeline/steps/decode.rs | 133 ++++++++++ src/pipeline/steps/face_resolution.rs | 98 ++++++++ src/pipeline/steps/mod.rs | 16 ++ src/pipeline/steps/resize.rs | 127 ++++++++++ src/pipeline/traits.rs | 224 +++++++++++++++++ src/web/handlers/processing.rs | 40 ++- src/web/state.rs | 127 +++++----- 16 files changed, 1544 insertions(+), 209 deletions(-) create mode 100644 src/pipeline/mod.rs create mode 100644 src/pipeline/steps/brightness.rs create mode 100644 src/pipeline/steps/crop.rs create mode 100644 src/pipeline/steps/decode.rs create mode 100644 src/pipeline/steps/face_resolution.rs create mode 100644 src/pipeline/steps/mod.rs create mode 100644 src/pipeline/steps/resize.rs create mode 100644 src/pipeline/traits.rs diff --git a/Cargo.toml b/Cargo.toml index 5f401d6..92ab29d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -42,6 +42,8 @@ chrono = { version = "0.4", features = ["serde"] } directories = "5" bytes = "1" tokio-util = "0.7" +async-trait = "0.1" +dotenvy = "0.15" [dev-dependencies] tokio-test = "0.4" diff --git a/frontend/src/lib/components/SettingsPanel.svelte b/frontend/src/lib/components/SettingsPanel.svelte index c13ffea..37e9370 100644 --- a/frontend/src/lib/components/SettingsPanel.svelte +++ b/frontend/src/lib/components/SettingsPanel.svelte @@ -19,6 +19,7 @@ ear_threshold: 0.2, max_workers: 4, keep_intermediates: false, + brightness: null, }, video: { framerate: 15, @@ -28,6 +29,17 @@ }, }); + // Helper to ensure brightness config exists + function ensureBrightnessConfig() { + if (!config.processing.brightness) { + config.processing.brightness = { + enabled: false, + min_brightness: 0.1, + max_brightness: 0.95, + }; + } + } + async function loadConfig() { loading = true; error = null; @@ -75,6 +87,7 @@ ear_threshold: 0.2, max_workers: 4, keep_intermediates: false, + brightness: null, }, video: { framerate: 15, @@ -196,6 +209,65 @@ {config.processing.pose_threshold}° + + +
+
+ Brightness Filter + { + ensureBrightnessConfig(); + config.processing.brightness.enabled = e.target.checked; + }} + /> +
+ + {#if config.processing.brightness?.enabled} +
+ +
+ { + ensureBrightnessConfig(); + config.processing.brightness.min_brightness = parseFloat(e.target.value); + }} + min="0" + max="0.5" + step="0.05" + /> + {(config.processing.brightness?.min_brightness ?? 0.1).toFixed(2)} +
+
+ +
+ +
+ { + ensureBrightnessConfig(); + config.processing.brightness.max_brightness = parseFloat(e.target.value); + }} + min="0.5" + max="1.0" + step="0.05" + /> + {(config.processing.brightness?.max_brightness ?? 0.95).toFixed(2)} +
+
+ {/if} +
{:else if activeTab === 'output'}
@@ -491,6 +563,39 @@ cursor: pointer; } + /* Settings section with header */ + .setting-section { + margin-top: 1rem; + padding-top: 1rem; + border-top: 1px solid #333; + } + + .section-header { + display: flex; + justify-content: space-between; + align-items: center; + margin-bottom: 0.5rem; + } + + .section-title { + font-size: 0.875rem; + font-weight: 600; + color: #e0e0e0; + } + + .section-header input[type='checkbox'] { + width: 1.25rem; + height: 1.25rem; + accent-color: #4f46e5; + cursor: pointer; + } + + .sub-setting { + padding-left: 1rem; + border-left: 2px solid #333; + margin-left: 0.5rem; + } + .actions { margin-top: 1.5rem; display: flex; diff --git a/src/config.rs b/src/config.rs index 034716f..0f20cc5 100644 --- a/src/config.rs +++ b/src/config.rs @@ -94,6 +94,35 @@ pub struct ProcessingConfig { /// Keep intermediate images (original, cropped) for inspection. pub keep_intermediates: bool, + + /// Brightness validation settings (optional). + #[serde(skip_serializing_if = "Option::is_none")] + pub brightness: Option, +} + +/// Brightness validation configuration. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BrightnessConfig { + /// Whether brightness validation is enabled. + pub enabled: bool, + + /// Minimum acceptable brightness (0.0-1.0). + /// Images darker than this will be skipped. + pub min_brightness: f32, + + /// Maximum acceptable brightness (0.0-1.0). + /// Images brighter than this will be skipped. + pub max_brightness: f32, +} + +impl Default for BrightnessConfig { + fn default() -> Self { + Self { + enabled: false, + min_brightness: 0.1, + max_brightness: 0.95, + } + } } impl Default for ProcessingConfig { @@ -107,6 +136,7 @@ impl Default for ProcessingConfig { right_eye_pos: (0.65, 0.4), max_workers: num_cpus(), keep_intermediates: false, + brightness: None, } } } diff --git a/src/job/mod.rs b/src/job/mod.rs index 46bb1f6..8d60cae 100644 --- a/src/job/mod.rs +++ b/src/job/mod.rs @@ -3,20 +3,18 @@ //! Handles the complete pipeline: //! 1. Fetching assets from Immich for a specific person //! 2. Downloading and processing images in parallel -//! 3. Cropping faces using bounding boxes +//! 3. Processing images through the extensible pipeline //! 4. Compiling processed images into a timelapse video use crate::config::Config; use crate::error::{Error, Result}; -use crate::face_processing::crop_face_with_intermediate; -use crate::face_processing::debug::draw_crop_debug; -use crate::face_processing::load_image_with_orientation; -use crate::face_processing::{AssetResult, ProcessedFace, SkipReason}; use crate::immich_api::{Asset, FaceData, ImmichClient}; +use crate::pipeline::{Pipeline, PipelineContext, PipelineResult}; use crate::utils::sanitize_folder_name; use crate::video::compile_timelapse; use crate::web::{AppState, AtomicSkipStats, JobStatus, Progress, SkipStats}; +use bytes::Bytes; use image::ImageFormat; use std::io::Cursor; use std::path::PathBuf; @@ -234,6 +232,10 @@ async fn run_job_inner( }) .await; + // Create the processing pipeline + let pipeline = Arc::new(Pipeline::with_default_steps()); + tracing::debug!("Pipeline steps: {:?}", pipeline.step_ids()); + // Process images in parallel with concurrency limit let completed = Arc::new(AtomicU32::new(0)); let semaphore = Arc::new(Semaphore::new(config.processing.max_workers)); @@ -266,6 +268,7 @@ async fn run_job_inner( let completed = completed.clone(); let state = state.clone(); let task_cancel_token = cancel_token.clone(); + let pipeline = pipeline.clone(); // Clone skip stats for this task let skip_stats = skip_stats.clone(); @@ -278,10 +281,8 @@ async fn run_job_inner( // Check cancellation at start of task if task_cancel_token.is_cancelled() { drop(permit); - return AssetResult::Skipped { + return AssetProcessResult::Cancelled { asset_id: asset.id.clone(), - reason: SkipReason::Cancelled, - detail: None, }; } @@ -292,14 +293,11 @@ async fn run_job_inner( &face_data, &output_dirs, &task_cancel_token, + &skip_stats, + &pipeline, ) .await; - // Update skip counters based on result - if let AssetResult::Skipped { reason, .. } = &result { - skip_stats.increment(reason); - } - // Update progress with current skip stats (only if not cancelled) if !task_cancel_token.is_cancelled() { let done = completed.fetch_add(1, Ordering::SeqCst) + 1; @@ -346,11 +344,11 @@ async fn run_job_inner( // Count results let successful = results .iter() - .filter(|r| matches!(r, AssetResult::Success(_))) + .filter(|r| matches!(r, AssetProcessResult::Success { .. })) .count(); let errors = results .iter() - .filter(|r| matches!(r, AssetResult::Error { .. })) + .filter(|r| matches!(r, AssetProcessResult::Error { .. })) .count(); let skipped = final_skip_stats.total(); @@ -434,7 +432,23 @@ fn find_face_for_person(asset: &Asset, person_id: &str) -> Option { None } -/// Process a single asset: download, crop face, save. +/// Result of processing a single asset. +/// +/// Fields are kept for debugging/logging purposes even if not currently read. +#[derive(Debug)] +#[allow(dead_code)] +enum AssetProcessResult { + /// Successfully processed. + Success { asset_id: String }, + /// Skipped for some reason. + Skipped { asset_id: String, reason: String }, + /// Error during processing. + Error { asset_id: String, error: String }, + /// Cancelled by user. + Cancelled { asset_id: String }, +} + +/// Process a single asset using the pipeline. async fn process_single_asset( client: &ImmichClient, config: &Config, @@ -442,45 +456,11 @@ async fn process_single_asset( face_data: &FaceData, output_dirs: &OutputDirs, cancel_token: &CancellationToken, -) -> AssetResult { + skip_stats: &Arc, + pipeline: &Pipeline, +) -> AssetProcessResult { let asset_id = &asset.id; - // Check face resolution (bounding box coordinates are already in pixels) - let face_width = face_data.bounding_box_x2 - face_data.bounding_box_x1; - let face_height = face_data.bounding_box_y2 - face_data.bounding_box_y1; - let face_size = face_width.min(face_height) as u32; - - if face_size < config.processing.face_resolution_threshold { - return AssetResult::Skipped { - asset_id: asset_id.clone(), - reason: SkipReason::FaceTooSmall, - detail: Some(format!( - "{}px (threshold: {}px)", - face_size, config.processing.face_resolution_threshold - )), - }; - } - - // Check before download (potentially slow) - if cancel_token.is_cancelled() { - return AssetResult::Skipped { - asset_id: asset_id.clone(), - reason: SkipReason::Cancelled, - detail: None, - }; - } - - // Download image - let image_bytes = match client.download_asset(asset_id).await { - Ok(bytes) => bytes, - Err(e) => { - return AssetResult::Error { - asset_id: asset_id.clone(), - error: format!("Download failed: {}", e), - }; - } - }; - // Generate timestamp-based filename for sorting let timestamp = asset .file_created_at @@ -489,99 +469,96 @@ async fn process_single_asset( .cloned() .unwrap_or_else(|| asset_id.clone()); - // Sanitize timestamp for filename (replace colons and other invalid chars) - let safe_timestamp: String = timestamp - .chars() - .map(|c| { - if c.is_alphanumeric() || c == '-' || c == '_' { - c - } else { - '_' - } - }) - .collect(); + // Create pipeline context + let mut ctx = PipelineContext::new( + asset_id.clone(), + timestamp.clone(), + face_data.clone(), + ); - let filename = format!("{}_{}.jpg", safe_timestamp, asset_id); - - // Check after download, before CPU-intensive processing + // Check before download (potentially slow) if cancel_token.is_cancelled() { - return AssetResult::Skipped { + return AssetProcessResult::Cancelled { asset_id: asset_id.clone(), - reason: SkipReason::Cancelled, - detail: None, }; } - // Decode image and apply EXIF orientation correction. - // This ensures bounding box coordinates from Immich align with the image pixels. - let img = match load_image_with_orientation(&image_bytes) { - Ok(img) => img, + // Download image + let image_bytes: Bytes = match client.download_asset(asset_id).await { + Ok(bytes) => bytes, Err(e) => { - return AssetResult::Error { + skip_stats.increment("download_failed"); + return AssetProcessResult::Error { asset_id: asset_id.clone(), - error: format!("Failed to decode image: {}", e), + error: format!("Download failed: {}", e), }; } }; - // Crop face - let (_cropped_full, final_image) = - match crop_face_with_intermediate(&img, face_data, config.processing.resize_size) { - Ok(result) => result, - Err(e) => { - return AssetResult::Error { - asset_id: asset_id.clone(), - error: format!("Failed to crop face: {}", e), + // Set raw bytes on context + ctx = ctx.with_bytes(image_bytes); + + // Determine debug directory + let debug_dir = if config.processing.keep_intermediates { + output_dirs.debug.as_ref().and_then(|d| d.crop.as_ref().map(|p| p.parent().unwrap().to_path_buf())) + } else { + None + }; + + // Execute the pipeline + let result = pipeline.execute( + ctx, + config, + cancel_token, + skip_stats, + debug_dir.as_ref(), + ).await; + + match result { + PipelineResult::Success { image, asset_id, timestamp, .. } => { + // Sanitize timestamp for filename + let safe_timestamp: String = timestamp + .chars() + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '_' { + c + } else { + '_' + } + }) + .collect(); + + let filename = format!("{}_{}.jpg", safe_timestamp, asset_id); + let output_path = output_dirs.images.join(&filename); + + // Encode and save final image + let mut buffer = Cursor::new(Vec::new()); + if let Err(e) = image.write_to(&mut buffer, ImageFormat::Jpeg) { + return AssetProcessResult::Error { + asset_id, + error: format!("Failed to encode image: {}", e), }; } - }; - // Save debug visualization if enabled - if let Some(ref debug) = output_dirs.debug { - if let Some(ref crop_dir) = debug.crop { - let debug_img = draw_crop_debug(&img, face_data); - let debug_path = crop_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); - } else if let Err(e) = tokio::fs::write(&debug_path, buffer.into_inner()).await { - tracing::warn!("Failed to save debug image: {}", e); + if let Err(e) = tokio::fs::write(&output_path, buffer.into_inner()).await { + return AssetProcessResult::Error { + asset_id, + error: format!("Failed to save image: {}", e), + }; } + + AssetProcessResult::Success { asset_id } + } + PipelineResult::Skipped { asset_id, reason, .. } => { + AssetProcessResult::Skipped { asset_id, reason } + } + PipelineResult::Error { asset_id, error } => { + AssetProcessResult::Error { asset_id, error } + } + PipelineResult::Cancelled { asset_id } => { + AssetProcessResult::Cancelled { asset_id } } } - - // Check before file I/O - if cancel_token.is_cancelled() { - return AssetResult::Skipped { - asset_id: asset_id.clone(), - reason: SkipReason::Cancelled, - detail: None, - }; - } - - let output_path = output_dirs.images.join(&filename); - - // Encode and save final image - let mut buffer = Cursor::new(Vec::new()); - if let Err(e) = final_image.write_to(&mut buffer, ImageFormat::Jpeg) { - return AssetResult::Error { - asset_id: asset_id.clone(), - error: format!("Failed to encode image: {}", e), - }; - } - - if let Err(e) = tokio::fs::write(&output_path, buffer.into_inner()).await { - return AssetResult::Error { - asset_id: asset_id.clone(), - error: format!("Failed to save image: {}", e), - }; - } - - AssetResult::Success(ProcessedFace { - image_data: Vec::new(), // We saved to file, so don't keep in memory - asset_id: asset_id.clone(), - timestamp, - }) } #[cfg(test)] diff --git a/src/lib.rs b/src/lib.rs index adeb73c..bf05f90 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,7 @@ pub mod error; pub mod face_processing; pub mod immich_api; pub mod job; +pub mod pipeline; pub mod utils; pub mod video; pub mod web; diff --git a/src/main.rs b/src/main.rs index d0317f6..623cccf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,14 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; #[tokio::main] async fn main() -> anyhow::Result<()> { + // Load .env file if present (before anything else) + if let Err(e) = dotenvy::dotenv() { + // Not an error if .env doesn't exist + if !matches!(e, dotenvy::Error::Io(_)) { + eprintln!("Warning: Failed to load .env file: {}", e); + } + } + // Initialize logging tracing_subscriber::registry() .with( diff --git a/src/pipeline/mod.rs b/src/pipeline/mod.rs new file mode 100644 index 0000000..fe3c79a --- /dev/null +++ b/src/pipeline/mod.rs @@ -0,0 +1,238 @@ +//! Extensible image processing pipeline. +//! +//! This module provides a trait-based pipeline architecture for processing images. +//! Steps can be validators (skip if condition fails), transformers (modify image), +//! or computers (calculate values for later steps). +//! +//! # Example +//! +//! ```ignore +//! use crate::pipeline::{Pipeline, PipelineContext}; +//! +//! let pipeline = Pipeline::default(); +//! let result = pipeline.execute(ctx, &config, &cancel_token, &skip_stats).await; +//! ``` + +mod traits; +pub mod steps; + +pub use traits::*; + +use crate::config::Config; +use crate::web::AtomicSkipStats; +use image::{DynamicImage, ImageFormat}; +use std::io::Cursor; +use std::path::PathBuf; +use std::sync::Arc; +use tokio_util::sync::CancellationToken; + +/// Result of pipeline execution. +#[derive(Debug)] +pub enum PipelineResult { + /// Successfully processed the image. + Success { + /// The final processed image. + image: DynamicImage, + /// Asset ID for reference. + asset_id: String, + /// Timestamp for file naming. + timestamp: String, + /// Debug images if keep_intermediates was enabled. + debug_images: Vec<(String, DynamicImage)>, + }, + /// Image was skipped. + Skipped { + asset_id: String, + reason: String, + detail: Option, + }, + /// Processing error. + Error { + asset_id: String, + error: String, + }, + /// Cancelled by user. + Cancelled { + asset_id: String, + }, +} + +/// Processing pipeline that executes steps sequentially. +pub struct Pipeline { + steps: Vec>, +} + +impl std::fmt::Debug for Pipeline { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Pipeline") + .field("steps", &self.steps.iter().map(|s| s.id()).collect::>()) + .finish() + } +} + +impl Default for Pipeline { + fn default() -> Self { + Self::new() + } +} + +impl Pipeline { + /// Create a new empty pipeline. + pub fn new() -> Self { + Self { steps: Vec::new() } + } + + /// Create the default processing pipeline with standard steps. + pub fn with_default_steps() -> Self { + use steps::*; + + let mut pipeline = Self::new(); + pipeline.add_step(Box::new(FaceResolutionStep)); + pipeline.add_step(Box::new(DecodeImageStep)); + pipeline.add_step(Box::new(BrightnessStep)); + pipeline.add_step(Box::new(CropFaceStep)); + pipeline.add_step(Box::new(ResizeStep)); + pipeline + } + + /// Add a step to the pipeline. + pub fn add_step(&mut self, step: Box) { + self.steps.push(step); + } + + /// Get the list of step IDs in this pipeline. + pub fn step_ids(&self) -> Vec<&'static str> { + self.steps.iter().map(|s| s.id()).collect() + } + + /// Execute the pipeline on the given context. + /// + /// Runs each step sequentially, checking for cancellation between steps. + /// Updates skip statistics for any skipped images. + pub async fn execute( + &self, + mut ctx: PipelineContext, + config: &Config, + cancel_token: &CancellationToken, + skip_stats: &Arc, + debug_dir: Option<&PathBuf>, + ) -> PipelineResult { + let asset_id = ctx.asset_id.clone(); + let timestamp = ctx.timestamp.clone(); + + for step in &self.steps { + // Check for cancellation before each step + if cancel_token.is_cancelled() { + return PipelineResult::Cancelled { asset_id }; + } + + tracing::trace!("Executing step '{}' for asset {}", step.id(), asset_id); + + match step.execute(ctx, config).await { + StepOutcome::Continue(new_ctx) => { + ctx = new_ctx; + + // Generate debug visualization if enabled + if config.processing.keep_intermediates { + if let Some(debug_img) = step.debug_visualize(&ctx) { + ctx.add_debug_image(step.id(), debug_img); + } + } + } + StepOutcome::Skip { reason, detail } => { + // Update skip stats with the step ID as reason + skip_stats.increment(&reason); + + tracing::debug!( + "Asset {} skipped at step '{}': {} ({})", + asset_id, + step.id(), + reason, + detail.as_deref().unwrap_or("no detail") + ); + + return PipelineResult::Skipped { + asset_id, + reason, + detail, + }; + } + StepOutcome::Error(error) => { + tracing::warn!( + "Asset {} error at step '{}': {}", + asset_id, + step.id(), + error + ); + + return PipelineResult::Error { asset_id, error }; + } + } + } + + // All steps completed successfully + let final_image = match ctx.image { + Some(img) => img, + None => { + return PipelineResult::Error { + asset_id, + error: "Pipeline completed but no image was produced".to_string(), + }; + } + }; + + // 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); + } + } + } + + PipelineResult::Success { + image: final_image, + asset_id, + timestamp, + debug_images, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pipeline_step_ids() { + let pipeline = Pipeline::with_default_steps(); + let ids = pipeline.step_ids(); + + assert!(ids.contains(&"face_resolution")); + assert!(ids.contains(&"decode")); + assert!(ids.contains(&"brightness")); + assert!(ids.contains(&"crop")); + assert!(ids.contains(&"resize")); + } +} diff --git a/src/pipeline/steps/brightness.rs b/src/pipeline/steps/brightness.rs new file mode 100644 index 0000000..8d5239c --- /dev/null +++ b/src/pipeline/steps/brightness.rs @@ -0,0 +1,227 @@ +//! Brightness validation step. +//! +//! Calculates average image brightness and skips images that are too dark or too bright. + +use crate::config::Config; +use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome}; +use async_trait::async_trait; +use image::DynamicImage; + +/// Validates image brightness and skips images outside acceptable range. +/// +/// This is a validation step that computes the average luminance of the image +/// and skips if it falls below `min_brightness` or above `max_brightness`. +/// +/// Brightness is normalized to 0.0-1.0 range. +pub struct BrightnessStep; + +impl BrightnessStep { + /// Calculate the average brightness of an image. + /// + /// Returns a value between 0.0 (pure black) and 1.0 (pure white). + fn calculate_brightness(image: &DynamicImage) -> f32 { + let rgb = image.to_rgb8(); + let pixels = rgb.pixels(); + let pixel_count = rgb.width() as u64 * rgb.height() as u64; + + if pixel_count == 0 { + return 0.0; + } + + // Calculate average luminance using standard RGB to luminance conversion + // Y = 0.299*R + 0.587*G + 0.114*B + let total_luminance: u64 = pixels + .map(|p| { + let [r, g, b] = p.0; + // Use integer math for speed, then convert + (299 * r as u64 + 587 * g as u64 + 114 * b as u64) / 1000 + }) + .sum(); + + (total_luminance as f32 / pixel_count as f32) / 255.0 + } +} + +#[async_trait] +impl ProcessingStep for BrightnessStep { + fn id(&self) -> &'static str { + "brightness" + } + + fn name(&self) -> &'static str { + "Brightness Check" + } + + async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { + let image = match &ctx.image { + Some(img) => img, + None => { + return StepOutcome::Error( + "No image available for brightness check".to_string(), + ); + } + }; + + let brightness = Self::calculate_brightness(image); + + // Store computed brightness for potential use by other steps + ctx.set_computed("brightness", ComputedValue::Float(brightness)); + + // Check brightness thresholds if configured + if let Some(ref brightness_config) = config.processing.brightness { + if !brightness_config.enabled { + return StepOutcome::Continue(ctx); + } + + if brightness < brightness_config.min_brightness { + return StepOutcome::Skip { + reason: "too_dark".to_string(), + detail: Some(format!( + "{:.2} (min: {:.2})", + brightness, brightness_config.min_brightness + )), + }; + } + + if brightness > brightness_config.max_brightness { + return StepOutcome::Skip { + reason: "too_bright".to_string(), + detail: Some(format!( + "{:.2} (max: {:.2})", + brightness, brightness_config.max_brightness + )), + }; + } + } + + StepOutcome::Continue(ctx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::BrightnessConfig; + use crate::immich_api::FaceData; + use image::{DynamicImage, RgbImage, Rgb}; + + fn make_ctx_with_image(image: DynamicImage) -> PipelineContext { + let face_data = FaceData { + bounding_box_x1: 0.0, + bounding_box_y1: 0.0, + bounding_box_x2: 10.0, + bounding_box_y2: 10.0, + image_width: 10, + image_height: 10, + }; + PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data) + .with_image(image) + } + + fn create_solid_image(r: u8, g: u8, b: u8) -> DynamicImage { + let img = RgbImage::from_fn(10, 10, |_, _| Rgb([r, g, b])); + DynamicImage::ImageRgb8(img) + } + + #[test] + fn test_calculate_brightness_black() { + let img = create_solid_image(0, 0, 0); + let brightness = BrightnessStep::calculate_brightness(&img); + assert!((brightness - 0.0).abs() < 0.01); + } + + #[test] + fn test_calculate_brightness_white() { + let img = create_solid_image(255, 255, 255); + let brightness = BrightnessStep::calculate_brightness(&img); + assert!((brightness - 1.0).abs() < 0.01); + } + + #[test] + fn test_calculate_brightness_gray() { + let img = create_solid_image(128, 128, 128); + let brightness = BrightnessStep::calculate_brightness(&img); + // Should be approximately 0.5 + assert!(brightness > 0.45 && brightness < 0.55); + } + + #[tokio::test] + async fn test_brightness_no_config() { + let step = BrightnessStep; + let img = create_solid_image(128, 128, 128); + let ctx = make_ctx_with_image(img); + let config = Config::default(); // No brightness config + + match step.execute(ctx, &config).await { + StepOutcome::Continue(new_ctx) => { + // Should still compute brightness even without thresholds + let brightness = new_ctx.get_computed("brightness") + .and_then(|v| v.as_float()) + .unwrap(); + assert!(brightness > 0.45 && brightness < 0.55); + } + _ => panic!("Expected Continue"), + } + } + + #[tokio::test] + async fn test_brightness_too_dark() { + let step = BrightnessStep; + let img = create_solid_image(10, 10, 10); // Very dark + let ctx = make_ctx_with_image(img); + + let mut config = Config::default(); + config.processing.brightness = Some(BrightnessConfig { + enabled: true, + min_brightness: 0.2, + max_brightness: 0.9, + }); + + match step.execute(ctx, &config).await { + StepOutcome::Skip { reason, .. } => { + assert_eq!(reason, "too_dark"); + } + _ => panic!("Expected Skip"), + } + } + + #[tokio::test] + async fn test_brightness_too_bright() { + let step = BrightnessStep; + let img = create_solid_image(250, 250, 250); // Very bright + let ctx = make_ctx_with_image(img); + + let mut config = Config::default(); + config.processing.brightness = Some(BrightnessConfig { + enabled: true, + min_brightness: 0.1, + max_brightness: 0.9, + }); + + match step.execute(ctx, &config).await { + StepOutcome::Skip { reason, .. } => { + assert_eq!(reason, "too_bright"); + } + _ => panic!("Expected Skip"), + } + } + + #[tokio::test] + async fn test_brightness_disabled() { + let step = BrightnessStep; + let img = create_solid_image(10, 10, 10); // Very dark + let ctx = make_ctx_with_image(img); + + let mut config = Config::default(); + config.processing.brightness = Some(BrightnessConfig { + enabled: false, // Disabled + min_brightness: 0.2, + max_brightness: 0.9, + }); + + match step.execute(ctx, &config).await { + StepOutcome::Continue(_) => {} // Should continue even though image is dark + _ => panic!("Expected Continue when disabled"), + } + } +} diff --git a/src/pipeline/steps/crop.rs b/src/pipeline/steps/crop.rs new file mode 100644 index 0000000..c41b7b6 --- /dev/null +++ b/src/pipeline/steps/crop.rs @@ -0,0 +1,144 @@ +//! Face cropping step. +//! +//! Crops the face region from the full image using the bounding box data. + +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 async_trait::async_trait; +use image::DynamicImage; + +/// Crops the face region from the full image. +/// +/// This transformer step extracts the face region using the bounding box +/// from Immich, adding padding around the face for context. +pub struct CropFaceStep; + +#[async_trait] +impl ProcessingStep for CropFaceStep { + fn id(&self) -> &'static str { + "crop" + } + + fn name(&self) -> &'static str { + "Face Crop" + } + + async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { + let image = match &ctx.image { + Some(img) => img, + None => { + return StepOutcome::Error("No image available for cropping".to_string()); + } + }; + + // Crop returns (full_res_crop, resized). We want 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.resize_size) { + Ok((cropped_full, _resized)) => { + // Use the full-resolution cropped image; resize step will handle final size + ctx.image = Some(cropped_full); + StepOutcome::Continue(ctx) + } + Err(e) => StepOutcome::Skip { + reason: "crop_failed".to_string(), + detail: Some(e.to_string()), + }, + } + } + + fn debug_visualize(&self, ctx: &PipelineContext) -> Option { + // Draw the crop region on the original image + // Note: This requires access to the original image before cropping, + // which we don't have here. For now, we'll create the debug image + // during execution if needed. This is a limitation of the current + // design that could be addressed by storing the original image. + + // For now, return None and handle debug visualization in the pipeline + // execution or via a separate mechanism + ctx.raw_bytes.as_ref()?; + + // If we had the original image, we could do: + // Some(draw_crop_debug(&original, &ctx.face_data)) + None + } +} + +/// Generates a debug visualization of the crop region. +/// +/// This can be called separately before the crop step to visualize +/// what will be cropped. +#[allow(dead_code)] +pub fn generate_crop_debug(image: &DynamicImage, ctx: &PipelineContext) -> DynamicImage { + draw_crop_debug(image, &ctx.face_data) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::immich_api::FaceData; + use image::{DynamicImage, RgbImage, Rgb}; + + fn make_ctx_with_image(image: DynamicImage) -> PipelineContext { + // Face in the center of a 100x100 image + let face_data = FaceData { + bounding_box_x1: 30.0, + bounding_box_y1: 30.0, + bounding_box_x2: 70.0, + bounding_box_y2: 70.0, + image_width: 100, + image_height: 100, + }; + PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data) + .with_image(image) + } + + fn create_test_image(width: u32, height: u32) -> DynamicImage { + let img = RgbImage::from_fn(width, height, |x, y| { + // Create a pattern so we can verify cropping + Rgb([(x % 256) as u8, (y % 256) as u8, 128]) + }); + DynamicImage::ImageRgb8(img) + } + + #[tokio::test] + async fn test_crop_success() { + let step = CropFaceStep; + let img = create_test_image(100, 100); + let ctx = make_ctx_with_image(img); + let config = Config::default(); + + match step.execute(ctx, &config).await { + StepOutcome::Continue(new_ctx) => { + assert!(new_ctx.image.is_some()); + // The cropped image should be square + let cropped = new_ctx.image.unwrap(); + assert_eq!(cropped.width(), cropped.height()); + } + _ => panic!("Expected Continue"), + } + } + + #[tokio::test] + async fn test_crop_no_image() { + let step = CropFaceStep; + let face_data = FaceData { + bounding_box_x1: 30.0, + bounding_box_y1: 30.0, + bounding_box_x2: 70.0, + bounding_box_y2: 70.0, + image_width: 100, + image_height: 100, + }; + let ctx = PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data); + let config = Config::default(); + + match step.execute(ctx, &config).await { + StepOutcome::Error(msg) => { + assert!(msg.contains("No image")); + } + _ => panic!("Expected Error"), + } + } +} diff --git a/src/pipeline/steps/decode.rs b/src/pipeline/steps/decode.rs new file mode 100644 index 0000000..9d46c41 --- /dev/null +++ b/src/pipeline/steps/decode.rs @@ -0,0 +1,133 @@ +//! Image decode step. +//! +//! Decodes raw image bytes and applies EXIF orientation correction. + +use crate::config::Config; +use crate::face_processing::load_image_with_orientation; +use crate::pipeline::{PipelineContext, ProcessingStep, StepOutcome}; +use async_trait::async_trait; + +/// Decodes image bytes and corrects EXIF orientation. +/// +/// This step expects `ctx.raw_bytes` to be set (from a previous download step +/// or passed in directly). It sets `ctx.image` with the decoded image. +pub struct DecodeImageStep; + +#[async_trait] +impl ProcessingStep for DecodeImageStep { + fn id(&self) -> &'static str { + "decode" + } + + fn name(&self) -> &'static str { + "Image Decode" + } + + async fn execute(&self, mut ctx: PipelineContext, _config: &Config) -> StepOutcome { + let raw_bytes = match &ctx.raw_bytes { + Some(bytes) => bytes, + None => { + return StepOutcome::Error( + "No raw bytes available for decoding".to_string(), + ); + } + }; + + // Decode image with EXIF orientation correction + match load_image_with_orientation(raw_bytes) { + Ok(image) => { + ctx.image = Some(image); + StepOutcome::Continue(ctx) + } + Err(e) => StepOutcome::Skip { + reason: "decode_failed".to_string(), + detail: Some(e.to_string()), + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bytes::Bytes; + use crate::immich_api::FaceData; + use image::{DynamicImage, RgbImage}; + use std::io::Cursor; + + fn make_test_jpeg() -> Bytes { + // Create a small test image + let img = RgbImage::from_fn(10, 10, |_, _| image::Rgb([128, 128, 128])); + let dynamic = DynamicImage::ImageRgb8(img); + + let mut buffer = Cursor::new(Vec::new()); + dynamic.write_to(&mut buffer, image::ImageFormat::Jpeg).unwrap(); + + Bytes::from(buffer.into_inner()) + } + + fn make_ctx_with_bytes(bytes: Bytes) -> PipelineContext { + let face_data = FaceData { + bounding_box_x1: 0.0, + bounding_box_y1: 0.0, + bounding_box_x2: 10.0, + bounding_box_y2: 10.0, + image_width: 10, + image_height: 10, + }; + PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data) + .with_bytes(bytes) + } + + #[tokio::test] + async fn test_decode_valid_jpeg() { + let step = DecodeImageStep; + let bytes = make_test_jpeg(); + let ctx = make_ctx_with_bytes(bytes); + let config = Config::default(); + + match step.execute(ctx, &config).await { + StepOutcome::Continue(new_ctx) => { + assert!(new_ctx.image.is_some()); + } + _ => panic!("Expected Continue"), + } + } + + #[tokio::test] + async fn test_decode_invalid_bytes() { + let step = DecodeImageStep; + let bytes = Bytes::from_static(b"not an image"); + let ctx = make_ctx_with_bytes(bytes); + let config = Config::default(); + + match step.execute(ctx, &config).await { + StepOutcome::Skip { reason, .. } => { + assert_eq!(reason, "decode_failed"); + } + _ => panic!("Expected Skip"), + } + } + + #[tokio::test] + async fn test_decode_no_bytes() { + let step = DecodeImageStep; + let face_data = FaceData { + bounding_box_x1: 0.0, + bounding_box_y1: 0.0, + bounding_box_x2: 10.0, + bounding_box_y2: 10.0, + image_width: 10, + image_height: 10, + }; + let ctx = PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data); + let config = Config::default(); + + match step.execute(ctx, &config).await { + StepOutcome::Error(msg) => { + assert!(msg.contains("No raw bytes")); + } + _ => panic!("Expected Error"), + } + } +} diff --git a/src/pipeline/steps/face_resolution.rs b/src/pipeline/steps/face_resolution.rs new file mode 100644 index 0000000..4e54f46 --- /dev/null +++ b/src/pipeline/steps/face_resolution.rs @@ -0,0 +1,98 @@ +//! Face resolution validation step. +//! +//! Skips images where the detected face is too small. + +use crate::config::Config; +use crate::pipeline::{ComputedValue, PipelineContext, ProcessingStep, StepOutcome}; +use async_trait::async_trait; + +/// Validates that the face bounding box meets minimum size requirements. +/// +/// This is a validation step that runs early in the pipeline (before downloading) +/// to quickly skip low-resolution faces. +pub struct FaceResolutionStep; + +#[async_trait] +impl ProcessingStep for FaceResolutionStep { + fn id(&self) -> &'static str { + "face_resolution" + } + + fn name(&self) -> &'static str { + "Face Resolution Check" + } + + async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { + let face_data = &ctx.face_data; + + // Calculate face size in pixels from bounding box + let face_width = face_data.bounding_box_x2 - face_data.bounding_box_x1; + let face_height = face_data.bounding_box_y2 - face_data.bounding_box_y1; + let face_size = face_width.min(face_height) as i32; + + // Store computed value for later steps or debugging + ctx.set_computed("face_size", ComputedValue::Int(face_size)); + + let threshold = config.processing.face_resolution_threshold as i32; + + if face_size < threshold { + return StepOutcome::Skip { + reason: "face_too_small".to_string(), + detail: Some(format!("{}px (threshold: {}px)", face_size, threshold)), + }; + } + + StepOutcome::Continue(ctx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::immich_api::FaceData; + + fn make_ctx(face_size: f32) -> PipelineContext { + let face_data = FaceData { + bounding_box_x1: 0.0, + bounding_box_y1: 0.0, + bounding_box_x2: face_size, + bounding_box_y2: face_size, + image_width: 1920, + image_height: 1080, + }; + PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data) + } + + #[tokio::test] + async fn test_face_too_small() { + let step = FaceResolutionStep; + let ctx = make_ctx(50.0); // 50px face + let config = Config::default(); // threshold is 80px + + match step.execute(ctx, &config).await { + StepOutcome::Skip { reason, detail } => { + assert_eq!(reason, "face_too_small"); + assert!(detail.unwrap().contains("50px")); + } + _ => panic!("Expected Skip"), + } + } + + #[tokio::test] + async fn test_face_large_enough() { + let step = FaceResolutionStep; + let ctx = make_ctx(100.0); // 100px face + let config = Config::default(); // threshold is 80px + + match step.execute(ctx, &config).await { + StepOutcome::Continue(new_ctx) => { + // Should have stored face_size + assert_eq!( + new_ctx.get_computed("face_size").and_then(|v| v.as_int()), + Some(100) + ); + } + _ => panic!("Expected Continue"), + } + } +} diff --git a/src/pipeline/steps/mod.rs b/src/pipeline/steps/mod.rs new file mode 100644 index 0000000..6ada8d2 --- /dev/null +++ b/src/pipeline/steps/mod.rs @@ -0,0 +1,16 @@ +//! Pipeline step implementations. +//! +//! Each step implements the `ProcessingStep` trait and performs a specific +//! operation in the image processing pipeline. + +mod face_resolution; +mod decode; +mod brightness; +mod crop; +mod resize; + +pub use face_resolution::FaceResolutionStep; +pub use decode::DecodeImageStep; +pub use brightness::BrightnessStep; +pub use crop::CropFaceStep; +pub use resize::ResizeStep; diff --git a/src/pipeline/steps/resize.rs b/src/pipeline/steps/resize.rs new file mode 100644 index 0000000..031287c --- /dev/null +++ b/src/pipeline/steps/resize.rs @@ -0,0 +1,127 @@ +//! Image resize step. +//! +//! Resizes the cropped face image to the configured output size. + +use crate::config::Config; +use crate::pipeline::{PipelineContext, ProcessingStep, StepOutcome}; +use async_trait::async_trait; +use image::imageops::FilterType; + +/// Resizes the image to the configured output size. +/// +/// This transformer step is typically the last step in the pipeline, +/// producing the final output image at the configured resolution. +pub struct ResizeStep; + +#[async_trait] +impl ProcessingStep for ResizeStep { + fn id(&self) -> &'static str { + "resize" + } + + fn name(&self) -> &'static str { + "Resize" + } + + async fn execute(&self, mut ctx: PipelineContext, config: &Config) -> StepOutcome { + let image = match ctx.image.take() { + Some(img) => img, + None => { + return StepOutcome::Error("No image available for resizing".to_string()); + } + }; + + let output_size = config.processing.resize_size; + + // Resize to square output + let resized = image.resize_exact(output_size, output_size, FilterType::Lanczos3); + + ctx.image = Some(resized); + StepOutcome::Continue(ctx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::immich_api::FaceData; + use image::{DynamicImage, RgbImage, Rgb}; + + fn make_ctx_with_image(image: DynamicImage) -> 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) + .with_image(image) + } + + fn create_test_image(width: u32, height: u32) -> DynamicImage { + let img = RgbImage::from_fn(width, height, |_, _| Rgb([128, 128, 128])); + DynamicImage::ImageRgb8(img) + } + + #[tokio::test] + async fn test_resize_success() { + let step = ResizeStep; + let img = create_test_image(200, 200); + let ctx = make_ctx_with_image(img); + + let mut config = Config::default(); + config.processing.resize_size = 512; + + match step.execute(ctx, &config).await { + StepOutcome::Continue(new_ctx) => { + let resized = new_ctx.image.unwrap(); + assert_eq!(resized.width(), 512); + assert_eq!(resized.height(), 512); + } + _ => panic!("Expected Continue"), + } + } + + #[tokio::test] + async fn test_resize_no_image() { + let step = ResizeStep; + 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, + }; + let ctx = PipelineContext::new("test".to_string(), "2024-01-01".to_string(), face_data); + let config = Config::default(); + + match step.execute(ctx, &config).await { + StepOutcome::Error(msg) => { + assert!(msg.contains("No image")); + } + _ => panic!("Expected Error"), + } + } + + #[tokio::test] + async fn test_resize_upscale() { + let step = ResizeStep; + let img = create_test_image(50, 50); // Small image + let ctx = make_ctx_with_image(img); + + let mut config = Config::default(); + config.processing.resize_size = 256; + + match step.execute(ctx, &config).await { + StepOutcome::Continue(new_ctx) => { + let resized = new_ctx.image.unwrap(); + assert_eq!(resized.width(), 256); + assert_eq!(resized.height(), 256); + } + _ => panic!("Expected Continue"), + } + } +} diff --git a/src/pipeline/traits.rs b/src/pipeline/traits.rs new file mode 100644 index 0000000..60b1708 --- /dev/null +++ b/src/pipeline/traits.rs @@ -0,0 +1,224 @@ +//! Pipeline traits and core types. +//! +//! Defines the `ProcessingStep` trait and associated types for building +//! extensible image processing pipelines. + +use crate::config::Config; +use crate::immich_api::FaceData; +use async_trait::async_trait; +use bytes::Bytes; +use image::DynamicImage; +use std::collections::HashMap; +use std::fmt; + +/// Outcome of a pipeline step execution. +#[derive(Debug)] +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 }, + /// An error occurred during processing. + Error(String), +} + +/// Values computed by pipeline steps that can be shared with subsequent steps. +#[derive(Debug, Clone)] +pub enum ComputedValue { + /// A floating-point value (e.g., brightness, EAR). + Float(f32), + /// An integer value (e.g., face size in pixels). + Int(i32), + /// A boolean value (e.g., eyes_open). + Bool(bool), + /// A string value. + String(String), +} + +impl ComputedValue { + /// Get as f32 if this is a Float variant. + pub fn as_float(&self) -> Option { + match self { + ComputedValue::Float(v) => Some(*v), + _ => None, + } + } + + /// Get as i32 if this is an Int variant. + pub fn as_int(&self) -> Option { + match self { + ComputedValue::Int(v) => Some(*v), + _ => None, + } + } + + /// Get as bool if this is a Bool variant. + pub fn as_bool(&self) -> Option { + match self { + ComputedValue::Bool(v) => Some(*v), + _ => None, + } + } + + /// Get as &str if this is a String variant. + pub fn as_str(&self) -> Option<&str> { + match self { + ComputedValue::String(v) => Some(v), + _ => None, + } + } +} + +/// Context passed through the pipeline, carrying data between steps. +#[derive(Debug)] +pub struct PipelineContext { + /// The asset ID from Immich. + pub asset_id: String, + /// Timestamp for sorting/naming output files. + pub timestamp: String, + /// Raw image bytes (before decoding). + pub raw_bytes: Option, + /// Decoded image (after decode step). + pub image: Option, + /// Face bounding box and metadata from Immich. + 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, +} + +impl PipelineContext { + /// Create a new pipeline context. + pub fn new(asset_id: String, timestamp: String, face_data: FaceData) -> Self { + Self { + asset_id, + timestamp, + raw_bytes: None, + image: None, + face_data, + computed: HashMap::new(), + debug_images: HashMap::new(), + } + } + + /// Set the raw image bytes. + pub fn with_bytes(mut self, bytes: Bytes) -> Self { + self.raw_bytes = Some(bytes); + self + } + + /// Set the decoded image. + pub fn with_image(mut self, image: DynamicImage) -> Self { + self.image = Some(image); + self + } + + /// Add a computed value. + pub fn set_computed(&mut self, key: impl Into, value: ComputedValue) { + self.computed.insert(key.into(), value); + } + + /// Get a computed value. + pub fn get_computed(&self, key: &str) -> Option<&ComputedValue> { + self.computed.get(key) + } + + /// 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); + } +} + +/// A single step in the processing pipeline. +/// +/// Steps can be: +/// - **Validators**: Skip the image if a condition fails (return `StepOutcome::Skip`) +/// - **Transformers**: Modify the image (update `ctx.image`) +/// - **Computers**: Calculate values for later steps (update `ctx.computed`) +#[async_trait] +pub trait ProcessingStep: Send + Sync { + /// Unique identifier for this step (used in skip stats and debug output). + fn id(&self) -> &'static str; + + /// Human-readable name for display. + fn name(&self) -> &'static str; + + /// Execute this step on the given context. + /// + /// Returns: + /// - `StepOutcome::Continue(ctx)` to proceed to the next step + /// - `StepOutcome::Skip { reason, detail }` to skip this image + /// - `StepOutcome::Error(msg)` if processing failed + async fn execute(&self, ctx: PipelineContext, config: &Config) -> StepOutcome; + + /// Generate a debug visualization for this step (optional). + /// + /// Called after `execute()` if debug mode is enabled and the step + /// returned `Continue`. The returned image is saved to the debug folder. + fn debug_visualize(&self, _ctx: &PipelineContext) -> Option { + None + } +} + +impl fmt::Debug for dyn ProcessingStep { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ProcessingStep") + .field("id", &self.id()) + .field("name", &self.name()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_computed_value_conversions() { + let float_val = ComputedValue::Float(0.5); + assert_eq!(float_val.as_float(), Some(0.5)); + assert_eq!(float_val.as_int(), None); + + let int_val = ComputedValue::Int(100); + assert_eq!(int_val.as_int(), Some(100)); + assert_eq!(int_val.as_float(), None); + + let bool_val = ComputedValue::Bool(true); + assert_eq!(bool_val.as_bool(), Some(true)); + + let str_val = ComputedValue::String("test".to_string()); + assert_eq!(str_val.as_str(), Some("test")); + } + + #[test] + fn test_pipeline_context_computed_values() { + 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: 1920, + image_height: 1080, + }; + + let mut ctx = PipelineContext::new( + "asset123".to_string(), + "2024-01-15".to_string(), + face_data, + ); + + ctx.set_computed("brightness", ComputedValue::Float(0.65)); + ctx.set_computed("face_size", ComputedValue::Int(150)); + + assert_eq!( + ctx.get_computed("brightness").and_then(|v| v.as_float()), + Some(0.65) + ); + assert_eq!( + ctx.get_computed("face_size").and_then(|v| v.as_int()), + Some(150) + ); + assert!(ctx.get_computed("nonexistent").is_none()); + } +} diff --git a/src/web/handlers/processing.rs b/src/web/handlers/processing.rs index ee0b93f..bca1a53 100644 --- a/src/web/handlers/processing.rs +++ b/src/web/handlers/processing.rs @@ -8,24 +8,31 @@ use axum::{ response::Json, }; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use super::StartResponse; /// Skip statistics for API response. +/// +/// Uses a dynamic HashMap to support arbitrary skip reasons from pipeline steps. #[derive(Serialize)] pub struct SkipStatsResponse { - pub face_too_small: u32, - pub eyes_closed: u32, - pub head_turned: u32, - pub too_dark: u32, - pub too_bright: u32, - pub no_face_detected: u32, - pub download_failed: u32, - pub decode_failed: u32, - pub crop_failed: u32, + /// Map of skip reason ID to count. + #[serde(flatten)] + pub counts: HashMap, + /// Total number of skipped images. pub total: u32, } +impl From<&SkipStats> for SkipStatsResponse { + fn from(stats: &SkipStats) -> Self { + Self { + counts: stats.counts().clone(), + total: stats.total(), + } + } +} + /// Progress response. #[derive(Serialize)] pub struct ProgressResponse { @@ -52,25 +59,12 @@ pub async fn get_progress(State(state): State) -> Json "error", }; - let skip_stats = &progress.skip_stats; - Json(ProgressResponse { status: status_str.to_string(), completed: progress.completed, total: progress.total, message: progress.message.clone(), - skip_stats: SkipStatsResponse { - face_too_small: skip_stats.face_too_small, - eyes_closed: skip_stats.eyes_closed, - head_turned: skip_stats.head_turned, - too_dark: skip_stats.too_dark, - too_bright: skip_stats.too_bright, - no_face_detected: skip_stats.no_face_detected, - download_failed: skip_stats.download_failed, - decode_failed: skip_stats.decode_failed, - crop_failed: skip_stats.crop_failed, - total: skip_stats.total(), - }, + skip_stats: SkipStatsResponse::from(&progress.skip_stats), person_id: progress.person_id.clone(), person_name: progress.person_name.clone(), }) diff --git a/src/web/state.rs b/src/web/state.rs index a16e697..4699a23 100644 --- a/src/web/state.rs +++ b/src/web/state.rs @@ -1,8 +1,9 @@ //! Application state for the web server. use crate::config::Config; +use std::collections::HashMap; use std::sync::atomic::{AtomicU32, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, RwLock as StdRwLock}; use tokio::sync::{broadcast, RwLock}; use tokio_util::sync::CancellationToken; @@ -27,89 +28,99 @@ impl JobStatus { } /// Detailed statistics for skipped images. +/// +/// Uses a HashMap to support dynamic skip reasons from pipeline steps. #[derive(Debug, Clone, Default)] pub struct SkipStats { - pub face_too_small: u32, - pub eyes_closed: u32, - pub head_turned: u32, - pub too_dark: u32, - pub too_bright: u32, - pub no_face_detected: u32, - pub download_failed: u32, - pub decode_failed: u32, - pub crop_failed: u32, + /// Map of skip reason ID to count. + counts: HashMap, } impl SkipStats { + /// Create new empty skip stats. + pub fn new() -> Self { + Self::default() + } + + /// Get the count for a specific reason. + pub fn get(&self, reason: &str) -> u32 { + self.counts.get(reason).copied().unwrap_or(0) + } + + /// Set the count for a specific reason. + pub fn set(&mut self, reason: impl Into, count: u32) { + self.counts.insert(reason.into(), count); + } + + /// Get all counts as a reference to the underlying HashMap. + pub fn counts(&self) -> &HashMap { + &self.counts + } + /// Total number of skipped images. pub fn total(&self) -> u32 { - self.face_too_small - + self.eyes_closed - + self.head_turned - + self.too_dark - + self.too_bright - + self.no_face_detected - + self.download_failed - + self.decode_failed - + self.crop_failed + self.counts.values().sum() } } /// Atomic version of SkipStats for thread-safe concurrent updates. /// -/// This consolidates all skip counters into a single struct that can be -/// shared across async tasks with Arc, avoiding the need to clone multiple -/// individual atomic counters. +/// Uses a HashMap with atomic counters to support dynamic skip reasons +/// from pipeline steps. New reasons are added on first increment. #[derive(Debug, Default)] pub struct AtomicSkipStats { - pub face_too_small: AtomicU32, - pub eyes_closed: AtomicU32, - pub head_turned: AtomicU32, - pub too_dark: AtomicU32, - pub too_bright: AtomicU32, - pub no_face_detected: AtomicU32, - pub download_failed: AtomicU32, - pub decode_failed: AtomicU32, - pub crop_failed: AtomicU32, + /// Map of skip reason to atomic counter. + /// Uses std::sync::RwLock for synchronous access (required for HashMap mutations). + counts: StdRwLock>>, } impl AtomicSkipStats { - /// Create a new AtomicSkipStats with all counters at zero. + /// Create a new AtomicSkipStats with no counters. pub fn new() -> Self { Self::default() } - /// Increment a specific counter and return the new value. - pub fn increment(&self, reason: &crate::face_processing::SkipReason) -> u32 { - use crate::face_processing::SkipReason; - let counter = match reason { - SkipReason::FaceTooSmall => &self.face_too_small, - SkipReason::EyesClosed => &self.eyes_closed, - SkipReason::HeadTurned => &self.head_turned, - SkipReason::TooDark => &self.too_dark, - SkipReason::TooBright => &self.too_bright, - SkipReason::NoFaceDetected => &self.no_face_detected, - SkipReason::DownloadFailed => &self.download_failed, - SkipReason::DecodeFailed => &self.decode_failed, - SkipReason::CropFailed => &self.crop_failed, - SkipReason::Cancelled => return 0, // Don't count cancellation - }; - counter.fetch_add(1, Ordering::SeqCst) + 1 + /// Increment a counter for the given reason string. + /// + /// If the reason doesn't exist yet, it's created with an initial count of 1. + /// Returns the new count value. + pub fn increment(&self, reason: &str) -> u32 { + // First, try to get existing counter with read lock + { + let counts = self.counts.read().unwrap(); + if let Some(counter) = counts.get(reason) { + return counter.fetch_add(1, Ordering::SeqCst) + 1; + } + } + + // Counter doesn't exist, need write lock to create it + let mut counts = self.counts.write().unwrap(); + + // Double-check after acquiring write lock (another thread may have created it) + if let Some(counter) = counts.get(reason) { + return counter.fetch_add(1, Ordering::SeqCst) + 1; + } + + // Create new counter starting at 1 + let counter = Arc::new(AtomicU32::new(1)); + counts.insert(reason.to_string(), counter); + 1 } /// Get a snapshot of the current stats as a non-atomic SkipStats. pub fn snapshot(&self) -> SkipStats { - SkipStats { - face_too_small: self.face_too_small.load(Ordering::SeqCst), - eyes_closed: self.eyes_closed.load(Ordering::SeqCst), - head_turned: self.head_turned.load(Ordering::SeqCst), - too_dark: self.too_dark.load(Ordering::SeqCst), - too_bright: self.too_bright.load(Ordering::SeqCst), - no_face_detected: self.no_face_detected.load(Ordering::SeqCst), - download_failed: self.download_failed.load(Ordering::SeqCst), - decode_failed: self.decode_failed.load(Ordering::SeqCst), - crop_failed: self.crop_failed.load(Ordering::SeqCst), + let counts = self.counts.read().unwrap(); + let mut stats = SkipStats::new(); + for (reason, counter) in counts.iter() { + stats.set(reason.clone(), counter.load(Ordering::SeqCst)); } + stats + } + + /// Reset all counters to zero. + pub fn reset(&self) { + let mut counts = self.counts.write().unwrap(); + counts.clear(); } }