danvics/immich-automated-selfie-timelapse
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Changed the face processing to generic tasks with traits. No need to hardcode them anymore.

Browse source
This commit is contained in:
Arnaud_Cayrol 2026-01-31 14:51:40 +01:00
parent 71b58db82a
commit 41d50167eb
16 changed files with 1544 additions and 209 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
Cargo.toml
View file

@ -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"

105
frontend/src/lib/components/SettingsPanel.svelte
View file

@ -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 @@
<span class="value">{config.processing.pose_threshold}°</span>
</div>
</div>
<!-- Brightness validation section -->
<div class="setting-section">
<div class="section-header">
<span class="section-title">Brightness Filter</span>
<input
type="checkbox"
checked={config.processing.brightness?.enabled ?? false}
onchange={(e) => {
ensureBrightnessConfig();
config.processing.brightness.enabled = e.target.checked;
}}
/>
</div>
{#if config.processing.brightness?.enabled}
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Min Brightness</span>
<span class="setting-hint">Skip images darker than this</span>
</label>
<div class="setting-control">
<input
type="range"
value={config.processing.brightness?.min_brightness ?? 0.1}
oninput={(e) => {
ensureBrightnessConfig();
config.processing.brightness.min_brightness = parseFloat(e.target.value);
}}
min="0"
max="0.5"
step="0.05"
/>
<span class="value">{(config.processing.brightness?.min_brightness ?? 0.1).toFixed(2)}</span>
</div>
</div>
<div class="setting-row sub-setting">
<label>
<span class="setting-label">Max Brightness</span>
<span class="setting-hint">Skip images brighter than this</span>
</label>
<div class="setting-control">
<input
type="range"
value={config.processing.brightness?.max_brightness ?? 0.95}
oninput={(e) => {
ensureBrightnessConfig();
config.processing.brightness.max_brightness = parseFloat(e.target.value);
}}
min="0.5"
max="1.0"
step="0.05"
/>
<span class="value">{(config.processing.brightness?.max_brightness ?? 0.95).toFixed(2)}</span>
</div>
</div>
{/if}
</div>
</fieldset>
{:else if activeTab === 'output'}
<fieldset disabled={disabled || saving}>
@ -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;

30
src/config.rs
View file

@ -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<BrightnessConfig>,
}
/// 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,
}
}
}

233
src/job/mod.rs
View file

@ -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<FaceData> {
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<AtomicSkipStats>,
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)]

1
src/lib.rs
View file

@ -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;

8
src/main.rs
View file

@ -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(

238
src/pipeline/mod.rs Normal file
View file

@ -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<String>,
},
/// 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<Box<dyn ProcessingStep>>,
}
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::<Vec<_>>())
.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<dyn ProcessingStep>) {
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<AtomicSkipStats>,
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::<String>();
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"));
}
}

227
src/pipeline/steps/brightness.rs Normal file
View file

@ -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"),
}
}
}

144
src/pipeline/steps/crop.rs Normal file
View file

@ -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<DynamicImage> {
// 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"),
}
}
}

133
src/pipeline/steps/decode.rs Normal file
View file

@ -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"),
}
}
}

98
src/pipeline/steps/face_resolution.rs Normal file
View file

@ -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"),
}
}
}

16
src/pipeline/steps/mod.rs Normal file
View file

@ -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;

127
src/pipeline/steps/resize.rs Normal file
View file

@ -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"),
}
}
}

224
src/pipeline/traits.rs Normal file
View file

@ -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<String> },
/// 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<f32> {
match self {
ComputedValue::Float(v) => Some(*v),
_ => None,
}
}
/// Get as i32 if this is an Int variant.
pub fn as_int(&self) -> Option<i32> {
match self {
ComputedValue::Int(v) => Some(*v),
_ => None,
}
}
/// Get as bool if this is a Bool variant.
pub fn as_bool(&self) -> Option<bool> {
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<Bytes>,
/// Decoded image (after decode step).
pub image: Option<DynamicImage>,
/// Face bounding box and metadata from Immich.
pub face_data: FaceData,
/// Values computed by previous steps (e.g., brightness, landmarks).
pub computed: HashMap<String, ComputedValue>,
/// Debug images generated by steps (step_id -> debug image).
pub debug_images: HashMap<String, DynamicImage>,
}
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<String>, 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<String>, 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<DynamicImage> {
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());
}
}

40
src/web/handlers/processing.rs
View file

@ -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<String, u32>,
/// 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<AppState>) -> Json<ProgressRespons
JobStatus::Error(_) => "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(),
})

127
src/web/state.rs
View file

@ -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<String, u32>,
}
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<String>, count: u32) {
self.counts.insert(reason.into(), count);
}
/// Get all counts as a reference to the underlying HashMap.
pub fn counts(&self) -> &HashMap<String, u32> {
&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<HashMap<String, Arc<AtomicU32>>>,
}
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();
}
}