diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 59f20d6c89..41ae51fc4c 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -35,9 +35,14 @@ on: required: false type: string default: '' + android-signing: + required: false + type: string + default: '' jobs: build: + if: ${{ inputs.mobile != 'android' || inputs.android-signing != 'false' }} runs-on: ${{ inputs.platform }} steps: - name: Checkout repository diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2c9e8e7d7c..fe178d09d8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,7 +6,24 @@ on: - main jobs: + # Forks do not have the Android signing secrets, so the Android build is skipped there. + android-signing: + runs-on: ubuntu-latest + outputs: + available: ${{ steps.check.outputs.available }} + steps: + - id: check + env: + ANDROID_KEY_BASE64: ${{ secrets.ANDROID_KEY_BASE64 }} + run: | + if [ -n "$ANDROID_KEY_BASE64" ]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + fi + build: + needs: android-signing strategy: fail-fast: false matrix: @@ -65,6 +82,7 @@ jobs: build-args: ${{ matrix.args }} asset-prefix: ${{ matrix.asset-prefix }} mobile: ${{ matrix.mobile }} + android-signing: ${{ needs.android-signing.outputs.available }} secrets: inherit permissions: contents: write diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 457518fce2..a1335c5cb0 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -5,7 +5,24 @@ on: types: [created] jobs: + # Forks do not have the Android signing secrets, so the Android build is skipped there. + android-signing: + runs-on: ubuntu-latest + outputs: + available: ${{ steps.check.outputs.available }} + steps: + - id: check + env: + ANDROID_KEY_BASE64: ${{ secrets.ANDROID_KEY_BASE64 }} + run: | + if [ -n "$ANDROID_KEY_BASE64" ]; then + echo "available=true" >> "$GITHUB_OUTPUT" + else + echo "available=false" >> "$GITHUB_OUTPUT" + fi + release: + needs: android-signing strategy: fail-fast: false matrix: @@ -66,6 +83,7 @@ jobs: asset-name-pattern: '[name]_v[version]_[platform]_[arch][ext]' asset-prefix: ${{ matrix.asset-prefix }} mobile: ${{ matrix.mobile }} + android-signing: ${{ needs.android-signing.outputs.available }} secrets: inherit permissions: contents: write diff --git a/src-tauri/src/ai_denoise.rs b/src-tauri/src/ai_denoise.rs index b76fd8c878..b0b328458a 100644 --- a/src-tauri/src/ai_denoise.rs +++ b/src-tauri/src/ai_denoise.rs @@ -307,6 +307,8 @@ pub fn apply_ai_denoise<'a>( (adjustments["aiDenoiseAmount"].as_f64().unwrap_or(100.0) as f32 / 100.0).clamp(0.0, 1.0); let detail = (adjustments["aiDenoiseDetail"].as_f64().unwrap_or(0.0) as f32 / 100.0).clamp(0.0, 1.0); + let color = + (adjustments["aiDenoiseColor"].as_f64().unwrap_or(100.0) as f32 / 100.0).clamp(0.0, 1.0); if amount <= 0.0 { return Cow::Borrowed(base); } @@ -336,7 +338,7 @@ pub fn apply_ai_denoise<'a>( let denoised: &[f16] = resized.as_deref().unwrap_or(&entry.data); let luma_keep = 1.0 - amount * (1.0 - detail); - let chroma_keep = 1.0 - amount; + let chroma_keep = 1.0 - amount * color; let mut out = base.to_rgb32f(); out.par_chunks_mut(BLEND_CHUNK) diff --git a/src-tauri/src/ai_processing.rs b/src-tauri/src/ai_processing.rs index 04490cc2bd..36a81bf328 100644 --- a/src-tauri/src/ai_processing.rs +++ b/src-tauri/src/ai_processing.rs @@ -26,6 +26,9 @@ use tauri::Manager; use tokenizers::Tokenizer; use tokio::sync::Mutex as TokioMutex; +use crate::face_parsing::{FaceModels, load_face_parser_session}; +use crate::face_refine::FaceCropCache; + use crate::face_masking::FaceAnalysis; const ENCODER_URL: &str = "https://huggingface.co/CyberTimon/RapidRAW-Models/resolve/main/sam_vit_b_01ec64_encoder.onnx?download=true"; @@ -124,6 +127,16 @@ const NORMAL_MAX_SIZE: u32 = 768; const NORMAL_NUM_TOKENS: i64 = 2500; const NORMAL_SHA256: &str = "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0"; +const FACE_DETECTOR_URL: &str = "https://huggingface.co/opencv/face_detection_yunet/resolve/main/face_detection_yunet_2023mar.onnx?download=true"; +const FACE_DETECTOR_FILENAME: &str = "face_detection_yunet_2023mar.onnx"; +const FACE_DETECTOR_SHA256: &str = + "8f2383e4dd3cfbb4553ea8718107fc0423210dc964f9f4280604804ed2552fa4"; + +const FACE_PARSER_URL: &str = + "https://github.com/yakhyo/face-parsing/releases/download/v0.0.1/resnet18.onnx"; +const FACE_PARSER_FILENAME: &str = "face_parsing_resnet18.onnx"; +const FACE_PARSER_SHA256: &str = "0d9bd318e46987c3bdbfacae9e2c0f461cae1c6ac6ea6d43bbe541a91727e33f"; + type SharedSession = Arc>; const GPU_IDLE: u8 = 0; @@ -140,6 +153,20 @@ struct ModelAsset { name: &'static str, } +const FACE_DETECTOR_ASSET: ModelAsset = ModelAsset { + filename: FACE_DETECTOR_FILENAME, + url: FACE_DETECTOR_URL, + sha256: Some(FACE_DETECTOR_SHA256), + name: "Face Detection Model", +}; + +const FACE_PARSER_ASSET: ModelAsset = ModelAsset { + filename: FACE_PARSER_FILENAME, + url: FACE_PARSER_URL, + sha256: Some(FACE_PARSER_SHA256), + name: "Face Regions Model", +}; + const SAM_ENCODER_ASSET: ModelAsset = ModelAsset { filename: ENCODER_FILENAME, url: ENCODER_URL, @@ -691,6 +718,8 @@ pub struct AiState { pub esam3_features: Option>, pub face_analysis: Option>, pub depth_map: Option, + pub face_models: Option>, + pub face_crops: Option>, } impl AiState { @@ -1356,6 +1385,51 @@ pub async fn get_or_init_lama_model( Ok(lama_model) } +pub async fn get_or_init_face_models( + app_handle: &tauri::AppHandle, + ai_state_mutex: &Mutex>, + ai_init_lock: &TokioMutex<()>, +) -> Result> { + let loaded = || { + ai_state_mutex + .lock() + .unwrap() + .as_ref() + .and_then(|state| state.face_models.clone()) + }; + if let Some(face_models) = loaded() { + return Ok(face_models); + } + + let _guard = ai_init_lock.lock().await; + + if let Some(face_models) = loaded() { + return Ok(face_models); + } + + let models_dir = get_models_dir(app_handle)?; + download_and_verify_model(app_handle, &models_dir, &FACE_DETECTOR_ASSET).await?; + download_and_verify_model(app_handle, &models_dir, &FACE_PARSER_ASSET).await?; + + let _ = ort::init().with_name("AI-Face").commit(); + let detector = Session::builder()?.commit_from_file(models_dir.join(FACE_DETECTOR_FILENAME))?; + let parser = load_face_parser_session(&models_dir.join(FACE_PARSER_FILENAME))?; + let face_models = Arc::new(FaceModels { + detector: Mutex::new(detector), + parser: Mutex::new(parser), + }); + + crate::register_exit_handler(); + + ai_state_mutex + .lock() + .unwrap() + .get_or_insert_with(AiState::default) + .face_models = Some(face_models.clone()); + + Ok(face_models) +} + pub async fn get_or_init_normal_model( app_handle: &tauri::AppHandle, ai_state_mutex: &Mutex>, diff --git a/src-tauri/src/app_settings.rs b/src-tauri/src/app_settings.rs index de6417b275..40581bcf8c 100644 --- a/src-tauri/src/app_settings.rs +++ b/src-tauri/src/app_settings.rs @@ -153,6 +153,7 @@ pub fn all_available_adjustments() -> HashSet { "aiDenoiseEnabled", "aiDenoiseMethod", "aiDenoiseAmount", + "aiDenoiseColor", "aiDenoiseDetail", "chromaticAberrationRedCyan", "chromaticAberrationBlueYellow", diff --git a/src-tauri/src/cache_utils.rs b/src-tauri/src/cache_utils.rs index b662d20908..30f83396f5 100644 --- a/src-tauri/src/cache_utils.rs +++ b/src-tauri/src/cache_utils.rs @@ -94,7 +94,12 @@ fn hash_ai_denoise(adjustments: &serde_json::Value, hasher: &mut DefaultHasher) if enabled { hash_effect_keys( adjustments, - &["aiDenoiseMethod", "aiDenoiseAmount", "aiDenoiseDetail"], + &[ + "aiDenoiseMethod", + "aiDenoiseAmount", + "aiDenoiseColor", + "aiDenoiseDetail", + ], hasher, ); crate::ai_denoise::cache_epoch().hash(hasher); diff --git a/src-tauri/src/face_parsing.rs b/src-tauri/src/face_parsing.rs new file mode 100644 index 0000000000..3d142649ab --- /dev/null +++ b/src-tauri/src/face_parsing.rs @@ -0,0 +1,767 @@ +use std::path::Path; +use std::sync::Mutex; + +use anyhow::{Result, anyhow}; +use image::RgbImage; +use image::imageops::{self, FilterType}; +use ndarray::Array4; +use ort::session::Session; +use ort::value::Tensor; +use rayon::prelude::*; + +const DETECTOR_SIZE: u32 = 640; +const DETECTOR_SCORE_THRESHOLD: f32 = 0.6; +const DETECTOR_CONFIDENT_SCORE: f32 = 0.75; +const DETECTOR_NMS_IOU: f32 = 0.3; +const MAX_FACES: usize = 10; + +const PARSER_SIZE: usize = 512; +const PARSER_CLASSES: usize = 19; + +// CelebAMask-HQ class indices produced by the parsing model. +const SKIN: usize = 1; +const LEFT_BROW: usize = 2; +const RIGHT_BROW: usize = 3; +const LEFT_EYE: usize = 4; +const RIGHT_EYE: usize = 5; +const GLASSES: usize = 6; +const LEFT_EAR: usize = 7; +const RIGHT_EAR: usize = 8; +const EARRING: usize = 9; +const NOSE: usize = 10; +const MOUTH: usize = 11; +const UPPER_LIP: usize = 12; +const LOWER_LIP: usize = 13; +const NECK: usize = 14; +const HAIR: usize = 17; +const HAT: usize = 18; + +pub struct FaceModels { + pub detector: Mutex, + pub parser: Mutex, +} + +/// Loads the parsing model on the GPU when that works and gives the same answer as the +/// CPU, and on the CPU otherwise. +pub fn load_face_parser_session(model_path: &Path) -> Result { + let cpu_session = Session::builder()?.commit_from_file(model_path)?; + + #[cfg(target_os = "windows")] + let cpu_session = { + let mut cpu_session = cpu_session; + match build_gpu_parser_session(model_path, &mut cpu_session) { + Ok(session) => { + log::info!("Face parsing model running on GPU (DirectML)"); + return Ok(session); + } + Err(e) => log::warn!("GPU face parsing unavailable, falling back to CPU: {}", e), + } + cpu_session + }; + + Ok(cpu_session) +} + +#[cfg(target_os = "windows")] +fn build_gpu_parser_session(model_path: &Path, cpu_session: &mut Session) -> Result { + use ort::execution_providers::DirectMLExecutionProvider; + + let mut gpu_session = Session::builder()? + .with_memory_pattern(false)? + .with_parallel_execution(false)? + .with_execution_providers([DirectMLExecutionProvider::default() + .build() + .error_on_failure()])? + .commit_from_file(model_path)?; + + // A fixed, face-free pattern is enough to tell a broken GPU path from a working one. + let probe = Array4::::from_shape_fn((1, 3, PARSER_SIZE, PARSER_SIZE), |(_, c, y, x)| { + ((x * 7 + y * 13 + c * 29) % 97) as f32 / 48.0 - 1.0 + }); + let run = |session: &mut Session| -> Result> { + let tensor = Tensor::from_array(probe.clone().into_dyn())?; + let outputs = session.run(ort::inputs![tensor])?; + Ok(outputs[0] + .try_extract_array::()? + .iter() + .copied() + .collect()) + }; + let expected = run(cpu_session)?; + let actual = run(&mut gpu_session)?; + + let area = PARSER_SIZE * PARSER_SIZE; + if actual.len() != expected.len() || actual.len() != PARSER_CLASSES * area { + return Err(anyhow!("GPU face parsing returned an unexpected shape")); + } + let winner = |logits: &[f32], i: usize| { + (0..PARSER_CLASSES) + .max_by(|&a, &b| logits[a * area + i].total_cmp(&logits[b * area + i])) + .unwrap_or(0) + }; + let disagreements = (0..area) + .filter(|&i| winner(&expected, i) != winner(&actual, i)) + .count(); + if disagreements * 100 > area { + return Err(anyhow!( + "GPU face parsing self-check failed ({} of {} pixels differ)", + disagreements, + area + )); + } + Ok(gpu_session) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FaceRegion { + Skin, + Eyes, + EyeWhites, + Teeth, + /// Everything on a head that is not skin. Used to protect features while retouching. + Features, + /// Forehead and cheeks: the skin where spots are removed. The nose, the nasolabial + /// folds, the area around the mouth, the chin, the ears and the neck are left out. + BlemishZone, +} + +impl FaceRegion { + fn classes(self) -> &'static [usize] { + match self { + Self::Skin => &[SKIN, NOSE, LEFT_EAR, RIGHT_EAR, NECK], + Self::Eyes | Self::EyeWhites => &[LEFT_EYE, RIGHT_EYE], + Self::Teeth => &[MOUTH], + Self::BlemishZone => &[SKIN], + Self::Features => &[ + LEFT_BROW, RIGHT_BROW, LEFT_EYE, RIGHT_EYE, GLASSES, EARRING, MOUTH, UPPER_LIP, + LOWER_LIP, HAIR, HAT, + ], + } + } + + /// Luma and saturation limits that separate teeth from the rest of the mouth and + /// the whites of the eyes from the iris. + fn bright_neutral_limits(self) -> Option<([f32; 2], [f32; 2])> { + match self { + Self::Teeth => Some(([0.18, 0.38], [0.30, 0.55])), + Self::EyeWhites => Some(([0.25, 0.45], [0.18, 0.38])), + _ => None, + } + } +} + +#[derive(Debug, Clone, Copy)] +struct Detection { + score: f32, + /// x, y, width, height + bbox: [f32; 4], + /// right eye, left eye, nose, right mouth corner, left mouth corner + landmarks: [[f32; 2]; 5], +} + +impl Detection { + fn map_points(&self, f: impl Fn([f32; 2]) -> [f32; 2]) -> Self { + let [x, y, w, h] = self.bbox; + let a = f([x, y]); + let b = f([x + w, y + h]); + Self { + score: self.score, + bbox: [ + a[0].min(b[0]), + a[1].min(b[1]), + (a[0] - b[0]).abs(), + (a[1] - b[1]).abs(), + ], + landmarks: self.landmarks.map(&f), + } + } + + fn iou(&self, other: &Self) -> f32 { + let [ax, ay, aw, ah] = self.bbox; + let [bx, by, bw, bh] = other.bbox; + let iw = ((ax + aw).min(bx + bw) - ax.max(bx)).max(0.0); + let ih = ((ay + ah).min(by + bh) - ay.max(by)).max(0.0); + let inter = iw * ih; + inter / (aw * ah + bw * bh - inter).max(1e-6) + } +} + +/// A face cut out of the image, upright and scaled the way the parsing model expects, +/// together with the per-class probabilities the model returned for it. +#[derive(Clone)] +pub struct ParsedFace { + center: [f32; 2], + x_dir: [f32; 2], + half: f32, + /// Approximate width of the face in image pixels. + pub width: f32, + /// The detector's landmarks in crop coordinates: right eye, left eye, nose, right + /// mouth corner, left mouth corner. + landmarks: [[f32; 2]; 5], + /// `PARSER_CLASSES` planes of `PARSER_SIZE * PARSER_SIZE` probabilities. + probs: Vec, +} + +impl ParsedFace { + fn y_dir(&self) -> [f32; 2] { + [-self.x_dir[1], self.x_dir[0]] + } + + fn crop_to_image(&self, u: f32, v: f32) -> [f32; 2] { + let half_size = PARSER_SIZE as f32 / 2.0; + let a = (u / half_size - 1.0) * self.half; + let b = (v / half_size - 1.0) * self.half; + let y_dir = self.y_dir(); + [ + self.center[0] + self.x_dir[0] * a + y_dir[0] * b, + self.center[1] + self.x_dir[1] * a + y_dir[1] * b, + ] + } + + fn image_to_crop(&self, point: [f32; 2]) -> [f32; 2] { + let half_size = PARSER_SIZE as f32 / 2.0; + let (dx, dy) = (point[0] - self.center[0], point[1] - self.center[1]); + let y_dir = self.y_dir(); + let a = (dx * self.x_dir[0] + dy * self.x_dir[1]) / self.half; + let b = (dx * y_dir[0] + dy * y_dir[1]) / self.half; + [(a + 1.0) * half_size, (b + 1.0) * half_size] + } + + /// How much the crop position (`u`, `v`) belongs to the forehead or the cheeks. + fn blemish_zone_weight(&self, u: f32, v: f32) -> f32 { + let [right_eye, left_eye, nose, right_mouth, left_mouth] = self.landmarks; + let eye_y = (right_eye[1] + left_eye[1]) * 0.5; + let mouth_y = (right_mouth[1] + left_mouth[1]) * 0.5; + // The eye distance collapses on turned heads, so the face height backs it up. + let unit = (left_eye[0] - right_eye[0]) + .abs() + .max((mouth_y - eye_y) * 0.8) + .max(8.0); + + let forehead = 1.0 - smoothstep(eye_y - 0.45 * unit, eye_y - 0.30 * unit, v); + let cheek_top = eye_y + 0.20 * unit; + let below_eyes = smoothstep(cheek_top, eye_y + 0.35 * unit, v); + let above_mouth = 1.0 - smoothstep(mouth_y - 0.15 * unit, mouth_y, v); + + // The cheek starts outside the nasolabial fold, which runs from the wing of the + // nose down to just outside the corner of the mouth. + let mouth_corner = if u < nose[0] { right_mouth } else { left_mouth }; + let at_nose = 0.50 * unit; + let at_mouth = (mouth_corner[0] - nose[0]).abs() + 0.22 * unit; + let fold = if v <= nose[1] { + let t = ((v - cheek_top) / (nose[1] - cheek_top).max(1.0)).clamp(0.0, 1.0); + 0.30 * unit + t * (at_nose - 0.30 * unit) + } else { + let t = ((v - nose[1]) / (mouth_y - nose[1]).max(1.0)).clamp(0.0, 1.0); + at_nose + t * (at_mouth - at_nose) + }; + let outside_fold = smoothstep(fold, fold + 0.12 * unit, (u - nose[0]).abs()); + + forehead.max(below_eyes * above_mouth * outside_fold) + } + + /// Image-space position that identifies this face. + pub fn center(&self) -> [f32; 2] { + self.center + } + + /// Small upright portrait of the face, for listing it in the interface. + pub fn thumbnail(&self, image: &RgbImage, size: u32) -> RgbImage { + let (width, height) = image.dimensions(); + RgbImage::from_fn(size, size, |i, j| { + let to_crop = + |t: u32| ((t as f32 + 0.5) / size as f32 * 0.6 + 0.2) * PARSER_SIZE as f32; + let [x, y] = self.crop_to_image(to_crop(i), to_crop(j)); + if x >= 0.0 && y >= 0.0 && x < width as f32 && y < height as f32 { + *image.get_pixel(x as u32, y as u32) + } else { + image::Rgb([0, 0, 0]) + } + }) + } + + /// Image-space bounding box of the analyzed crop: min x, min y, max x, max y. + pub fn bounds(&self) -> [f32; 4] { + let size = PARSER_SIZE as f32; + let corners = [[0.0, 0.0], [size, 0.0], [0.0, size], [size, size]] + .map(|[u, v]| self.crop_to_image(u, v)); + let mut b = [f32::MAX, f32::MAX, f32::MIN, f32::MIN]; + for [x, y] in corners { + b = [b[0].min(x), b[1].min(y), b[2].max(x), b[3].max(y)]; + } + b + } + + fn region_plane(&self, region: FaceRegion) -> Vec { + let area = PARSER_SIZE * PARSER_SIZE; + let mut plane = vec![0.0_f32; area]; + for &class in region.classes() { + let probs = &self.probs[class * area..(class + 1) * area]; + plane + .iter_mut() + .zip(probs) + .for_each(|(out, &p)| *out += p as f32 / 255.0); + } + plane.iter_mut().for_each(|v| *v = v.min(1.0)); + if region == FaceRegion::BlemishZone { + for (i, value) in plane.iter_mut().enumerate() { + let (u, v) = ((i % PARSER_SIZE) as f32, (i / PARSER_SIZE) as f32); + *value *= self.blemish_zone_weight(u + 0.5, v + 0.5); + } + } + plane + } +} + +fn run_detector(session: &Mutex, canvas: &RgbImage) -> Result> { + let size = DETECTOR_SIZE as usize; + let mut input = Array4::::zeros((1, 3, size, size)); + for (x, y, pixel) in canvas.enumerate_pixels() { + let (x, y) = (x as usize, y as usize); + input[[0, 0, y, x]] = pixel[2] as f32; + input[[0, 1, y, x]] = pixel[1] as f32; + input[[0, 2, y, x]] = pixel[0] as f32; + } + let tensor = Tensor::from_array(input.into_dyn())?; + + let mut session = session.lock().unwrap(); + let outputs = session.run(ort::inputs![tensor])?; + + let mut detections = Vec::new(); + for stride in [8_usize, 16, 32] { + let cols = size / stride; + let read = |name: String| -> Result> { + Ok(outputs[name.as_str()] + .try_extract_array::()? + .iter() + .copied() + .collect()) + }; + let cls = read(format!("cls_{stride}"))?; + let obj = read(format!("obj_{stride}"))?; + let bbox = read(format!("bbox_{stride}"))?; + let kps = read(format!("kps_{stride}"))?; + if cls.len() != cols * cols || bbox.len() != cls.len() * 4 || kps.len() != cls.len() * 10 { + return Err(anyhow!("Unexpected face detector output shape")); + } + + let stride_f = stride as f32; + for i in 0..cls.len() { + let score = (cls[i].clamp(0.0, 1.0) * obj[i].clamp(0.0, 1.0)).sqrt(); + if score < DETECTOR_SCORE_THRESHOLD { + continue; + } + let col = (i % cols) as f32; + let row = (i / cols) as f32; + let b = &bbox[i * 4..i * 4 + 4]; + let cx = (col + b[0]) * stride_f; + let cy = (row + b[1]) * stride_f; + let w = b[2].exp() * stride_f; + let h = b[3].exp() * stride_f; + + let k = &kps[i * 10..i * 10 + 10]; + let mut landmarks = [[0.0_f32; 2]; 5]; + for (n, point) in landmarks.iter_mut().enumerate() { + *point = [(k[n * 2] + col) * stride_f, (k[n * 2 + 1] + row) * stride_f]; + } + detections.push(Detection { + score, + bbox: [cx - w / 2.0, cy - h / 2.0, w, h], + landmarks, + }); + } + } + Ok(detections) +} + +/// Finds faces in any of the four 90 degree orientations, since raw files are analyzed +/// before their orientation is applied. +fn detect_faces(image: &RgbImage, session: &Mutex) -> Result> { + let (width, height) = image.dimensions(); + let scale = DETECTOR_SIZE as f32 / width.max(height) as f32; + let small_w = ((width as f32 * scale).round() as u32).clamp(1, DETECTOR_SIZE); + let small_h = ((height as f32 * scale).round() as u32).clamp(1, DETECTOR_SIZE); + let small = imageops::resize(image, small_w, small_h, FilterType::Triangle); + let (sw, sh) = (small_w as f32, small_h as f32); + + let mut found: Vec = Vec::new(); + for turns in 0..4 { + let rotated = match turns { + 0 => small.clone(), + 1 => imageops::rotate90(&small), + 2 => imageops::rotate180(&small), + _ => imageops::rotate270(&small), + }; + let mut canvas = RgbImage::new(DETECTOR_SIZE, DETECTOR_SIZE); + imageops::replace(&mut canvas, &rotated, 0, 0); + + for detection in run_detector(session, &canvas)? { + let upright = detection.map_points(|[x, y]| match turns { + 0 => [x, y], + 1 => [y, sh - x], + 2 => [sw - x, sh - y], + _ => [sw - y, x], + }); + found.push(upright.map_points(|[x, y]| [x / scale, y / scale])); + } + + if turns == 0 && found.iter().any(|d| d.score >= DETECTOR_CONFIDENT_SCORE) { + break; + } + } + + found.sort_by(|a, b| b.score.total_cmp(&a.score)); + let mut kept: Vec = Vec::new(); + for detection in found { + if kept.iter().all(|k| k.iou(&detection) <= DETECTOR_NMS_IOU) { + kept.push(detection); + } + } + + let min_size = width.max(height) as f32 * 0.02; + kept.retain(|d| d.bbox[2].max(d.bbox[3]) >= min_size); + kept.sort_by(|a, b| (b.bbox[2] * b.bbox[3]).total_cmp(&(a.bbox[2] * a.bbox[3]))); + kept.truncate(MAX_FACES); + Ok(kept) +} + +/// Same framing as the aligned portraits the parsing model was trained on. +fn align_face(detection: &Detection) -> ParsedFace { + let [right_eye, left_eye, _, right_mouth, left_mouth] = detection.landmarks; + let eye_avg = [ + (right_eye[0] + left_eye[0]) * 0.5, + (right_eye[1] + left_eye[1]) * 0.5, + ]; + let mouth_avg = [ + (right_mouth[0] + left_mouth[0]) * 0.5, + (right_mouth[1] + left_mouth[1]) * 0.5, + ]; + let eye_to_eye = [left_eye[0] - right_eye[0], left_eye[1] - right_eye[1]]; + let eye_to_mouth = [mouth_avg[0] - eye_avg[0], mouth_avg[1] - eye_avg[1]]; + + let x = [ + eye_to_eye[0] + eye_to_mouth[1], + eye_to_eye[1] - eye_to_mouth[0], + ]; + let x_len = x[0].hypot(x[1]).max(1e-3); + let eye_dist = eye_to_eye[0].hypot(eye_to_eye[1]); + let mouth_dist = eye_to_mouth[0].hypot(eye_to_mouth[1]); + + let mut face = ParsedFace { + center: [ + eye_avg[0] + eye_to_mouth[0] * 0.1, + eye_avg[1] + eye_to_mouth[1] * 0.1, + ], + x_dir: [x[0] / x_len, x[1] / x_len], + half: (1.125 * (eye_dist * 2.0).max(mouth_dist * 1.8)).max(8.0), + width: detection.bbox[2].max(detection.bbox[3] * 0.75), + landmarks: [[0.0; 2]; 5], + probs: Vec::new(), + }; + face.landmarks = detection.landmarks.map(|point| face.image_to_crop(point)); + face +} + +fn parse_face(image: &RgbImage, face: &mut ParsedFace, session: &Mutex) -> Result<()> { + const MEAN: [f32; 3] = [0.485, 0.456, 0.406]; + const STD: [f32; 3] = [0.229, 0.224, 0.225]; + + let (width, height) = image.dimensions(); + let raw = image.as_raw(); + let samples = ((face.half * 2.0 / PARSER_SIZE as f32).ceil() as usize).clamp(1, 4); + let area = PARSER_SIZE * PARSER_SIZE; + + let mut pixels = vec![[0.0_f32; 3]; area]; + pixels + .par_chunks_exact_mut(PARSER_SIZE) + .enumerate() + .for_each(|(v, row)| { + for (u, out) in row.iter_mut().enumerate() { + let mut sum = [0.0_f32; 3]; + for j in 0..samples { + for i in 0..samples { + let [x, y] = face.crop_to_image( + u as f32 + (i as f32 + 0.5) / samples as f32, + v as f32 + (j as f32 + 0.5) / samples as f32, + ); + if x >= 0.0 && y >= 0.0 && x < width as f32 && y < height as f32 { + let idx = (y as usize * width as usize + x as usize) * 3; + for c in 0..3 { + sum[c] += raw[idx + c] as f32; + } + } + } + } + let norm = 255.0 * (samples * samples) as f32; + for c in 0..3 { + out[c] = (sum[c] / norm - MEAN[c]) / STD[c]; + } + } + }); + + let mut input = Array4::::zeros((1, 3, PARSER_SIZE, PARSER_SIZE)); + for (i, pixel) in pixels.iter().enumerate() { + let (v, u) = (i / PARSER_SIZE, i % PARSER_SIZE); + for c in 0..3 { + input[[0, c, v, u]] = pixel[c]; + } + } + let tensor = Tensor::from_array(input.into_dyn())?; + + let logits: Vec = { + let mut session = session.lock().unwrap(); + let outputs = session.run(ort::inputs![tensor])?; + outputs[0] + .try_extract_array::()? + .iter() + .copied() + .collect() + }; + if logits.len() != PARSER_CLASSES * area { + return Err(anyhow!("Unexpected face parsing output shape")); + } + + let mut probs = vec![0u8; PARSER_CLASSES * area]; + for i in 0..area { + let mut max = f32::MIN; + for c in 0..PARSER_CLASSES { + max = max.max(logits[c * area + i]); + } + let mut exp = [0.0_f32; PARSER_CLASSES]; + let mut sum = 0.0; + for c in 0..PARSER_CLASSES { + exp[c] = (logits[c * area + i] - max).exp(); + sum += exp[c]; + } + for c in 0..PARSER_CLASSES { + probs[c * area + i] = (exp[c] / sum * 255.0).round() as u8; + } + } + face.probs = probs; + Ok(()) +} + +pub fn analyze_faces(image: &RgbImage, models: &FaceModels) -> Result> { + let detections = detect_faces(image, &models.detector)?; + let mut faces = Vec::with_capacity(detections.len()); + for detection in &detections { + let mut face = align_face(detection); + parse_face(image, &mut face, &models.parser)?; + faces.push(face); + } + Ok(faces) +} + +/// Renders the probability (0..1) of `region` for the `width` x `height` window of the +/// image whose top-left corner is at `origin`. Faces are combined with a maximum. +pub fn render_region( + faces: &[ParsedFace], + region: FaceRegion, + origin: (u32, u32), + width: u32, + height: u32, +) -> Vec { + let (w, h) = (width as usize, height as usize); + let mut out = vec![0.0_f32; w * h]; + let half_size = PARSER_SIZE as f32 / 2.0; + let last = (PARSER_SIZE - 1) as f32; + + for face in faces { + let plane = face.region_plane(region); + let [min_x, min_y, max_x, max_y] = face.bounds(); + let x0 = (min_x.floor() as i64 - origin.0 as i64).clamp(0, w as i64) as usize; + let x1 = (max_x.ceil() as i64 + 1 - origin.0 as i64).clamp(0, w as i64) as usize; + let y0 = (min_y.floor() as i64 - origin.1 as i64).clamp(0, h as i64) as usize; + let y1 = (max_y.ceil() as i64 + 1 - origin.1 as i64).clamp(0, h as i64) as usize; + if x0 >= x1 || y0 >= y1 { + continue; + } + let y_dir = face.y_dir(); + + out[y0 * w..y1 * w] + .par_chunks_exact_mut(w) + .enumerate() + .for_each(|(row_idx, row)| { + let dy = (y0 + row_idx) as f32 + origin.1 as f32 + 0.5 - face.center[1]; + for (x, value) in row.iter_mut().enumerate().take(x1).skip(x0) { + let dx = x as f32 + origin.0 as f32 + 0.5 - face.center[0]; + let a = (dx * face.x_dir[0] + dy * face.x_dir[1]) / face.half; + let b = (dx * y_dir[0] + dy * y_dir[1]) / face.half; + let u = (a + 1.0) * half_size - 0.5; + let v = (b + 1.0) * half_size - 0.5; + if u < -0.5 || v < -0.5 || u > last + 0.5 || v > last + 0.5 { + continue; + } + let (u, v) = (u.clamp(0.0, last), v.clamp(0.0, last)); + let (u0, v0) = (u as usize, v as usize); + let (u1, v1) = ((u0 + 1).min(PARSER_SIZE - 1), (v0 + 1).min(PARSER_SIZE - 1)); + let (fu, fv) = (u - u0 as f32, v - v0 as f32); + let top = plane[v0 * PARSER_SIZE + u0] * (1.0 - fu) + + plane[v0 * PARSER_SIZE + u1] * fu; + let bottom = plane[v1 * PARSER_SIZE + u0] * (1.0 - fu) + + plane[v1 * PARSER_SIZE + u1] * fu; + *value = value.max(top * (1.0 - fv) + bottom * fv); + } + }); + } + out +} + +#[inline] +fn smoothstep(edge0: f32, edge1: f32, x: f32) -> f32 { + let t = ((x - edge0) / (edge1 - edge0)).clamp(0.0, 1.0); + t * t * (3.0 - 2.0 * t) +} + +/// Narrows a rendered `Teeth` or `EyeWhites` plane down to the bright, neutral pixels of +/// `rgb`, the sRGB pixels the plane is aligned with. Other regions are left untouched. +pub fn keep_bright_neutral(values: &mut [f32], rgb: &[u8], region: FaceRegion) { + let Some((luma_limits, saturation_limits)) = region.bright_neutral_limits() else { + return; + }; + values + .par_iter_mut() + .zip(rgb.par_chunks_exact(3)) + .for_each(|(value, px)| { + if *value <= 0.0 { + return; + } + let [r, g, b] = [px[0], px[1], px[2]].map(|c| c as f32 / 255.0); + let luma = 0.299 * r + 0.587 * g + 0.114 * b; + let max = r.max(g).max(b); + let saturation = if max > 1e-4 { + (max - r.min(g).min(b)) / max + } else { + 0.0 + }; + *value *= smoothstep(luma_limits[0], luma_limits[1], luma) + * (1.0 - smoothstep(saturation_limits[0], saturation_limits[1], saturation)); + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn detection(landmarks: [[f32; 2]; 5]) -> Detection { + Detection { + score: 0.9, + bbox: [290.0, 250.0, 480.0, 680.0], + landmarks, + } + } + + #[test] + fn alignment_matches_training_framing() { + // Landmarks of an aligned 1024px training portrait must map back onto itself. + let face = align_face(&detection([ + [406.0, 493.0], + [634.0, 492.0], + [502.0, 644.0], + [415.0, 741.0], + [621.0, 741.0], + ])); + assert!((face.half - 512.0).abs() < 20.0, "half {}", face.half); + assert!((face.center[0] - 512.0).abs() < 20.0); + assert!((face.center[1] - 512.0).abs() < 20.0); + assert!(face.x_dir[0] > 0.99); + } + + #[test] + fn alignment_follows_sideways_faces() { + // The same face rotated 90 degrees clockwise. + let face = align_face(&detection([ + [531.0, 406.0], + [532.0, 634.0], + [380.0, 502.0], + [283.0, 415.0], + [283.0, 621.0], + ])); + assert!(face.x_dir[1] > 0.99, "x_dir {:?}", face.x_dir); + let top_left = face.crop_to_image(0.0, 0.0); + assert!(top_left[0] > 900.0 && top_left[1] < 120.0, "{top_left:?}"); + } + + /// Runs the real models. Point `FACE_TEST_DIR` at a folder holding `yunet.onnx`, + /// `parsing.onnx` and `face.jpg`. + #[test] + #[ignore] + fn finds_a_real_face_in_any_orientation() { + let dir = std::path::PathBuf::from(std::env::var("FACE_TEST_DIR").unwrap()); + let _ = ort::init().with_name("face-test").commit(); + let models = FaceModels { + detector: Mutex::new( + Session::builder() + .unwrap() + .commit_from_file(dir.join("yunet.onnx")) + .unwrap(), + ), + parser: Mutex::new(load_face_parser_session(&dir.join("parsing.onnx")).unwrap()), + }; + let upright = image::open(dir.join("face.jpg")).unwrap().to_rgb8(); + + let mut skin_pixels = Vec::new(); + for image in [upright.clone(), imageops::rotate90(&upright)] { + let faces = analyze_faces(&image, &models).unwrap(); + assert_eq!(faces.len(), 1); + let (w, h) = image.dimensions(); + let skin = render_region(&faces, FaceRegion::Skin, (0, 0), w, h); + skin_pixels.push(skin.iter().filter(|&&p| p > 0.5).count()); + } + assert!(skin_pixels[0] as f32 > 0.03 * (upright.width() * upright.height()) as f32); + let difference = skin_pixels[0].abs_diff(skin_pixels[1]); + assert!(difference * 50 < skin_pixels[0], "{skin_pixels:?}"); + } + + #[test] + fn blemish_zone_is_forehead_and_cheeks_only() { + let face = align_face(&detection([ + [406.0, 493.0], + [634.0, 492.0], + [502.0, 644.0], + [415.0, 741.0], + [621.0, 741.0], + ])); + // Positions on the aligned 1024px portrait, which is twice the crop size. + let weight = |x: f32, y: f32| face.blemish_zone_weight(x / 2.0, y / 2.0); + + assert!(weight(512.0, 330.0) > 0.99, "forehead"); + assert!(weight(360.0, 620.0) > 0.99, "right cheek"); + assert!(weight(680.0, 620.0) > 0.99, "left cheek"); + assert!(weight(505.0, 640.0) < 0.01, "nose"); + assert!(weight(505.0, 700.0) < 0.01, "under the nose"); + assert!(weight(430.0, 700.0) < 0.01, "right nasolabial fold"); + assert!(weight(600.0, 700.0) < 0.01, "left nasolabial fold"); + assert!(weight(445.0, 650.0) < 0.01, "wing of the nose"); + assert!(weight(512.0, 860.0) < 0.01, "chin"); + assert!(weight(400.0, 500.0) < 0.01, "eye"); + } + + #[test] + fn renders_region_back_into_image_space() { + let mut face = align_face(&detection([ + [406.0, 493.0], + [634.0, 492.0], + [502.0, 644.0], + [415.0, 741.0], + [621.0, 741.0], + ])); + let area = PARSER_SIZE * PARSER_SIZE; + face.probs = vec![0u8; PARSER_CLASSES * area]; + for v in 200..300 { + for u in 100..200 { + face.probs[SKIN * area + v * PARSER_SIZE + u] = 255; + } + } + + let [inside_x, inside_y] = face.crop_to_image(150.0, 250.0); + let [outside_x, outside_y] = face.crop_to_image(400.0, 250.0); + let origin = (40_u32, 60_u32); + let rendered = render_region(&[face], FaceRegion::Skin, origin, 1000, 900); + let at = |x: f32, y: f32| { + rendered[(y as usize - origin.1 as usize) * 1000 + (x as usize - origin.0 as usize)] + }; + assert!(at(inside_x, inside_y) > 0.99); + assert!(at(outside_x, outside_y) < 0.01); + } +} diff --git a/src-tauri/src/face_refine.rs b/src-tauri/src/face_refine.rs new file mode 100644 index 0000000000..e45880dff1 --- /dev/null +++ b/src-tauri/src/face_refine.rs @@ -0,0 +1,458 @@ +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; +use std::io::Cursor; +use std::sync::Arc; + +use base64::{Engine as _, engine::general_purpose}; +use image::{GrayImage, RgbImage, imageops}; +use rayon::prelude::*; +use serde_json::Value; + +use crate::ai_processing::get_or_init_face_models; +use crate::app_state::AppState; +use crate::face_parsing::{ + FaceRegion, ParsedFace, analyze_faces, keep_bright_neutral, render_region, +}; +use crate::inpainting::{encode_patch_result, prepare_source_image}; +use crate::skin_retouch::{ + FaceHints, FeatureEnhance, SkinRetouchParams, enhance_eyes_and_teeth, retouch_skin_scaled, +}; + +/// Everything needed to retouch one face again and again without going back to the +/// full image: its pixels and what the parsing model found in them. +pub struct FaceCrop { + center: [f32; 2], + width: f32, + x: u32, + y: u32, + w: u32, + h: u32, + rgb: Vec, + /// The skin that gets retouched. + selection: Vec, + skin: Vec, + features: Vec, + blemish_zone: Vec, + eyes: Vec, + eye_whites: Vec, + teeth: Vec, + thumbnail: String, +} + +/// The face crops of the loaded image, valid while `key` matches. +pub struct FaceCropCache { + pub key: String, + is_raw: bool, + crops: Vec, +} + +pub struct FaceRefineSettings { + pub skin: SkinRetouchParams, + pub eyes: f32, + pub teeth: f32, +} + +fn to_u8(plane: &[f32]) -> Vec { + plane + .par_iter() + .map(|&v| (v.clamp(0.0, 1.0) * 255.0).round() as u8) + .collect() +} + +fn build_face_crop(face: &ParsedFace, source: &RgbImage) -> Option { + let (img_w, img_h) = source.dimensions(); + let face_slice = std::slice::from_ref(face); + + // First pass over everything the model looked at, to find where the skin actually is. + let [min_x, min_y, max_x, max_y] = face.bounds(); + let bx = min_x.clamp(0.0, img_w as f32 - 1.0) as u32; + let by = min_y.clamp(0.0, img_h as f32 - 1.0) as u32; + let bw = (max_x.clamp(0.0, img_w as f32) as u32).saturating_sub(bx); + let bh = (max_y.clamp(0.0, img_h as f32) as u32).saturating_sub(by); + if bw == 0 || bh == 0 { + return None; + } + let coarse = render_region(face_slice, FaceRegion::Skin, (bx, by), bw, bh); + let (mut x0, mut y0, mut x1, mut y1) = (u32::MAX, u32::MAX, 0, 0); + for (i, &p) in coarse.iter().enumerate() { + if p > 0.3 { + let (px, py) = (i as u32 % bw, i as u32 / bw); + x0 = x0.min(px); + y0 = y0.min(py); + x1 = x1.max(px); + y1 = y1.max(py); + } + } + if x0 > x1 || y0 > y1 { + return None; + } + drop(coarse); + + let pad = (face.width * 0.1 + 16.0).ceil() as u32; + let x = (bx + x0).saturating_sub(pad); + let y = (by + y0).saturating_sub(pad); + let w = (bx + x1 + pad + 1).min(img_w) - x; + let h = (by + y1 + pad + 1).min(img_h) - y; + let origin = (x, y); + + let rgb = imageops::crop_imm(source, x, y, w, h).to_image().into_raw(); + let skin = render_region(face_slice, FaceRegion::Skin, origin, w, h); + let selection: Vec = skin + .par_iter() + .map(|&p| { + let t = ((p - 0.3) / 0.4).clamp(0.0, 1.0); + (t * t * (3.0 - 2.0 * t) * 255.0) as u8 + }) + .collect(); + let eyes = render_region(face_slice, FaceRegion::Eyes, origin, w, h); + let mut eye_whites = eyes.clone(); + keep_bright_neutral(&mut eye_whites, &rgb, FaceRegion::EyeWhites); + let mut teeth = render_region(face_slice, FaceRegion::Teeth, origin, w, h); + keep_bright_neutral(&mut teeth, &rgb, FaceRegion::Teeth); + + let mut jpeg = Cursor::new(Vec::new()); + face.thumbnail(source, 96) + .write_with_encoder(image::codecs::jpeg::JpegEncoder::new_with_quality( + &mut jpeg, 85, + )) + .ok()?; + + Some(FaceCrop { + center: face.center(), + width: face.width, + x, + y, + w, + h, + rgb, + selection, + skin: to_u8(&skin), + features: to_u8(&render_region( + face_slice, + FaceRegion::Features, + origin, + w, + h, + )), + blemish_zone: to_u8(&render_region( + face_slice, + FaceRegion::BlemishZone, + origin, + w, + h, + )), + eyes: to_u8(&eyes), + eye_whites: to_u8(&eye_whites), + teeth: to_u8(&teeth), + thumbnail: format!( + "data:image/jpeg;base64,{}", + general_purpose::STANDARD.encode(jpeg.get_ref()) + ), + }) +} + +/// Retouches one cached face. Returns the new pixels of the crop and where they apply. +fn refine_face_crop(crop: &FaceCrop, settings: &FaceRefineSettings) -> (Vec, Vec) { + let (w, h) = (crop.w as usize, crop.h as usize); + let scale = crop.width * 1.2; + + let mut color = retouch_skin_scaled( + &crop.rgb, + &crop.selection, + w, + h, + scale, + &settings.skin, + Some(FaceHints { + skin: &crop.skin, + features: &crop.features, + blemish_zone: Some(&crop.blemish_zone), + }), + ); + let mut mask = crop.selection.clone(); + + if settings.eyes > 0.0 || settings.teeth > 0.0 { + enhance_eyes_and_teeth( + &mut color, + w, + h, + scale, + &FeatureEnhance { + eyes: &crop.eyes, + eye_whites: &crop.eye_whites, + teeth: &crop.teeth, + eye_amount: settings.eyes, + teeth_amount: settings.teeth, + }, + ); + // Eyes and teeth are outside the skin selection, so the patch has to cover them too. + mask.par_iter_mut().enumerate().for_each(|(i, m)| { + if settings.eyes > 0.0 { + *m = (*m).max(crop.eyes[i]); + } + if settings.teeth > 0.0 { + *m = (*m).max(crop.teeth[i]); + } + }); + } + (color, mask) +} + +/// The image the faces are read from must not contain Facial Refine's own patches. +fn without_face_patches(adjustments: &Value) -> Value { + let mut stripped = adjustments.clone(); + if let Some(patches) = stripped.get_mut("aiPatches").and_then(|v| v.as_array_mut()) { + patches.retain(|p| p.get("faceRefine").and_then(|v| v.as_bool()) != Some(true)); + } + stripped +} + +/// Changes whenever the image or any other patch painted on it changes. +fn source_key(path: &str, stripped_adjustments: &Value) -> String { + let mut hasher = DefaultHasher::new(); + if let Some(patches) = stripped_adjustments + .get("aiPatches") + .and_then(|v| v.as_array()) + { + for patch in patches { + let text = |pointer: &str| { + patch + .pointer(pointer) + .and_then(|v| v.as_str()) + .unwrap_or_default() + }; + text("/id").hash(&mut hasher); + patch + .get("visible") + .and_then(|v| v.as_bool()) + .hash(&mut hasher); + text("/patchData/color").len().hash(&mut hasher); + text("/patchData/mask").len().hash(&mut hasher); + } + } + format!("{}|{:x}", path, hasher.finish()) +} + +async fn face_crops( + current_adjustments: &Value, + state: &tauri::State<'_, AppState>, + app_handle: &tauri::AppHandle, +) -> Result, String> { + let path = state + .original_image + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .map(|img| img.path.clone()) + .ok_or("No original image loaded")?; + let stripped = without_face_patches(current_adjustments); + let key = source_key(&path, &stripped); + + let cached = state + .ai_state + .lock() + .unwrap() + .as_ref() + .and_then(|ai_state| ai_state.face_crops.clone()) + .filter(|cache| cache.key == key); + if let Some(cache) = cached { + return Ok(cache); + } + + let models = get_or_init_face_models(app_handle, &state.ai_state, &state.ai_init_lock) + .await + .map_err(|e| e.to_string())?; + + let (source, is_raw) = { + let (source_dynamic, is_raw) = prepare_source_image("", &stripped, state)?; + (source_dynamic.to_rgb8(), is_raw) + }; + let faces = analyze_faces(&source, &models).map_err(|e| e.to_string())?; + let mut crops: Vec = faces + .iter() + .filter_map(|face| build_face_crop(face, &source)) + .collect(); + crops.sort_by(|a, b| a.center[0].total_cmp(&b.center[0])); + + let cache = Arc::new(FaceCropCache { key, is_raw, crops }); + if let Some(ai_state) = state.ai_state.lock().unwrap().as_mut() { + ai_state.face_crops = Some(cache.clone()); + } + Ok(cache) +} + +#[derive(serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DetectedFace { + center_x: f32, + center_y: f32, + width: f32, + thumbnail: String, +} + +/// Lists the faces of the loaded image, in source image coordinates, for Facial Refine. +#[tauri::command] +pub async fn detect_faces_for_refine( + current_adjustments: Value, + state: tauri::State<'_, AppState>, + app_handle: tauri::AppHandle, +) -> Result, String> { + let cache = face_crops(¤t_adjustments, &state, &app_handle).await?; + Ok(cache + .crops + .iter() + .map(|crop| DetectedFace { + center_x: crop.center[0], + center_y: crop.center[1], + width: crop.width, + thumbnail: crop.thumbnail.clone(), + }) + .collect()) +} + +/// Retouches the face at `face_center` and returns it as patch data. +#[tauri::command] +pub async fn generate_face_refine_patch( + current_adjustments: Value, + face_center: (f32, f32), + parameters: Value, + state: tauri::State<'_, AppState>, + app_handle: tauri::AppHandle, +) -> Result { + let cache = face_crops(¤t_adjustments, &state, &app_handle).await?; + + let distance = + |crop: &FaceCrop| (crop.center[0] - face_center.0).hypot(crop.center[1] - face_center.1); + let crop = cache + .crops + .iter() + .filter(|crop| distance(crop) < crop.width) + .min_by(|a, b| distance(a).total_cmp(&distance(b))) + .ok_or("The selected face was not found in this image.")?; + + let percent = |key: &str, default: f32| { + parameters + .get(key) + .and_then(|v| v.as_f64()) + .map_or(default, |v| (v as f32 / 100.0).clamp(0.0, 1.0)) + }; + let settings = FaceRefineSettings { + skin: SkinRetouchParams { + smoothing: percent("intensity", 0.4), + texture: percent("texture", 0.75), + blemish: percent("blemish", 0.5), + shine: percent("shine", 0.4), + even_tone: percent("evenTone", 0.35), + skin_protection: 0.85, + }, + eyes: percent("eyes", 0.0), + teeth: percent("teeth", 0.0), + }; + + let started = std::time::Instant::now(); + let (color, mask) = refine_face_crop(crop, &settings); + log::info!( + "facial refine ({}x{} face crop) took {:.2?}", + crop.w, + crop.h, + started.elapsed() + ); + + encode_patch_result( + &RgbImage::from_raw(crop.w, crop.h, color).ok_or("Invalid face crop")?, + &GrayImage::from_raw(crop.w, crop.h, mask).ok_or("Invalid face crop")?, + crop.x, + crop.y, + crop.w, + crop.h, + cache.is_raw, + 100, + false, + ) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use ort::session::Session; + + use super::*; + use crate::face_parsing::{FaceModels, load_face_parser_session}; + + /// Runs the real models. Point `FACE_TEST_DIR` at a folder holding `yunet.onnx`, + /// `parsing.onnx` and `face.jpg`; the results are written next to them. + #[test] + #[ignore] + fn refines_a_real_portrait() { + let dir = std::path::PathBuf::from(std::env::var("FACE_TEST_DIR").unwrap()); + let _ = ort::init().with_name("face-test").commit(); + let models = FaceModels { + detector: Mutex::new( + Session::builder() + .unwrap() + .commit_from_file(dir.join("yunet.onnx")) + .unwrap(), + ), + parser: Mutex::new(load_face_parser_session(&dir.join("parsing.onnx")).unwrap()), + }; + let small = image::open(dir.join("face.jpg")).unwrap().to_rgb8(); + let settings = FaceRefineSettings { + skin: SkinRetouchParams { + smoothing: 0.4, + texture: 0.75, + blemish: 1.0, + shine: 0.4, + even_tone: 0.35, + skin_protection: 0.85, + }, + eyes: 0.6, + teeth: 0.6, + }; + + // The same portrait as shot, and blown up to the size of a close-up on a 24MP frame. + let large = imageops::resize(&small, 3072, 3072, imageops::FilterType::CatmullRom); + for (name, image) in [("small", small), ("large", large)] { + let faces = analyze_faces(&image, &models).unwrap(); + assert_eq!(faces.len(), 1, "{name}"); + let crop = build_face_crop(&faces[0], &image).unwrap(); + for (region, plane) in [ + ("selection", &crop.selection), + ("blemish zone", &crop.blemish_zone), + ("eyes", &crop.eyes), + ("eye whites", &crop.eye_whites), + ("teeth", &crop.teeth), + ] { + assert!(plane.iter().any(|&p| p > 127), "{name} {region}"); + } + + image::GrayImage::from_raw(crop.w, crop.h, crop.blemish_zone.clone()) + .unwrap() + .save(dir.join(format!("{name}_blemish_zone.png"))) + .unwrap(); + + let started = std::time::Instant::now(); + let (color, mask) = refine_face_crop(&crop, &settings); + println!( + "{name}: {}x{} crop refined in {:?}", + crop.w, + crop.h, + started.elapsed() + ); + assert_eq!(mask.len(), color.len() / 3); + + let mut result = image.clone(); + for (i, px) in color.chunks_exact(3).enumerate() { + let (cx, cy) = (crop.x + i as u32 % crop.w, crop.y + i as u32 / crop.w); + let alpha = mask[i] as f32 / 255.0; + let dst = result.get_pixel_mut(cx, cy); + for c in 0..3 { + dst[c] = (dst[c] as f32 * (1.0 - alpha) + px[c] as f32 * alpha) as u8; + } + } + result + .save(dir.join(format!("{name}_refined.png"))) + .unwrap(); + } + } +} diff --git a/src-tauri/src/file_management.rs b/src-tauri/src/file_management.rs index d9f4c1f181..164b7927dc 100644 --- a/src-tauri/src/file_management.rs +++ b/src-tauri/src/file_management.rs @@ -10,7 +10,7 @@ use std::path::{Path, PathBuf}; #[cfg(any(target_os = "windows", target_os = "macos", target_os = "linux"))] use std::process::Command; use std::sync::Arc; -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicU64, Ordering}; use std::thread; use anyhow::Result; @@ -806,6 +806,115 @@ pub async fn list_images_in_dir( .map_err(|e| format!("Task failed: {}", e))? } +// Cloud placeholders are skipped so a background count never forces downloads. +fn collect_flagged_paths( + root_path: &Path, + settings: &AppSettings, + is_cancelled: &(dyn Fn() -> bool + Sync), +) -> HashMap { + let enable_xmp_sync = settings.enable_xmp_sync.unwrap_or(false); + let mut images = Vec::new(); + let mut sidecars_by_path: HashMap>> = HashMap::new(); + + for entry in WalkDir::new(root_path).into_iter().filter_map(Result::ok) { + if is_cancelled() { + return HashMap::new(); + } + + let file_type = entry.file_type(); + if !(file_type.is_file() || (file_type.is_symlink() && entry.path().is_file())) { + continue; + } + + let file_name = entry.file_name().to_string_lossy(); + if let Some((source_filename, copy_id)) = parse_sidecar_name(&file_name) { + sidecars_by_path + .entry(entry.path().with_file_name(source_filename)) + .or_default() + .push(copy_id.map(str::to_string)); + } else if is_supported_image_file(file_name.as_ref()) { + images.push(entry.into_path()); + } + } + + let tasks: Vec<_> = images + .into_iter() + .filter_map(|image_path| { + let sidecars = sidecars_by_path.remove(&image_path); + // Without a sidecar only an XMP file can hold a flag. + (sidecars.is_some() || enable_xmp_sync).then_some((image_path, sidecars)) + }) + .collect(); + + tasks + .into_par_iter() + .flat_map_iter(|(image_path, sidecars)| { + let mut flagged = Vec::new(); + if is_cancelled() + || (enable_xmp_sync + && resolve_xmp_path(&image_path).is_some_and(|p| is_cloud_placeholder(&p))) + { + return flagged; + } + + let path_str = image_path.to_string_lossy().into_owned(); + let file_name = image_path + .file_name() + .unwrap_or_default() + .to_string_lossy() + .into_owned(); + let sidecar_exists = sidecars.is_some(); + + for copy_id in sidecars.unwrap_or_else(|| vec![None]) { + let (virtual_path, sidecar_filename) = match copy_id { + Some(id) => ( + format!("{}?vc={}", path_str, id), + format!("{}.{}.rrdata", file_name, id), + ), + None => (path_str.clone(), format!("{}.rrdata", file_name)), + }; + let sidecar_path = image_path.with_file_name(sidecar_filename); + if sidecar_exists && is_cloud_placeholder(&sidecar_path) { + continue; + } + + let metadata = resolve_image_metadata( + &image_path, + &sidecar_path, + sidecar_exists, + enable_xmp_sync, + settings, + ); + if let Some(flag) = metadata.flag { + flagged.push((virtual_path, flag)); + } + } + + flagged + }) + .collect() +} + +static FLAG_SCAN_GENERATION: AtomicU64 = AtomicU64::new(0); + +// Only the newest scan matters: starting one makes any scan still running stop early. +#[tauri::command] +pub async fn get_flagged_paths( + path: String, + app_handle: AppHandle, +) -> Result, String> { + let generation = FLAG_SCAN_GENERATION.fetch_add(1, Ordering::SeqCst) + 1; + + tauri::async_runtime::spawn_blocking(move || { + let settings = load_settings(app_handle).unwrap_or_default(); + collect_flagged_paths(Path::new(&path), &settings, &|| { + FLAG_SCAN_GENERATION.load(Ordering::SeqCst) != generation + }) + }) + .await + .map_err(|e| e.to_string()) +} + #[tauri::command] pub async fn list_images_recursive( path: String, @@ -4717,3 +4826,87 @@ pub fn sync_metadata_to_xmp(source_path: &Path, metadata: &ImageMetadata, create let _ = fs::write(&xmp_file, content); } } + +#[cfg(test)] +mod flagged_paths_tests { + use super::*; + + fn write_sidecar(path: &Path, flag: Option) { + let metadata = ImageMetadata { + flag, + ..Default::default() + }; + fs::write(path, serde_json::to_string(&metadata).unwrap()).unwrap(); + } + + fn create_library(root: &Path) { + let sub = root.join("session").join("day2"); + fs::create_dir_all(&sub).unwrap(); + + for image in [ + root.join("picked.jpg"), + root.join("no_sidecar.jpg"), + root.join("unflagged.jpg"), + sub.join("rejected.jpg"), + ] { + fs::write(image, b"").unwrap(); + } + + write_sidecar(&root.join("picked.jpg.rrdata"), Some(ImageFlag::Pick)); + write_sidecar( + &root.join("picked.jpg.abc123.rrdata"), + Some(ImageFlag::Reject), + ); + write_sidecar(&root.join("unflagged.jpg.rrdata"), None); + write_sidecar(&sub.join("rejected.jpg.rrdata"), Some(ImageFlag::Reject)); + } + + #[test] + fn finds_flags_in_subfolders_and_virtual_copies() { + let dir = tempfile::tempdir().unwrap(); + create_library(dir.path()); + + let flagged = collect_flagged_paths(dir.path(), &AppSettings::default(), &|| false); + + let picked = dir.path().join("picked.jpg").to_string_lossy().into_owned(); + let rejected = dir + .path() + .join("session") + .join("day2") + .join("rejected.jpg") + .to_string_lossy() + .into_owned(); + + assert_eq!(flagged.len(), 3); + assert_eq!(flagged.get(&picked), Some(&ImageFlag::Pick)); + assert_eq!(flagged.get(&rejected), Some(&ImageFlag::Reject)); + assert_eq!( + flagged.get(&format!("{}?vc=abc123", picked)), + Some(&ImageFlag::Reject) + ); + } + + #[test] + fn subfolder_scan_only_sees_its_own_images() { + let dir = tempfile::tempdir().unwrap(); + create_library(dir.path()); + + let flagged = collect_flagged_paths( + &dir.path().join("session"), + &AppSettings::default(), + &|| false, + ); + + assert_eq!(flagged.len(), 1); + } + + #[test] + fn cancelled_scan_returns_nothing() { + let dir = tempfile::tempdir().unwrap(); + create_library(dir.path()); + + let flagged = collect_flagged_paths(dir.path(), &AppSettings::default(), &|| true); + + assert!(flagged.is_empty()); + } +} diff --git a/src-tauri/src/inpainting.rs b/src-tauri/src/inpainting.rs index 9aff0e5e81..b9559cffab 100644 --- a/src-tauri/src/inpainting.rs +++ b/src-tauri/src/inpainting.rs @@ -13,8 +13,9 @@ use crate::image_loader::prepare_base_image_cow; use crate::image_processing::apply_linear_to_srgb; use crate::mask_generation::{AiPatchDefinition, MaskDefinition, generate_mask_bitmap}; use crate::resolve_warped_image_for_masks; +use crate::skin_retouch::SkinRetouchParams; -fn prepare_source_image( +pub(crate) fn prepare_source_image( patch_id: &str, current_adjustments: &Value, state: &tauri::State<'_, AppState>, @@ -113,7 +114,7 @@ fn calculate_mask_bounds( } #[allow(clippy::too_many_arguments)] -fn encode_patch_result( +pub(crate) fn encode_patch_result( color_image: &RgbImage, mask_image: &image::GrayImage, offset_x: u32, @@ -1144,7 +1145,14 @@ pub async fn generate_retouch_patch( let mut max_x = 0.0_f32; let mut max_y = 0.0_f32; let mut max_radius = 0.0_f32; - let mut intensity = 50.0_f32; + let mut retouch_params = SkinRetouchParams { + smoothing: 0.4, + texture: 0.75, + blemish: 0.5, + shine: 0.4, + even_tone: 0.35, + skin_protection: 0.85, + }; let sub_masks_val = serde_json::to_value(&patch_definition.sub_masks).unwrap_or(Value::Null); let mut mask_canvas = vec![0u8; (img_w * img_h) as usize]; @@ -1154,10 +1162,20 @@ pub async fn generate_retouch_patch( if sm.get("type").and_then(|v| v.as_str()) == Some("retouch") && let Some(params) = sm.get("parameters") { - intensity = params - .get("intensity") - .and_then(|v| v.as_f64()) - .unwrap_or(50.0) as f32; + let percent = |key: &str, default: f32| { + params + .get(key) + .and_then(|v| v.as_f64()) + .map_or(default, |v| (v as f32 / 100.0).clamp(0.0, 1.0)) + }; + retouch_params = SkinRetouchParams { + smoothing: percent("intensity", retouch_params.smoothing), + texture: percent("texture", retouch_params.texture), + blemish: percent("blemish", retouch_params.blemish), + shine: percent("shine", retouch_params.shine), + even_tone: percent("evenTone", retouch_params.even_tone), + skin_protection: percent("skinProtection", retouch_params.skin_protection), + }; if let Some(lines) = params.get("lines").and_then(|v| v.as_array()) { for line in lines { @@ -1239,7 +1257,9 @@ pub async fn generate_retouch_patch( return Err("No brush strokes found for Retouch.".to_string()); } - let pad = (max_radius * 3.0 + 30.0).ceil(); + let painted_area: f32 = mask_canvas.par_iter().map(|&v| v as f32 / 255.0).sum(); + let retouch_scale = painted_area.sqrt().max(max_radius * 6.0); + let pad = (max_radius + retouch_scale / 12.0 + 16.0).ceil(); let min_x_u32 = (min_x - pad).clamp(0.0, img_w as f32 - 1.0) as u32; let min_y_u32 = (min_y - pad).clamp(0.0, img_h as f32 - 1.0) as u32; let max_x_u32 = (max_x + pad).clamp(0.0, img_w as f32 - 1.0) as u32; @@ -1259,112 +1279,15 @@ pub async fn generate_retouch_patch( } } - let norm_intensity = (intensity / 100.0).clamp(0.05, 1.0); - let sig_s = (1.0 + norm_intensity * 6.0).clamp(1.2, 8.0); - let sig_c = (8.0 + norm_intensity * 20.0).clamp(8.0, 28.0); - - let sig_s_sq = sig_s * sig_s; - let sig_c_sq = sig_c * sig_c; - - let r = (sig_s * 2.0).ceil() as i32; - let window_size = (2 * r + 1) as usize; - - let mut spatial_weights = vec![0.0_f32; window_size * window_size]; - for dy in -r..=r { - for dx in -r..=r { - let dist_s = (dx * dx + dy * dy) as f32; - let w = (-dist_s / (2.0 * sig_s_sq)).exp(); - let idx = ((dy + r) as usize) * window_size + ((dx + r) as usize); - spatial_weights[idx] = w; - } - } - - let max_color_dist_sq = 3 * 255 * 255; - let color_weights: Vec = (0..=max_color_dist_sq) - .map(|i| (-(i as f32) / (2.0 * sig_c_sq)).exp()) - .collect(); - - let mut color_pixels = vec![0u8; (crop_w * crop_h * 3) as usize]; - let crop_w_i32 = crop_w as i32; - let crop_h_i32 = crop_h as i32; - let crop_w_usize = crop_w as usize; - - color_pixels - .par_chunks_mut(crop_w_usize * 3) - .enumerate() - .for_each(|(y, row_out)| { - let y_i32 = y as i32; - - for x in 0..crop_w_usize { - let m_val = mask_pixels[y * crop_w_usize + x]; - let px_idx = (y * crop_w_usize + x) * 3; - - if m_val == 0 { - row_out[x * 3] = orig_raw[px_idx]; - row_out[x * 3 + 1] = orig_raw[px_idx + 1]; - row_out[x * 3 + 2] = orig_raw[px_idx + 2]; - continue; - } - - let m = m_val as f32 / 255.0; - let smooth_m = m * m * (3.0 - 2.0 * m); - let blend_factor = smooth_m * (0.3 + norm_intensity * 0.7); - - let ctr_r = orig_raw[px_idx] as i32; - let ctr_g = orig_raw[px_idx + 1] as i32; - let ctr_b = orig_raw[px_idx + 2] as i32; - - let mut sum_w = 0.0_f32; - let mut sum_r = 0.0_f32; - let mut sum_g = 0.0_f32; - let mut sum_b = 0.0_f32; - - let x_i32 = x as i32; - let y_min = (y_i32 - r).max(0); - let y_max = (y_i32 + r).min(crop_h_i32 - 1); - let x_min = (x_i32 - r).max(0); - let x_max = (x_i32 + r).min(crop_w_i32 - 1); - - for ny in y_min..=y_max { - let dy = ny - y_i32; - let wy_offset = (dy + r) as usize * window_size; - let row_offset = (ny as usize) * crop_w_usize; - - for nx in x_min..=x_max { - let dx = nx - x_i32; - let spatial_w = spatial_weights[wy_offset + (dx + r) as usize]; - - let n_idx = (row_offset + nx as usize) * 3; - let pr = orig_raw[n_idx] as i32; - let pg = orig_raw[n_idx + 1] as i32; - let pb = orig_raw[n_idx + 2] as i32; - - let dr = pr - ctr_r; - let dg = pg - ctr_g; - let db = pb - ctr_b; - - let color_dist_sq = (dr * dr + dg * dg + db * db) as usize; - let w = spatial_w * color_weights[color_dist_sq]; - - sum_w += w; - sum_r += (pr as f32) * w; - sum_g += (pg as f32) * w; - sum_b += (pb as f32) * w; - } - } - - let base_r = sum_r / sum_w; - let base_g = sum_g / sum_w; - let base_b = sum_b / sum_w; - - row_out[x * 3] = ((ctr_r as f32) * (1.0 - blend_factor) + base_r * blend_factor) - .clamp(0.0, 255.0) as u8; - row_out[x * 3 + 1] = ((ctr_g as f32) * (1.0 - blend_factor) + base_g * blend_factor) - .clamp(0.0, 255.0) as u8; - row_out[x * 3 + 2] = ((ctr_b as f32) * (1.0 - blend_factor) + base_b * blend_factor) - .clamp(0.0, 255.0) as u8; - } - }); + let color_pixels = crate::skin_retouch::retouch_skin_scaled( + orig_raw, + &mask_pixels, + crop_w as usize, + crop_h as usize, + retouch_scale, + &retouch_params, + None, + ); let color_image = RgbImage::from_raw(crop_w, crop_h, color_pixels).unwrap(); let mask_image = image::GrayImage::from_raw(crop_w, crop_h, mask_pixels).unwrap(); diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 4285ad8569..f729205588 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -40,6 +40,8 @@ mod edit_in; mod effect_maps; mod exif_processing; mod export_processing; +mod face_parsing; +mod face_refine; mod face_masking; mod file_management; mod focus_stacking; @@ -67,6 +69,7 @@ mod preset_converter; mod preview_jpeg; mod raw_processing; mod relight; +mod skin_retouch; mod storage; mod tagging; mod tagging_utils; @@ -2853,6 +2856,8 @@ pub fn run() { inpainting::generate_manual_cleanup_patch, inpainting::generate_liquify_patch, inpainting::generate_retouch_patch, + face_refine::detect_faces_for_refine, + face_refine::generate_face_refine_patch, focus_stacking::stitch_focus_stack, focus_stacking::save_focus_stack, image_loader::load_image, @@ -2874,6 +2879,7 @@ pub fn run() { file_management::read_exif_for_paths, file_management::list_images_in_dir, file_management::list_images_recursive, + file_management::get_flagged_paths, file_management::get_folder_tree, file_management::get_folder_children, file_management::get_pinned_folder_trees, diff --git a/src-tauri/src/skin_retouch.rs b/src-tauri/src/skin_retouch.rs new file mode 100644 index 0000000000..819674f2a8 --- /dev/null +++ b/src-tauri/src/skin_retouch.rs @@ -0,0 +1,1169 @@ +use rayon::prelude::*; + +/// All values are normalized to `0.0..=1.0`. +#[derive(Debug, Clone, Copy)] +pub struct SkinRetouchParams { + /// Evens out mid-frequency tonal unevenness (the "dodge & burn" part). + pub smoothing: f32, + /// How much of the fine skin texture (pores) is kept. + pub texture: f32, + /// Automatic removal of small spots (pimples, moles, dark dots). + pub blemish: f32, + /// Reduction of specular shine and sweat sparkles. + pub shine: f32, + /// Evens out color blotches and redness. + pub even_tone: f32, + /// How strongly non-skin pixels (eyes, lips, brows, hair) are left untouched. + pub skin_protection: f32, +} + +/// What a face parsing model knows about the pixels being retouched, as `0..=255` +/// probabilities aligned with the image buffer. +#[derive(Clone, Copy)] +pub struct FaceHints<'a> { + /// Skin, including the parts washed out by shine that color alone cannot identify. + pub skin: &'a [u8], + /// Eyes, brows, lips, mouth and hair, which must never be retouched as skin. + pub features: &'a [u8], + /// Where spots may be removed. Everywhere on the skin when absent. + pub blemish_zone: Option<&'a [u8]>, +} + +#[inline] +fn smoothstep(edge0: f32, edge1: f32, x: f32) -> f32 { + let t = ((x - edge0) / (edge1 - edge0)).clamp(0.0, 1.0); + t * t * (3.0 - 2.0 * t) +} + +fn box_1d(src: &[f32], dst: &mut [f32], r: usize) { + let n = src.len(); + let mut sum = 0.0_f32; + let mut lo = 0; + let mut hi = 0; + for (x, out) in dst.iter_mut().enumerate() { + let want_hi = (x + r + 1).min(n); + while hi < want_hi { + sum += src[hi]; + hi += 1; + } + let want_lo = x.saturating_sub(r); + while lo < want_lo { + sum -= src[lo]; + lo += 1; + } + *out = sum / (hi - lo) as f32; + } +} + +fn box_h(src: &[f32], dst: &mut [f32], w: usize, r: usize) { + dst.par_chunks_exact_mut(w) + .zip(src.par_chunks_exact(w)) + .for_each(|(dst_row, src_row)| box_1d(src_row, dst_row, r)); +} + +fn box_v(src: &[f32], dst: &mut [f32], w: usize, h: usize, r: usize) { + let mut sum = vec![0.0_f32; w]; + let mut lo = 0; + let mut hi = 0; + for y in 0..h { + let want_hi = (y + r + 1).min(h); + while hi < want_hi { + let row = &src[hi * w..(hi + 1) * w]; + sum.iter_mut().zip(row).for_each(|(s, &v)| *s += v); + hi += 1; + } + let want_lo = y.saturating_sub(r); + while lo < want_lo { + let row = &src[lo * w..(lo + 1) * w]; + sum.iter_mut().zip(row).for_each(|(s, &v)| *s -= v); + lo += 1; + } + let inv = 1.0 / (hi - lo) as f32; + dst[y * w..(y + 1) * w] + .iter_mut() + .zip(&sum) + .for_each(|(d, &s)| *d = s * inv); + } +} + +fn box_blur(src: &[f32], w: usize, h: usize, r: usize) -> Vec { + let mut tmp = vec![0.0_f32; src.len()]; + let mut out = vec![0.0_f32; src.len()]; + box_h(src, &mut tmp, w, r); + box_v(&tmp, &mut out, w, h, r); + out +} + +/// Gaussian approximation using three box blur passes. +fn gauss_blur(src: &[f32], w: usize, h: usize, sigma: f32) -> Vec { + if sigma < 0.6 { + return src.to_vec(); + } + let ideal_width = (4.0 * sigma * sigma + 1.0).sqrt(); + let r = (((ideal_width - 1.0) / 2.0).round() as usize).max(1); + + let mut a = src.to_vec(); + let mut b = vec![0.0_f32; src.len()]; + for _ in 0..3 { + box_h(&a, &mut b, w, r); + box_v(&b, &mut a, w, h, r); + } + a +} + +/// Normalized convolution: blurs `src` while ignoring pixels with low `weight`. +fn weighted_blur(src: &[f32], weight: &[f32], w: usize, h: usize, sigma: f32) -> Vec { + let weighted: Vec = src + .par_iter() + .zip(weight.par_iter()) + .map(|(&v, &k)| v * k) + .collect(); + let mut num = gauss_blur(&weighted, w, h, sigma); + let den = gauss_blur(weight, w, h, sigma); + num.par_iter_mut() + .zip(den.par_iter()) + .zip(src.par_iter()) + .for_each(|((n, &d), &s)| *n = if d > 1e-5 { *n / d } else { s }); + num +} + +struct GuidedFilter<'a> { + guide: &'a [f32], + mean_i: Vec, + var_i: Vec, + w: usize, + h: usize, + r: usize, +} + +impl<'a> GuidedFilter<'a> { + fn new(guide: &'a [f32], w: usize, h: usize, r: usize) -> Self { + let mean_i = box_blur(guide, w, h, r); + let sq: Vec = guide.par_iter().map(|&v| v * v).collect(); + let mut var_i = box_blur(&sq, w, h, r); + var_i + .par_iter_mut() + .zip(mean_i.par_iter()) + .for_each(|(v, &m)| *v = (*v - m * m).max(0.0)); + Self { + guide, + mean_i, + var_i, + w, + h, + r, + } + } + + fn filter(&self, p: &[f32], eps: f32) -> Vec { + let mean_p = box_blur(p, self.w, self.h, self.r); + let ip: Vec = self + .guide + .par_iter() + .zip(p.par_iter()) + .map(|(&i, &v)| i * v) + .collect(); + let mean_ip = box_blur(&ip, self.w, self.h, self.r); + + let mut a = vec![0.0_f32; p.len()]; + let mut b = vec![0.0_f32; p.len()]; + a.par_iter_mut() + .zip(b.par_iter_mut()) + .enumerate() + .for_each(|(i, (a_out, b_out))| { + let cov = mean_ip[i] - self.mean_i[i] * mean_p[i]; + let a_val = cov / (self.var_i[i] + eps); + *a_out = a_val; + *b_out = mean_p[i] - a_val * self.mean_i[i]; + }); + + let mean_a = box_blur(&a, self.w, self.h, self.r); + let mut out = box_blur(&b, self.w, self.h, self.r); + out.par_iter_mut() + .zip(mean_a.par_iter()) + .zip(self.guide.par_iter()) + .for_each(|((o, &ma), &g)| *o += ma * g); + out + } +} + +/// Gaussian skin-chroma model fitted to the pixels the user painted over. +struct SkinModel { + mean: [f32; 2], + inv_cov: [f32; 3], +} + +impl SkinModel { + #[inline] + fn dist_sq(&self, cb: f32, cr: f32) -> f32 { + let dx = cb - self.mean[0]; + let dy = cr - self.mean[1]; + self.inv_cov[0] * dx * dx + 2.0 * self.inv_cov[1] * dx * dy + self.inv_cov[2] * dy * dy + } + + #[inline] + fn probability(&self, cb: f32, cr: f32) -> f32 { + 1.0 - smoothstep(2.5, 4.5, self.dist_sq(cb, cr).sqrt()) + } + + /// Like `probability`, but also accepts the skin color washed out towards white by up + /// to `min_saturation`, which is how specular shine on skin looks. + #[inline] + fn probability_desaturated(&self, cb: f32, cr: f32, min_saturation: f32) -> f32 { + let [mx, my] = self.mean; + let [a, b, c] = self.inv_cov; + let along = a * cb * mx + b * (cb * my + cr * mx) + c * cr * my; + let len = a * mx * mx + 2.0 * b * mx * my + c * my * my; + let t = (along / len.max(1e-9)).clamp(min_saturation, 1.0); + let dx = cb - t * mx; + let dy = cr - t * my; + let dist_sq = a * dx * dx + 2.0 * b * dx * dy + c * dy * dy; + 1.0 - smoothstep(2.5, 4.5, dist_sq.sqrt()) + } + + fn from_moments(n: f64, sx: f64, sy: f64, sxx: f64, sxy: f64, syy: f64) -> Self { + const MIN_STD: f64 = 0.014; + let mx = sx / n; + let my = sy / n; + let cxx = sxx / n - mx * mx + MIN_STD * MIN_STD; + let cyy = syy / n - my * my + MIN_STD * MIN_STD; + let cxy = sxy / n - mx * my; + let det = (cxx * cyy - cxy * cxy).max(1e-12); + Self { + mean: [mx as f32, my as f32], + inv_cov: [(cyy / det) as f32, (-cxy / det) as f32, (cxx / det) as f32], + } + } + + fn fit(y: &[f32], cb: &[f32], cr: &[f32], mask: &[u8]) -> Self { + let fit_where = |pred: &dyn Fn(usize) -> bool| -> Option<(Self, usize)> { + let (mut n, mut sx, mut sy, mut sxx, mut sxy, mut syy) = + (0usize, 0.0_f64, 0.0_f64, 0.0_f64, 0.0_f64, 0.0_f64); + for i in 0..y.len() { + if mask[i] > 76 && pred(i) { + let (x, v) = (cb[i] as f64, cr[i] as f64); + n += 1; + sx += x; + sy += v; + sxx += x * x; + sxy += x * v; + syy += v * v; + } + } + (n > 0).then(|| (Self::from_moments(n as f64, sx, sy, sxx, sxy, syy), n)) + }; + + let Some((all_painted, painted)) = fit_where(&|i| y[i] > 0.08) else { + return Self::from_moments(1.0, -0.06, 0.08, 0.0036, -0.0048, 0.0064); + }; + + // Broad range that covers skin of any tone under reasonable white balance. + let looks_like_skin = |i: usize| { + y[i] > 0.08 + && y[i] < 0.98 + && cr[i] > 0.012 + && cr[i] < 0.25 + && cb[i] > -0.23 + && cb[i] < 0.02 + }; + let mut model = match fit_where(&looks_like_skin) { + Some((m, n)) if n >= (painted / 20).max(64) => m, + _ => all_painted, + }; + + for _ in 0..2 { + let refined = fit_where(&|i| y[i] > 0.08 && model.dist_sq(cb[i], cr[i]) < 6.25); + if let Some((refined, n)) = refined + && n >= 64 + { + model = refined; + } + } + model + } +} + +/// Finds small compact spots that differ from the surrounding skin and fills them +/// with the surrounding tone plus texture borrowed from a clean neighbour. +#[allow(clippy::too_many_arguments)] +fn heal_blemishes( + y: &mut [f32], + cb: &mut [f32], + cr: &mut [f32], + skin: &[f32], + mask: &[u8], + w: usize, + h: usize, + scale: f32, + amount: f32, + zone: Option<&[u8]>, +) { + let n = w * h; + let sigma_b = (scale / 80.0).clamp(2.0, 24.0); + let skin_weight: Vec = skin.par_iter().map(|&v| v + 0.02).collect(); + + let fine_y = gauss_blur(y, w, h, (sigma_b / 3.0).max(0.8)); + let bg_y = weighted_blur(y, &skin_weight, w, h, sigma_b * 2.0); + let fine_cr = gauss_blur(cr, w, h, (sigma_b / 3.0).max(0.8)); + let bg_cr = weighted_blur(cr, &skin_weight, w, h, sigma_b * 2.0); + + let (mut dev_y, mut dev_c, mut count) = (0.0_f64, 0.0_f64, 0usize); + for i in 0..n { + if mask[i] > 76 && skin[i] > 0.5 { + dev_y += (fine_y[i] - bg_y[i]).abs() as f64; + dev_c += (fine_cr[i] - bg_cr[i]).abs() as f64; + count += 1; + } + } + if count < 64 { + return; + } + let sig_y = ((1.2533 * dev_y / count as f64) as f32).max(0.006); + let sig_c = ((1.2533 * dev_c / count as f64) as f32).max(0.003); + + // Nostrils, mouth corners and ear canals are far darker than any blemish. They are + // anatomy, so they and their rims are kept out of the spot search. + let deep_shadow: Vec = (0..n) + .into_par_iter() + .map(|i| 1.0 - smoothstep(0.55, 0.75, fine_y[i] / bg_y[i].max(1e-3))) + .collect(); + let mut shadow_guard = gauss_blur(&deep_shadow, w, h, sigma_b * 1.2); + shadow_guard + .par_iter_mut() + .for_each(|v| *v = (*v * 8.0).min(1.0)); + + // Spots are darker and/or redder than their surroundings. + let score: Vec = (0..n) + .into_par_iter() + .map(|i| { + let spot = ((bg_y[i] - fine_y[i]) / sig_y).max((fine_cr[i] - bg_cr[i]) / sig_c); + spot * (1.0 - shadow_guard[i]) + }) + .collect(); + // How blob-like the score is around a pixel, from its Hessian: 1 for a round spot, 0 + // for a line, and nothing at all where the score is not a local bump. + let roundness_of = |smooth: &[f32], step: usize, px: usize, py: usize| -> Option { + let (x0, x1) = (px.saturating_sub(step), (px + step).min(w - 1)); + let (y0, y1) = (py.saturating_sub(step), (py + step).min(h - 1)); + let c = smooth[py * w + px]; + let lxx = smooth[py * w + x0] + smooth[py * w + x1] - 2.0 * c; + let lyy = smooth[y0 * w + px] + smooth[y1 * w + px] - 2.0 * c; + let lxy = + (smooth[y1 * w + x1] + smooth[y0 * w + x0] - smooth[y0 * w + x1] - smooth[y1 * w + x0]) + * 0.25; + let trace = lxx + lyy; + (trace < 0.0).then(|| 4.0 * (lxx * lyy - lxy * lxy) / (trace * trace)) + }; + + // A spot is round both up close and from further away. Folds, wrinkles, lash lines + // and lip edges can look like a string of spots up close, but from further away they + // are lines, and lines are never touched. + let near = gauss_blur(&score, w, h, sigma_b * 0.6); + let near_step = ((sigma_b * 0.6).round() as usize).max(1); + let far = gauss_blur(&score, w, h, sigma_b * 1.6); + let far_step = ((sigma_b * 1.6).round() as usize).max(1); + let threshold = 4.5 - 2.5 * amount; + + let mut spots = vec![0.0_f32; n]; + spots + .par_chunks_exact_mut(w) + .enumerate() + .for_each(|(py, row)| { + for (px, out) in row.iter_mut().enumerate() { + let i = py * w + px; + if mask[i] == 0 { + continue; + } + let in_zone = zone.map_or(1.0, |zone| zone[i] as f32 / 255.0); + let strength = smoothstep(threshold, threshold + 1.0, score[i]) * skin[i] * in_zone; + if strength <= 0.0 { + continue; + } + let Some(roundness) = roundness_of(&near, near_step, px, py) + .zip(roundness_of(&far, far_step, px, py)) + .map(|(near, far)| near.min(far)) + else { + continue; + }; + *out = strength * smoothstep(0.3, 0.6, roundness); + } + }); + + let mut cover = gauss_blur(&spots, w, h, sigma_b * 0.8); + cover + .par_iter_mut() + .zip(spots.par_iter()) + .zip(shadow_guard.par_iter()) + .for_each(|((c, &s), &guard)| *c = (*c * 4.0).max(s).min(1.0) * (1.0 - guard)); + + let clean_weight: Vec = cover + .par_iter() + .zip(skin_weight.par_iter()) + .map(|(&c, &s)| (1.0 - c) * s) + .collect(); + + let base_y = weighted_blur(y, &clean_weight, w, h, sigma_b * 1.5); + let base_cb = weighted_blur(cb, &clean_weight, w, h, sigma_b * 1.5); + let base_cr = weighted_blur(cr, &clean_weight, w, h, sigma_b * 1.5); + + let low = gauss_blur(y, w, h, (scale / 250.0).clamp(1.0, 6.0)); + let texture: Vec = y + .par_iter() + .zip(low.par_iter()) + .map(|(&v, &l)| v - l) + .collect(); + let reach = (sigma_b * 3.0).ceil() as isize; + let offsets = [(reach, 0), (-reach, 0), (0, reach), (0, -reach)]; + + let healed_y: Vec = (0..n) + .into_par_iter() + .map(|i| { + let c = cover[i]; + if c <= 0.003 { + return y[i]; + } + let (px, py) = ((i % w) as isize, (i / w) as isize); + let mut best_weight = 0.05_f32; + let mut borrowed = 0.0_f32; + for (dx, dy) in offsets { + let (qx, qy) = (px + dx, py + dy); + if qx < 0 || qy < 0 || qx >= w as isize || qy >= h as isize { + continue; + } + let q = qy as usize * w + qx as usize; + if clean_weight[q] > best_weight { + best_weight = clean_weight[q]; + borrowed = texture[q]; + } + } + y[i] + c * (base_y[i] + borrowed - y[i]) + }) + .collect(); + + y.copy_from_slice(&healed_y); + cb.par_iter_mut() + .zip(base_cb.par_iter()) + .zip(cover.par_iter()) + .for_each(|((v, &b), &c)| *v += c * (b - *v)); + cr.par_iter_mut() + .zip(base_cr.par_iter()) + .zip(cover.par_iter()) + .for_each(|((v, &b), &c)| *v += c * (b - *v)); +} + +fn percentile(hist: &[u32], total: u32, fraction: f32) -> f32 { + let target = (total as f32 * fraction) as u32; + let mut acc = 0; + for (bin, &count) in hist.iter().enumerate() { + acc += count; + if acc >= target { + return bin as f32 / (hist.len() - 1) as f32; + } + } + 1.0 +} + +/// Portrait skin retouch on an sRGB encoded RGB8 buffer. +/// +/// `mask` is the painted area and `scale` the approximate size in pixels of the region +/// being retouched (roughly the face width), which sets every filter radius. +pub fn retouch_skin( + rgb: &[u8], + mask: &[u8], + w: usize, + h: usize, + scale: f32, + params: &SkinRetouchParams, + hints: Option, +) -> Vec { + let n = w * h; + let scale = scale.clamp(80.0, 4000.0); + + let mut y = vec![0.0_f32; n]; + let mut cb = vec![0.0_f32; n]; + let mut cr = vec![0.0_f32; n]; + y.par_iter_mut() + .zip(cb.par_iter_mut()) + .zip(cr.par_iter_mut()) + .zip(rgb.par_chunks_exact(3)) + .for_each(|(((y_out, cb_out), cr_out), px)| { + let r = px[0] as f32 / 255.0; + let g = px[1] as f32 / 255.0; + let b = px[2] as f32 / 255.0; + let luma = 0.299 * r + 0.587 * g + 0.114 * b; + *y_out = luma; + *cb_out = (b - luma) * 0.564; + *cr_out = (r - luma) * 0.713; + }); + + let model = SkinModel::fit(&y, &cb, &cr, mask); + let mut luma_hist = [0u32; 256]; + let mut skin_count = 0u32; + for i in 0..n { + if mask[i] > 76 && model.probability(cb[i], cr[i]) > 0.5 { + luma_hist[(y[i].clamp(0.0, 1.0) * 255.0) as usize] += 1; + skin_count += 1; + } + } + let skin_luma = if skin_count > 0 { + percentile(&luma_hist, skin_count, 0.5) + } else { + 0.5 + }; + + let skin_raw: Vec = (0..n) + .into_par_iter() + .map(|i| { + // Only pixels clearly brighter than the skin may be desaturated shine. + let washout = smoothstep(skin_luma + 0.08, skin_luma + 0.30, y[i]); + let by_color = model.probability_desaturated(cb[i], cr[i], 1.0 - 0.7 * washout) + * smoothstep(0.03, 0.10, y[i]); + match hints { + Some(hints) => { + let skin = hints.skin[i] as f32 / 255.0; + let features = hints.features[i] as f32 / 255.0; + (by_color + (1.0 - by_color) * 0.6 * skin) * (1.0 - features) + } + None => by_color, + } + }) + .collect(); + let skin = gauss_blur(&skin_raw, w, h, (scale / 250.0).clamp(1.0, 4.0)); + let skin_weight: Vec = skin.par_iter().map(|&v| v + 0.02).collect(); + + if params.blemish > 0.001 { + heal_blemishes( + &mut y, + &mut cb, + &mut cr, + &skin, + mask, + w, + h, + scale, + params.blemish, + hints.and_then(|hints| hints.blemish_zone), + ); + } + + // Frequency separation: fine texture is kept apart while tone is evened out below it. + let low = gauss_blur(&y, w, h, (scale / 250.0).clamp(1.0, 6.0)); + let texture_gain = 0.25 + 0.75 * params.texture; + let mut high: Vec = y + .par_iter() + .zip(low.par_iter()) + .map(|(&v, &l)| (v - l) * texture_gain) + .collect(); + + let radius = ((scale / 30.0).round() as usize).clamp(4, 96); + let guided = GuidedFilter::new(&low, w, h, radius); + let eps = (0.02 + 0.06 * params.smoothing).powi(2); + + let low_smooth = guided.filter(&low, eps); + let mut y_low: Vec = low + .par_iter() + .zip(low_smooth.par_iter()) + .map(|(&l, &s)| l + params.smoothing * (s - l)) + .collect(); + + if params.even_tone > 0.001 { + for plane in [&mut cb, &mut cr] { + let local = guided.filter(plane, eps); + let broad = weighted_blur(plane, &skin_weight, w, h, radius as f32 * 2.0); + plane + .par_iter_mut() + .zip(local.par_iter()) + .zip(broad.par_iter()) + .for_each(|((v, &l), &b)| { + *v += params.even_tone * (l - *v) + 0.5 * params.even_tone * (b - l); + }); + } + } + + if params.shine > 0.001 { + let mut hist = [0u32; 1024]; + let (mut total, mut high_dev) = (0u32, 0.0_f64); + for i in 0..n { + if mask[i] > 76 && skin[i] > 0.5 { + hist[(y_low[i].clamp(0.0, 1.0) * 1023.0) as usize] += 1; + high_dev += high[i].abs() as f64; + total += 1; + } + } + + if total >= 64 { + let median = percentile(&hist, total, 0.5); + let peak = percentile(&hist, total, 0.995); + let range = ((peak - median) * 0.6).max(0.04); + let strength = params.shine; + + // Shine is whatever rises above the broad local skin brightness. + let ambient = weighted_blur(&y_low, &skin_weight, w, h, scale / 8.0); + let mut shine = vec![0.0_f32; n]; + y_low + .par_iter_mut() + .zip(shine.par_iter_mut()) + .zip(ambient.par_iter()) + .for_each(|((v, s), &a)| { + let floor = a + 0.02; + let excess = *v - floor; + if excess > 0.0 { + *s = smoothstep(0.0, range, excess); + *v = floor + excess / (1.0 + 3.0 * strength * excess / range); + } + }); + + // Specular highlights are desaturated, so bring back the surrounding skin color. + let matte_weight: Vec = skin_weight + .par_iter() + .zip(shine.par_iter()) + .map(|(&k, &s)| k * (1.0 - s)) + .collect(); + for plane in [&mut cb, &mut cr] { + let reference = weighted_blur(plane, &matte_weight, w, h, radius as f32 * 1.5); + plane + .par_iter_mut() + .zip(reference.par_iter()) + .zip(shine.par_iter()) + .for_each(|((v, &r), &s)| *v += 0.8 * strength * s * (r - *v)); + } + + // Sweat sparkles live in the texture layer as isolated bright outliers. + let sig_h = ((1.2533 * high_dev / total as f64) as f32).max(0.002); + high.par_iter_mut().for_each(|v| { + *v *= 1.0 - 0.85 * strength * smoothstep(2.5, 4.5, *v / sig_h); + }); + } + } + + let mut out = vec![0u8; n * 3]; + out.par_chunks_exact_mut(3) + .zip(rgb.par_chunks_exact(3)) + .enumerate() + .for_each(|(i, (dst, src))| { + if mask[i] == 0 { + dst.copy_from_slice(src); + return; + } + let m = mask[i] as f32 / 255.0; + let blend = m * m * (3.0 - 2.0 * m) * (1.0 - params.skin_protection * (1.0 - skin[i])); + + let luma = y_low[i] + high[i]; + let result = [ + luma + 1.403 * cr[i], + luma - 0.344 * cb[i] - 0.714 * cr[i], + luma + 1.773 * cb[i], + ]; + for c in 0..3 { + let original = src[c] as f32; + let target = result[c] * 255.0; + dst[c] = (original + blend * (target - original)) + .round() + .clamp(0.0, 255.0) as u8; + } + }); + out +} + +/// Faces wider than this are retouched on a reduced copy; see `retouch_skin_scaled`. +const WORKING_SCALE: f32 = 640.0; + +/// Area-averaging reduction of an interleaved 8-bit buffer. +fn shrink( + src: &[u8], + w: usize, + h: usize, + channels: usize, + small_w: usize, + small_h: usize, +) -> Vec { + let mut out = vec![0u8; small_w * small_h * channels]; + out.par_chunks_exact_mut(small_w * channels) + .enumerate() + .for_each(|(sy, row)| { + let y0 = sy * h / small_h; + let y1 = ((sy + 1) * h / small_h).max(y0 + 1).min(h); + for sx in 0..small_w { + let x0 = sx * w / small_w; + let x1 = ((sx + 1) * w / small_w).max(x0 + 1).min(w); + let count = ((x1 - x0) * (y1 - y0)) as u32; + for c in 0..channels { + let mut sum = 0u32; + for y in y0..y1 { + for x in x0..x1 { + sum += src[(y * w + x) * channels + c] as u32; + } + } + row[sx * channels + c] = ((sum + count / 2) / count) as u8; + } + } + }); + out +} + +/// Same result as `retouch_skin`, but large faces are processed on a reduced copy and only +/// the change is brought back to full size, so the cost no longer grows with the face. +/// +/// Everything the retouch alters is coarser than the reduced copy can represent except +/// the finest skin texture, which is handled separately at full size. +pub fn retouch_skin_scaled( + rgb: &[u8], + mask: &[u8], + w: usize, + h: usize, + scale: f32, + params: &SkinRetouchParams, + hints: Option, +) -> Vec { + let factor = scale / WORKING_SCALE; + if factor <= 1.25 { + return retouch_skin(rgb, mask, w, h, scale, params, hints); + } + + let small_w = ((w as f32 / factor).round() as usize).max(8); + let small_h = ((h as f32 / factor).round() as usize).max(8); + let small = shrink(rgb, w, h, 3, small_w, small_h); + let small_mask = shrink(mask, w, h, 1, small_w, small_h); + let small_hints = hints.map(|hints| { + ( + shrink(hints.skin, w, h, 1, small_w, small_h), + shrink(hints.features, w, h, 1, small_w, small_h), + hints + .blemish_zone + .map(|zone| shrink(zone, w, h, 1, small_w, small_h)), + ) + }); + + let retouched = retouch_skin( + &small, + &small_mask, + small_w, + small_h, + WORKING_SCALE, + params, + small_hints + .as_ref() + .map(|(skin, features, blemish_zone)| FaceHints { + skin, + features, + blemish_zone: blemish_zone.as_deref(), + }), + ); + + let texture_gain = 0.25 + 0.75 * params.texture; + let (step_x, step_y) = (small_w as f32 / w as f32, small_h as f32 / h as f32); + let mut out = vec![0u8; rgb.len()]; + out.par_chunks_exact_mut(w * 3) + .enumerate() + .for_each(|(y, row)| { + let fy = ((y as f32 + 0.5) * step_y - 0.5).clamp(0.0, (small_h - 1) as f32); + let y0 = fy as usize; + let y1 = (y0 + 1).min(small_h - 1); + let ty = fy - y0 as f32; + + for x in 0..w { + let i = y * w + x; + let src = &rgb[i * 3..i * 3 + 3]; + let dst = &mut row[x * 3..x * 3 + 3]; + if mask[i] == 0 { + dst.copy_from_slice(src); + continue; + } + + let fx = ((x as f32 + 0.5) * step_x - 0.5).clamp(0.0, (small_w - 1) as f32); + let x0 = fx as usize; + let x1 = (x0 + 1).min(small_w - 1); + let tx = fx - x0 as f32; + let sample = |buf: &[u8], c: usize| { + let at = |sx: usize, sy: usize| buf[(sy * small_w + sx) * 3 + c] as f32; + let top = at(x0, y0) + (at(x1, y0) - at(x0, y0)) * tx; + let bottom = at(x0, y1) + (at(x1, y1) - at(x0, y1)) * tx; + top + (bottom - top) * ty + }; + + // Pores finer than the reduced copy follow the texture setting on skin only. + let fine_weight = match hints { + Some(hints) => { + let skin = + hints.skin[i] as f32 * (255 - hints.features[i]) as f32 / 65025.0; + mask[i] as f32 / 255.0 * skin * (texture_gain - 1.0) + } + None => 0.0, + }; + for c in 0..3 { + let original = src[c] as f32; + let coarse = sample(&small, c); + let change = sample(&retouched, c) - coarse; + dst[c] = (original + change + fine_weight * (original - coarse)) + .round() + .clamp(0.0, 255.0) as u8; + } + } + }); + out +} + +/// Eye and teeth regions of a face, as `0..=255` planes aligned with the image buffer, +/// and how strongly (`0.0..=1.0`) each should be enhanced. +pub struct FeatureEnhance<'a> { + pub eyes: &'a [u8], + pub eye_whites: &'a [u8], + pub teeth: &'a [u8], + pub eye_amount: f32, + pub teeth_amount: f32, +} + +/// Adds clarity to the eyes, cleans up their whites and whitens teeth, in place on an +/// sRGB encoded RGB8 buffer. +pub fn enhance_eyes_and_teeth( + rgb: &mut [u8], + w: usize, + h: usize, + scale: f32, + features: &FeatureEnhance, +) { + let luma_of = + |px: &[u8]| (0.299 * px[0] as f32 + 0.587 * px[1] as f32 + 0.114 * px[2] as f32) / 255.0; + + // Clarity compares each eye pixel with a blurred copy; only the box around the eyes + // is blurred, since that is the only place the result is used. + let sigma = (scale / 150.0).clamp(1.0, 6.0); + let mut eye_box = None; + if features.eye_amount > 0.0 { + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0, 0); + for (i, &eye) in features.eyes.iter().enumerate() { + if eye > 0 { + x0 = x0.min(i % w); + x1 = x1.max(i % w); + y0 = y0.min(i / w); + y1 = y1.max(i / w); + } + } + if x0 <= x1 { + let margin = (sigma * 3.0).ceil() as usize; + let (x0, y0) = (x0.saturating_sub(margin), y0.saturating_sub(margin)); + let (x1, y1) = ((x1 + margin + 1).min(w), (y1 + margin + 1).min(h)); + let (bw, bh) = (x1 - x0, y1 - y0); + let mut luma = vec![0.0_f32; bw * bh]; + luma.par_chunks_exact_mut(bw) + .enumerate() + .for_each(|(by, row)| { + for (bx, value) in row.iter_mut().enumerate() { + let i = ((y0 + by) * w + x0 + bx) * 3; + *value = luma_of(&rgb[i..i + 3]); + } + }); + eye_box = Some((x0, y0, bw, bh, gauss_blur(&luma, bw, bh, sigma))); + } + } + + rgb.par_chunks_exact_mut(3).enumerate().for_each(|(i, px)| { + let clarity = features.eyes[i] as f32 / 255.0 * features.eye_amount; + let whites = features.eye_whites[i] as f32 / 255.0 * features.eye_amount; + let teeth = features.teeth[i] as f32 / 255.0 * features.teeth_amount; + if clarity <= 0.0 && teeth <= 0.0 { + return; + } + + let (r, b) = (px[0] as f32 / 255.0, px[2] as f32 / 255.0); + let mut y = luma_of(px); + let mut cb = (b - y) * 0.564; + let mut cr = (r - y) * 0.713; + + if let Some((x0, y0, bw, bh, soft)) = &eye_box { + let (x, row) = (i % w, i / w); + if clarity > 0.0 && x >= *x0 && row >= *y0 && x < x0 + bw && row < y0 + bh { + y += 0.9 * clarity * (y - soft[(row - y0) * bw + (x - x0)]); + } + } + y *= 1.0 + 0.22 * whites + 0.16 * teeth; + let keep_color = (1.0 - 0.45 * whites) * (1.0 - 0.6 * teeth); + cb *= keep_color; + cr *= keep_color; + + let result = [y + 1.403 * cr, y - 0.344 * cb - 0.714 * cr, y + 1.773 * cb]; + for c in 0..3 { + px[c] = (result[c] * 255.0).round().clamp(0.0, 255.0) as u8; + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + + const SIZE: usize = 240; + + fn params() -> SkinRetouchParams { + SkinRetouchParams { + smoothing: 0.5, + texture: 0.75, + blemish: 0.6, + shine: 0.5, + even_tone: 0.4, + skin_protection: 0.85, + } + } + + fn skin_image() -> Vec { + let mut seed = 12345_u32; + let mut rgb = Vec::with_capacity(SIZE * SIZE * 3); + for _ in 0..SIZE * SIZE { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + let noise = ((seed >> 24) as i32 % 7) - 3; + for base in [205_i32, 155, 130] { + rgb.push((base + noise) as u8); + } + } + rgb + } + + fn paint_disc(rgb: &mut [u8], cx: usize, cy: usize, radius: i32, color: [u8; 3]) { + for dy in -radius..=radius { + for dx in -radius..=radius { + if dx * dx + dy * dy <= radius * radius { + let i = ((cy as i32 + dy) as usize * SIZE + (cx as i32 + dx) as usize) * 3; + rgb[i..i + 3].copy_from_slice(&color); + } + } + } + } + + fn luma_at(rgb: &[u8], x: usize, y: usize) -> f32 { + let i = (y * SIZE + x) * 3; + 0.299 * rgb[i] as f32 + 0.587 * rgb[i + 1] as f32 + 0.114 * rgb[i + 2] as f32 + } + + #[test] + fn removes_dark_spot_and_keeps_unmasked_pixels() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 120, 120, 4, [150, 95, 80]); + + let mut mask = vec![255u8; SIZE * SIZE]; + for row in mask.chunks_exact_mut(SIZE).take(20) { + row.fill(0); + } + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, ¶ms(), None); + + assert_eq!(&out[..20 * SIZE * 3], &rgb[..20 * SIZE * 3]); + + let surround = luma_at(&rgb, 60, 120); + let before = surround - luma_at(&rgb, 120, 120); + let after = surround - luma_at(&out, 120, 120); + assert!(after < before * 0.4, "spot went from {before} to {after}"); + } + + #[test] + fn tames_specular_highlight() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 120, 120, 14, [250, 240, 235]); + let mask = vec![255u8; SIZE * SIZE]; + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, ¶ms(), None); + assert!(luma_at(&out, 120, 120) < luma_at(&rgb, 120, 120) - 8.0); + } + + #[test] + fn face_hints_protect_skin_colored_features() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 120, 120, 4, [150, 95, 80]); + let mask = vec![255u8; SIZE * SIZE]; + let skin = vec![255u8; SIZE * SIZE]; + let mut features = vec![0u8; SIZE * SIZE]; + for y in 100..140 { + features[y * SIZE + 100..y * SIZE + 140].fill(255); + } + let hints = FaceHints { + skin: &skin, + features: &features, + blemish_zone: None, + }; + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, ¶ms(), Some(hints)); + assert!((luma_at(&out, 120, 120) - luma_at(&rgb, 120, 120)).abs() < 6.0); + } + + #[test] + fn whitens_teeth_only_where_marked() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 120, 120, 10, [200, 185, 140]); + let original = rgb.clone(); + let none = vec![0u8; SIZE * SIZE]; + let mut teeth = vec![0u8; SIZE * SIZE]; + for y in 112..128 { + teeth[y * SIZE + 112..y * SIZE + 128].fill(255); + } + + enhance_eyes_and_teeth( + &mut rgb, + SIZE, + SIZE, + 480.0, + &FeatureEnhance { + eyes: &none, + eye_whites: &none, + teeth: &teeth, + eye_amount: 1.0, + teeth_amount: 1.0, + }, + ); + + let i = (120 * SIZE + 120) * 3; + let yellowness = |px: &[u8]| px[0] as i32 - px[2] as i32; + assert!(yellowness(&rgb[i..i + 3]) < yellowness(&original[i..i + 3]) / 2); + assert!(luma_at(&rgb, 120, 120) > luma_at(&original, 120, 120)); + assert_eq!(&rgb[..60 * SIZE * 3], &original[..60 * SIZE * 3]); + } + + #[test] + fn keeps_nostril_like_deep_shadows() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 120, 120, 5, [70, 38, 30]); + paint_disc(&mut rgb, 60, 120, 4, [150, 95, 80]); + let mask = vec![255u8; SIZE * SIZE]; + let strong = SkinRetouchParams { + blemish: 1.0, + ..params() + }; + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, &strong, None); + + let nostril_before = luma_at(&rgb, 120, 120); + assert!( + luma_at(&out, 120, 120) < nostril_before + 12.0, + "nostril went from {nostril_before} to {}", + luma_at(&out, 120, 120) + ); + // A real blemish next to it is still removed. + let surround = luma_at(&rgb, 60, 60); + assert!(surround - luma_at(&out, 60, 120) < (surround - luma_at(&rgb, 60, 120)) * 0.4); + } + + #[test] + fn scaled_retouch_matches_full_size_retouch() { + const BIG: usize = 720; + let mut seed = 99_u32; + let mut rgb = Vec::with_capacity(BIG * BIG * 3); + for _ in 0..BIG * BIG { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + let noise = ((seed >> 24) as i32 % 7) - 3; + for base in [205_i32, 155, 130] { + rgb.push((base + noise) as u8); + } + } + let center = (BIG / 2) as i32; + for dy in -12_i32..=12 { + for dx in -12_i32..=12 { + if dx * dx + dy * dy <= 144 { + let i = ((center + dy) as usize * BIG + (center + dx) as usize) * 3; + rgb[i..i + 3].copy_from_slice(&[150, 95, 80]); + } + } + } + let mut mask = vec![255u8; BIG * BIG]; + mask[..40 * BIG].fill(0); + let luma = |buf: &[u8], x: usize, y: usize| { + let i = (y * BIG + x) * 3; + 0.299 * buf[i] as f32 + 0.587 * buf[i + 1] as f32 + 0.114 * buf[i + 2] as f32 + }; + + let scale = 1440.0; + let full = retouch_skin(&rgb, &mask, BIG, BIG, scale, ¶ms(), None); + let scaled = retouch_skin_scaled(&rgb, &mask, BIG, BIG, scale, ¶ms(), None); + + assert_eq!(&scaled[..40 * BIG * 3], &rgb[..40 * BIG * 3]); + let c = BIG / 2; + let spot = luma(&rgb, 60, c) - luma(&rgb, c, c); + assert!(luma(&rgb, 60, c) - luma(&scaled, c, c) < spot * 0.4); + assert!((luma(&scaled, c, c) - luma(&full, c, c)).abs() < 12.0); + } + + #[test] + fn removes_spots_only_inside_the_blemish_zone() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 60, 120, 4, [150, 95, 80]); + paint_disc(&mut rgb, 180, 120, 4, [150, 95, 80]); + let mask = vec![255u8; SIZE * SIZE]; + let skin = vec![255u8; SIZE * SIZE]; + let features = vec![0u8; SIZE * SIZE]; + let mut zone = vec![0u8; SIZE * SIZE]; + for row in zone.chunks_exact_mut(SIZE) { + row[..120].fill(255); + } + let hints = FaceHints { + skin: &skin, + features: &features, + blemish_zone: Some(&zone), + }; + let only_blemish = SkinRetouchParams { + smoothing: 0.0, + texture: 1.0, + blemish: 1.0, + shine: 0.0, + even_tone: 0.0, + skin_protection: 0.85, + }; + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, &only_blemish, Some(hints)); + + let surround = luma_at(&rgb, 120, 60); + let spot = surround - luma_at(&rgb, 60, 120); + assert!( + surround - luma_at(&out, 60, 120) < spot * 0.4, + "inside the zone" + ); + assert!( + surround - luma_at(&out, 180, 120) > spot * 0.9, + "outside the zone" + ); + } + + #[test] + fn keeps_fold_like_shadows() { + let mut rgb = skin_image(); + // A soft dark crease, like a nasolabial fold, and a real spot away from it. + for step in 0..90 { + let (x, y) = (80 + step / 3, 70 + step); + paint_disc(&mut rgb, x, y, 3, [178, 128, 104]); + } + paint_disc(&mut rgb, 180, 120, 4, [150, 95, 80]); + let mask = vec![255u8; SIZE * SIZE]; + let only_blemish = SkinRetouchParams { + smoothing: 0.0, + texture: 1.0, + blemish: 1.0, + shine: 0.0, + even_tone: 0.0, + skin_protection: 0.85, + }; + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, &only_blemish, None); + + let surround = luma_at(&rgb, 30, 30); + for step in [20, 45, 70] { + let (x, y) = (80 + step / 3, 70 + step); + let before = surround - luma_at(&rgb, x, y); + let after = surround - luma_at(&out, x, y); + assert!( + after > before * 0.8, + "fold at {x},{y} went from {before} to {after}" + ); + } + let spot = surround - luma_at(&rgb, 180, 120); + assert!(surround - luma_at(&out, 180, 120) < spot * 0.4, "spot"); + } + + #[test] + fn protects_non_skin_colors() { + let mut rgb = skin_image(); + paint_disc(&mut rgb, 120, 120, 12, [40, 90, 160]); + let mask = vec![255u8; SIZE * SIZE]; + + let out = retouch_skin(&rgb, &mask, SIZE, SIZE, 480.0, ¶ms(), None); + let i = (120 * SIZE + 120) * 3; + for c in 0..3 { + assert!((out[i + c] as i32 - rgb[i + c] as i32).abs() <= 12); + } + } +} diff --git a/src/App.tsx b/src/App.tsx index 633499cb5c..5756388b6f 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -356,6 +356,7 @@ function App() { handleImageSelect, handleSelectSubfolder, handleSelectAlbum, + handleSelectFlagView, handleOpenFolder, handleNavBack, handleNavForward, @@ -730,6 +731,7 @@ function App() { onFolderSelect={(path, skipHistory) => handleSelectSubfolder(path, false, undefined, true, false, skipHistory) } + onSelectFlagView={handleSelectFlagView} onToggleFolder={handleToggleFolder} onOpenFolder={handleOpenFolder} onNavBack={handleNavBack} @@ -776,6 +778,7 @@ function App() { handleAlbumTreeContextMenu, handleSelectAlbum, handleSelectSubfolder, + handleSelectFlagView, handleToggleFolder, handleOpenFolder, handleNavBack, diff --git a/src/components/adjustments/Details.tsx b/src/components/adjustments/Details.tsx index 519a1747d3..13e63186bf 100644 --- a/src/components/adjustments/Details.tsx +++ b/src/components/adjustments/Details.tsx @@ -216,6 +216,19 @@ function DetailsPanel({ onDragStateChange={onDragStateChange} fillOrigin="min" /> + + handleAdjustmentChange(DetailsAdjustment.AiDenoiseColor, String(e.target.value)) + } + step={1} + value={adjustments.aiDenoiseColor ?? 100} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> )} + {!isAiFree && !hiddenTools.includes('facialRefine') && ( + + + + )} + {!hiddenTools.includes('lut') && ( void; + setAdjustments(adjustments: Partial | ((prev: Adjustments) => Adjustments)): void; +} + +const FACE_SLIDERS = [ + { key: 'intensity', defaultValue: 40 }, + { key: 'blemish', defaultValue: 50 }, + { key: 'shine', defaultValue: 40 }, + { key: 'evenTone', defaultValue: 35 }, + { key: 'texture', defaultValue: 75 }, + { key: 'eyes', defaultValue: 0 }, + { key: 'teeth', defaultValue: 0 }, +] as const; + +type FaceSliderKey = (typeof FACE_SLIDERS)[number]['key']; +type FaceSettings = Partial>; + +const REFINE_DEBOUNCE_MS = 150; + +const faceParameters = (patch: AiPatch) => patch.subMasks[0]?.parameters ?? {}; + +const withFaceParameters = (patch: AiPatch, parameters: Record): AiPatch => ({ + ...patch, + subMasks: [{ ...patch.subMasks[0], parameters: { ...faceParameters(patch), ...parameters } }], +}); + +const isSameFace = (patch: AiPatch, face: DetectedFace) => { + const center = faceParameters(patch).faceCenter; + return !!center && Math.hypot(center.x - face.centerX, center.y - face.centerY) < face.width * 0.5; +}; + +// The backend only needs the edits that change the pixels faces are read from, so the +// heavy maps and Facial Refine's own patches stay out of the request. +const sourceAdjustments = (adjustments: Adjustments) => ({ + ...adjustments, + aiPatches: (adjustments.aiPatches || []).filter((patch) => !patch.faceRefine), + fogDepthMap: null, + lensBlurDepthMap: null, + lutData: null, + masks: [], + relightNormalMap: null, +}); + +const requestFacePatch = async (center: { x: number; y: number }, parameters: Record) => + JSON.parse( + await invoke('generate_face_refine_patch', { + currentAdjustments: sourceAdjustments(useEditorStore.getState().adjustments), + faceCenter: [center.x, center.y], + parameters, + }), + ); + +export default function FacialRefine({ adjustments, setAdjustments, onDragStateChange }: FacialRefineProps) { + const { t } = useTranslation(); + const [isDetecting, setIsDetecting] = useState(false); + const [selectedFaceId, setSelectedFaceId] = useState(null); + const [busyFaces, setBusyFaces] = useState>(new Set()); + // Slider edits live here until their result is ready, so the image is only reprocessed + // once per change, when the settings and the retouched face are committed together. + const [drafts, setDrafts] = useState>({}); + const aiModelDownloadStatus = useProcessStore(selectAiModelDownloadStatus); + + const draftsRef = useRef(drafts); + draftsRef.current = drafts; + const queue = useRef>(Promise.resolve()); + const queued = useRef(new Set()); + const timers = useRef(new Map>()); + + const faces = (adjustments.aiPatches || []).filter((patch) => patch.faceRefine); + const selectedFace = faces.find((face) => face.id === selectedFaceId) ?? faces[0]; + + useEffect(() => { + const pendingTimers = timers.current; + return () => pendingTimers.forEach((timer) => clearTimeout(timer)); + }, []); + + const reportFailure = useCallback( + (err: unknown) => toast.error(t('adjustments.effects.facialRefineFailed', { error: String(err) })), + [t], + ); + + const refineFace = useCallback( + async (id: string) => { + const patch = useEditorStore.getState().adjustments.aiPatches?.find((p) => p.id === id); + if (!patch) return; + + const used = { ...draftsRef.current[id] }; + const parameters = { ...faceParameters(patch), ...used }; + setBusyFaces((prev) => new Set(prev).add(id)); + try { + const patchData = await requestFacePatch(parameters.faceCenter, parameters); + useEditorStore.getState().patchesSentToBackend.delete(id); + setAdjustments((prev: Adjustments) => ({ + ...prev, + aiPatches: (prev.aiPatches || []).map((p) => + p.id === id ? { ...withFaceParameters(p, used), patchData } : p, + ), + })); + // Anything moved again while this was running stays pending for the next pass. + setDrafts((prev) => { + const remaining = { ...prev[id] }; + (Object.keys(used) as FaceSliderKey[]).forEach((key) => { + if (remaining[key] === used[key]) delete remaining[key]; + }); + const next = { ...prev }; + if (Object.keys(remaining).length > 0) next[id] = remaining; + else delete next[id]; + return next; + }); + } catch (err) { + reportFailure(err); + } finally { + setBusyFaces((prev) => { + const next = new Set(prev); + next.delete(id); + return next; + }); + } + }, + [reportFailure, setAdjustments], + ); + + // Faces are refined one after another, each from the settings it has when its turn comes. + const requestRefine = useCallback( + (id: string) => { + if (queued.current.has(id)) return; + queued.current.add(id); + queue.current = queue.current.then(() => { + queued.current.delete(id); + return refineFace(id); + }); + }, + [refineFace], + ); + + const scheduleRefine = (id: string) => { + clearTimeout(timers.current.get(id)); + timers.current.set( + id, + setTimeout(() => requestRefine(id), REFINE_DEBOUNCE_MS), + ); + }; + + const handleDetect = async () => { + setIsDetecting(true); + try { + const current = useEditorStore.getState().adjustments; + const detected = await invoke('detect_faces_for_refine', { + currentAdjustments: sourceAdjustments(current), + }); + if (detected.length === 0) { + toast.info(t('adjustments.effects.facialRefineNoFaces')); + return; + } + + const existing = (current.aiPatches || []).filter((p) => p.faceRefine); + const nextFaces: AiPatch[] = []; + for (const [index, face] of detected.entries()) { + const location = { faceCenter: { x: face.centerX, y: face.centerY }, thumbnail: face.thumbnail }; + const name = t('adjustments.effects.facialRefineFace', { index: index + 1 }); + const known = existing.find((patch) => isSameFace(patch, face)); + if (known) { + nextFaces.push({ ...withFaceParameters(known, location), name }); + continue; + } + + const parameters: Record = { lines: [], ...location }; + FACE_SLIDERS.forEach((slider) => (parameters[slider.key] = slider.defaultValue)); + nextFaces.push({ + id: uuidv4(), + faceRefine: true, + invert: false, + isLoading: false, + name, + patchData: await requestFacePatch(location.faceCenter, parameters), + prompt: '', + visible: true, + subMasks: [ + { + id: uuidv4(), + invert: false, + mode: SubMaskMode.Additive, + opacity: 100, + parameters, + type: Mask.Retouch, + visible: true, + }, + ], + }); + } + + // Every new face is committed at once, and face patches stay last so they are + // applied on top of every other edit. + const { patchesSentToBackend } = useEditorStore.getState(); + nextFaces.forEach((face) => patchesSentToBackend.delete(face.id)); + setAdjustments((prev: Adjustments) => ({ + ...prev, + aiPatches: [...(prev.aiPatches || []).filter((p) => !p.faceRefine), ...nextFaces], + })); + if (!nextFaces.some((face) => face.id === selectedFaceId)) { + setSelectedFaceId(nextFaces[0].id); + } + } catch (err) { + reportFailure(err); + } finally { + setIsDetecting(false); + } + }; + + const handleSliderChange = (face: AiPatch, key: FaceSliderKey, value: number) => { + setDrafts((prev) => ({ ...prev, [face.id]: { ...prev[face.id], [key]: value } })); + scheduleRefine(face.id); + }; + + const handleToggleFace = (face: AiPatch, visible: boolean) => { + setAdjustments((prev: Adjustments) => ({ + ...prev, + aiPatches: (prev.aiPatches || []).map((patch) => (patch.id === face.id ? { ...patch, visible } : patch)), + })); + if (visible && !face.patchData) scheduleRefine(face.id); + }; + + const handleRemoveFace = (id: string) => { + setAdjustments((prev: Adjustments) => ({ + ...prev, + aiPatches: (prev.aiPatches || []).filter((patch) => patch.id !== id), + })); + }; + + const settingOf = (face: AiPatch, key: FaceSliderKey, defaultValue: number): number => + drafts[face.id]?.[key] ?? faceParameters(face)[key] ?? defaultValue; + + const handleApplyToAll = (source: AiPatch) => { + const shared: FaceSettings = {}; + FACE_SLIDERS.forEach((slider) => (shared[slider.key] = settingOf(source, slider.key, slider.defaultValue))); + + const others = faces.filter((face) => face.id !== source.id); + setDrafts((prev) => { + const next = { ...prev }; + others.forEach((face) => (next[face.id] = { ...shared })); + return next; + }); + others.forEach((face) => scheduleRefine(face.id)); + }; + + return ( +
+ + + {isDetecting && aiModelDownloadStatus && ( + + {t('editor.masks.settings.aiModelDownloading')} + {aiModelDownloadStatus} + + )} + + {faces.length === 0 && !isDetecting && ( + + {t('adjustments.effects.facialRefineHint')} + + )} + + {faces.length > 0 && selectedFace && ( + <> +
+ {faces.map((face) => ( + + ))} +
+ +
+
+ handleToggleFace(selectedFace, visible)} + /> +
+ +
+ +
+ {FACE_SLIDERS.map((slider) => ( + handleSliderChange(selectedFace, slider.key, parseInt(e.target.value, 10))} + onDragStateChange={onDragStateChange} + step={1} + value={settingOf(selectedFace, slider.key, slider.defaultValue)} + /> + ))} +
+ + {faces.length > 1 && ( + + )} + + )} +
+ ); +} diff --git a/src/components/panel/MainLibrary.tsx b/src/components/panel/MainLibrary.tsx index 3e8f914d65..d297fac299 100644 --- a/src/components/panel/MainLibrary.tsx +++ b/src/components/panel/MainLibrary.tsx @@ -31,6 +31,7 @@ import { RawStatus, EditedStatus, FlagStatus, + ImageFlag, LibraryDisplayMode, } from '../ui/AppProperties'; import { GroupBadgeInfo, GroupId } from '../../utils/imageGrouping'; @@ -194,6 +195,9 @@ export default function MainLibrary(props: MainLibraryProps) { }; const searchCriteria = useLibraryStore((state) => state.searchCriteria); + const flagView = useLibraryStore((state) => state.flagView); + const activeAlbumId = useLibraryStore((state) => state.activeAlbumId); + const isViewLoading = useLibraryStore((state) => state.isViewLoading); const setLibrary = useLibraryStore((state) => state.setLibrary); const translatedRatingFilterOptions = useMemo( @@ -576,6 +580,12 @@ export default function MainLibrary(props: MainLibraryProps) { {props.currentFolderPath ? ( {props.currentFolderPath} + {flagView && + ` · ${t( + flagView === ImageFlag.Pick + ? 'library.folders.flagViews.picked' + : 'library.folders.flagViews.rejected', + )}`} ) : (

@@ -663,6 +673,7 @@ export default function MainLibrary(props: MainLibraryProps) { ) : ( {t('library.status.moment')} + ) : flagView && isViewLoading ? ( +
+ +
) : searchCriteria.tags.length > 0 || searchCriteria.text ? (
- {t('library.filters.noMatch')} + {flagView ? t('library.folders.flagViews.empty') : t('library.filters.noMatch')}
)} {props.isAndroid && ( diff --git a/src/components/panel/right/AIPanel.tsx b/src/components/panel/right/AIPanel.tsx index 91d7457d0a..5c87e27ac6 100644 --- a/src/components/panel/right/AIPanel.tsx +++ b/src/components/panel/right/AIPanel.tsx @@ -1,4 +1,4 @@ -import { useState, useEffect, useRef, useCallback } from 'react'; +import { useState, useEffect, useRef, useCallback, useMemo } from 'react'; import { v4 as uuidv4 } from 'uuid'; import { motion, AnimatePresence } from 'framer-motion'; import { useTranslation } from 'react-i18next'; @@ -59,7 +59,7 @@ import { getSubMaskName, NewMaskDropZone, } from './Masks'; -import { Adjustments, AiPatch } from '../../../utils/adjustments'; +import { Adjustments, AiPatch, insertAiPatch } from '../../../utils/adjustments'; import { OPTION_SEPARATOR } from '../../ui/AppProperties'; import { createSubMask } from '../../../utils/maskUtils'; import Text from '../../ui/Text'; @@ -109,7 +109,14 @@ const SUB_MASK_CONFIG: any = { }, [Mask.Retouch]: { showBrushTools: true, - parameters: [{ key: 'intensity', min: 1, max: 100, step: 1, defaultValue: 40 }], + parameters: [ + { key: 'intensity', min: 0, max: 100, step: 1, defaultValue: 40 }, + { key: 'texture', min: 0, max: 100, step: 1, defaultValue: 75 }, + { key: 'blemish', min: 0, max: 100, step: 1, defaultValue: 50 }, + { key: 'shine', min: 0, max: 100, step: 1, defaultValue: 40 }, + { key: 'evenTone', min: 0, max: 100, step: 1, defaultValue: 35 }, + { key: 'skinProtection', min: 0, max: 100, step: 1, defaultValue: 85 }, + ], }, [Mask.Linear]: { parameters: [] }, [Mask.AiSubject]: { @@ -336,7 +343,14 @@ export default function AIPanel() { const { t } = useTranslation(); const activePatchContainerId = useEditorStore((s) => s.activeAiPatchContainerId); const activeSubMaskId = useEditorStore((s) => s.activeAiSubMaskId); - const adjustments = useEditorStore((s) => s.adjustments); + const storeAdjustments = useEditorStore((s) => s.adjustments); + const adjustments = useMemo( + () => ({ + ...storeAdjustments, + aiPatches: (storeAdjustments.aiPatches || []).filter((p: AiPatch) => !p.faceRefine), + }), + [storeAdjustments], + ); const brushSettings = useEditorStore((s) => s.brushSettings); const isAIConnectorConnected = useEditorStore((s) => s.isAIConnectorConnected); const selectedImage = useEditorStore((s) => s.selectedImage); @@ -500,7 +514,10 @@ export default function AIPanel() { }); }); handleDeselect(); - setAdjustments((prev: Adjustments) => ({ ...prev, aiPatches: [] })); + setAdjustments((prev: Adjustments) => ({ + ...prev, + aiPatches: (prev.aiPatches || []).filter((p: AiPatch) => p.faceRefine), + })); }; const createMaskLogic = (type: Mask, mode: SubMaskMode = SubMaskMode.Additive) => { @@ -586,7 +603,10 @@ export default function AIPanel() { visible: true, }; - setAdjustments((prev: Adjustments) => ({ ...prev, aiPatches: [...(prev.aiPatches || []), newContainer] })); + setAdjustments((prev: Adjustments) => ({ + ...prev, + aiPatches: insertAiPatch(prev.aiPatches || [], newContainer), + })); onSelectPatchContainer(newContainer.id); const isStandalone = isStandaloneMask(type); @@ -777,11 +797,7 @@ export default function AIPanel() { const insertPatchContainer = (container: AiPatch, insertIndex?: number) => { setAdjustments((prev: Adjustments) => { - const newPatches = [...(prev.aiPatches || [])]; - const targetIndex = Math.max(0, Math.min(insertIndex ?? newPatches.length, newPatches.length)); - - newPatches.splice(targetIndex, 0, container); - return { ...prev, aiPatches: newPatches }; + return { ...prev, aiPatches: insertAiPatch(prev.aiPatches || [], container, insertIndex) }; }); onSelectPatchContainer(container.id); @@ -975,7 +991,7 @@ export default function AIPanel() { const oldIndex = prev.aiPatches.findIndex((p) => p.id === dragData.item!.id); let newIndex = -1; - if (overId === 'ai-list-root') newIndex = prev.aiPatches.length - 1; + if (overId === 'ai-list-root') newIndex = prev.aiPatches.filter((p) => !p.faceRefine).length - 1; else if (overData?.type === 'Container') newIndex = prev.aiPatches.findIndex((p) => p.id === overId); else if (overData?.type === 'SubMask') newIndex = prev.aiPatches.findIndex((p) => p.id === overData.parentId); diff --git a/src/components/panel/right/FolderTree.tsx b/src/components/panel/right/FolderTree.tsx index fcd3058872..007357f5ae 100644 --- a/src/components/panel/right/FolderTree.tsx +++ b/src/components/panel/right/FolderTree.tsx @@ -39,8 +39,17 @@ import { useShallow } from 'zustand/react/shallow'; import { useLibraryStore } from '../../../store/useLibraryStore'; import { useSettingsStore } from '../../../store/useSettingsStore'; import { useUIStore } from '../../../store/useUIStore'; -import { AlbumItem, AlbumGroup, Album, Invokes, FolderTreeSort, SortDirection } from '../../ui/AppProperties'; +import { + AlbumItem, + AlbumGroup, + Album, + Invokes, + FolderTreeSort, + ImageFlag, + SortDirection, +} from '../../ui/AppProperties'; import { useLibraryActions } from '../../../hooks/useLibraryActions'; +import { countFlaggedImages, FLAG_ICONS } from '../../../utils/imageFlags'; import ImmichSection from '../../immich/ImmichSection'; export interface FolderTree { @@ -60,6 +69,7 @@ interface FolderTreeProps { onAlbumContextMenu(event: any, item: AlbumItem | null): void; onFolderSelect(folder: string, skipHistory?: boolean): void; onSelectAlbum(albumId: string, albumName: string, images: string[], skipHistory?: boolean): void; + onSelectFlagView(flagView: ImageFlag | null): void; onToggleFolder(folder: string): void; onOpenFolder(): void; onNavBack(): void; @@ -103,6 +113,11 @@ const ALBUM_ICONS: Record = { briefcase: Briefcase, }; +const FLAG_VIEWS = [ + { flag: ImageFlag.Pick, labelKey: 'library.folders.flagViews.picked' }, + { flag: ImageFlag.Reject, labelKey: 'library.folders.flagViews.rejected' }, +] as const; + const filterTree = (node: FolderTree | null, query: string): FolderTree | null => { if (!node) { return null; @@ -338,6 +353,28 @@ function SectionHeader({ title, isOpen, onToggle }: { title: string; isOpen: boo ); } +function FlagViewCount({ flag }: { flag: ImageFlag }) { + const imageList = useLibraryStore((state) => state.imageList); + const scopeFlags = useLibraryStore((state) => state.scopeFlags); + const count = useMemo(() => countFlaggedImages(scopeFlags, imageList, flag), [scopeFlags, imageList, flag]); + + return ( + <> + {count !== null && ( + + {count} + + )} +