From 39b9e50ff645a44f4a91175669ec8d652bcdd7d7 Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Sat, 26 Sep 2026 16:18:13 -0500 Subject: [PATCH 01/16] fix(presets): support multi-preset toggle and add active border indicator --- src/components/panel/right/PresetsPanel.tsx | 186 ++++++++++++++++---- src/store/useEditorStore.ts | 16 +- 2 files changed, 165 insertions(+), 37 deletions(-) diff --git a/src/components/panel/right/PresetsPanel.tsx b/src/components/panel/right/PresetsPanel.tsx index b1e7e65313..8914c54cf2 100644 --- a/src/components/panel/right/PresetsPanel.tsx +++ b/src/components/panel/right/PresetsPanel.tsx @@ -43,7 +43,7 @@ import Slider from '../../ui/Slider'; import { TextColors, TextVariants, TextWeights } from '../../../types/typography'; import { Adjustments, INITIAL_ADJUSTMENTS, ADJUSTMENT_GROUPS } from '../../../utils/adjustments'; import { Invokes, OPTION_SEPARATOR, Panel, Preset, SelectedImage } from '../../ui/AppProperties'; -import { useEditorStore } from '../../../store/useEditorStore'; +import { useEditorStore, ActivePresetState } from '../../../store/useEditorStore'; import { useUIStore } from '../../../store/useUIStore'; import { useEditorActions } from '../../../hooks/useEditorActions'; @@ -188,7 +188,6 @@ const mixAdjustments = (presetObj: any, intensity: number, initialObj: any = INI const fraction = intensity / 100; if (fraction === 1) return { ...presetObj }; - if (fraction === 0) return { ...initialObj }; const result: any = {}; const keys = Object.keys(presetObj); @@ -257,6 +256,26 @@ const mixAdjustments = (presetObj: any, intensity: number, initialObj: any = INI return result; }; +const applyPresetStack = (base: Adjustments, activeList: ActivePresetState[]): Adjustments => { + let current = { ...base }; + for (const { preset, intensity } of activeList) { + if (!preset.adjustments) continue; + if (intensity === 100) { + current = { + ...current, + ...preset.adjustments, + }; + } else { + const mixed = mixAdjustments(preset.adjustments, intensity, current); + current = { + ...current, + ...mixed, + }; + } + } + return current; +}; + function PresetItemDisplay({ preset, previewUrl, @@ -283,7 +302,11 @@ function PresetItemDisplay({ }, [supportsMasks, supportsGeometry, t]); return ( -
+
s.selectedImage); const adjustments = useEditorStore((s) => s.adjustments); + const activePresets = useEditorStore((s) => s.activePresets); + const presetBaseAdjustments = useEditorStore((s) => s.presetBaseAdjustments); const setEditor = useEditorStore((s) => s.setEditor); const { setAdjustments } = useEditorActions(); @@ -580,9 +605,20 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp const [folderPreviewsGenerated, setFolderPreviewsGenerated] = useState>(new Set()); const [deletingItemId, setDeletingItemId] = useState(null); - const [activePresetId, setActivePresetId] = useState(null); - const [presetIntensity, setPresetIntensity] = useState(100); - const [baseAdjustments, setBaseAdjustments] = useState(null); + const activePresetsRef = useRef(activePresets); + activePresetsRef.current = activePresets; + const presetBaseAdjustmentsRef = useRef(presetBaseAdjustments); + presetBaseAdjustmentsRef.current = presetBaseAdjustments; + const adjustmentsRef = useRef(adjustments); + adjustmentsRef.current = adjustments; + + const activePresetMap = useMemo(() => { + const map = new Map(); + for (const item of activePresets) { + map.set(item.preset.id, item); + } + return map; + }, [activePresets]); const previewsRef = useRef(previews); previewsRef.current = previews; @@ -862,8 +898,10 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp setPreviews({}); setFolderPreviewsGenerated(new Set()); - setActivePresetId(null); - setBaseAdjustments(null); + setEditor({ + activePresets: [], + presetBaseAdjustments: null, + }); if (isPathChanged && selectedImage?.path) { currentImagePathRef.current = selectedImage.path; @@ -883,38 +921,65 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp generateRootPreviews, generateFolderPreviews, expandedFolders, + setEditor, ]); - const handleApplyPreset = (preset: Preset) => { - if (activePresetId === preset.id) { - setActivePresetId(null); - if (baseAdjustments) { - setAdjustments(baseAdjustments); + const handleApplyPreset = useCallback( + (preset: Preset) => { + const currentActive = activePresetsRef.current; + const isAlreadyActive = currentActive.some((p) => p.preset.id === preset.id); + + if (isAlreadyActive) { + // Deactivate this preset + const nextActive = currentActive.filter((p) => p.preset.id !== preset.id); + + if (nextActive.length === 0) { + if (presetBaseAdjustmentsRef.current) { + setAdjustments(presetBaseAdjustmentsRef.current); + } + setEditor({ + activePresets: [], + presetBaseAdjustments: null, + }); + } else { + const base = presetBaseAdjustmentsRef.current ?? adjustmentsRef.current; + const nextAdjustments = applyPresetStack(base, nextActive); + setAdjustments(nextAdjustments); + setEditor({ + activePresets: nextActive, + }); + } + return; } - setBaseAdjustments(null); - return; - } - setBaseAdjustments(adjustments); - setActivePresetId(preset.id); - setPresetIntensity(100); + // Activate this preset + const base = presetBaseAdjustmentsRef.current ?? adjustmentsRef.current; + const nextActive = [...currentActive, { preset, intensity: 100 }]; + const nextAdjustments = applyPresetStack(base, nextActive); - setAdjustments((prevAdjustments: Adjustments) => ({ - ...prevAdjustments, - ...preset.adjustments, - })); - }; + setAdjustments(nextAdjustments); + setEditor({ + activePresets: nextActive, + ...(presetBaseAdjustmentsRef.current ? {} : { presetBaseAdjustments: base }), + }); + }, + [setAdjustments, setEditor], + ); const handleIntensityChange = useCallback( (preset: Preset, intensity: number) => { - setPresetIntensity(intensity); - const mixed = mixAdjustments(preset.adjustments, intensity); - setAdjustments((prev: Adjustments) => ({ - ...prev, - ...mixed, - })); + const base = presetBaseAdjustmentsRef.current ?? adjustmentsRef.current; + const nextActive = activePresetsRef.current.map((p) => + p.preset.id === preset.id ? { ...p, intensity } : p, + ); + const nextAdjustments = applyPresetStack(base, nextActive); + setAdjustments(nextAdjustments); + setEditor({ + activePresets: nextActive, + ...(presetBaseAdjustmentsRef.current ? {} : { presetBaseAdjustments: base }), + }); }, - [setAdjustments], + [setAdjustments, setEditor], ); const handleSaveConfiguredPreset = async ( @@ -932,6 +997,17 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp presetType, ); if (updated) { + if (activePresetsRef.current.some((p) => p.preset.id === updated.id)) { + const nextActive = activePresetsRef.current.map((p) => + p.preset.id === updated.id ? { ...p, preset: updated } : p, + ); + const base = presetBaseAdjustmentsRef.current ?? adjustmentsRef.current; + const nextAdjustments = applyPresetStack(base, nextActive); + setAdjustments(nextAdjustments); + setEditor({ + activePresets: nextActive, + }); + } await generateSinglePreview(updated); } } else { @@ -961,6 +1037,36 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp return; } + const currentActive = activePresetsRef.current; + let nextActive = currentActive; + + if (isFolder) { + const folderItem = presets.find((item: any) => item.folder && item.folder.id === id); + const childIds = new Set((folderItem?.folder?.children || []).map((c: any) => c.id)); + nextActive = currentActive.filter((p) => !childIds.has(p.preset.id)); + } else { + nextActive = currentActive.filter((p) => p.preset.id !== id); + } + + if (nextActive.length !== currentActive.length) { + if (nextActive.length === 0) { + if (presetBaseAdjustmentsRef.current) { + setAdjustments(presetBaseAdjustmentsRef.current); + } + setEditor({ + activePresets: [], + presetBaseAdjustments: null, + }); + } else { + const base = presetBaseAdjustmentsRef.current ?? adjustmentsRef.current; + const nextAdjustments = applyPresetStack(base, nextActive); + setAdjustments(nextAdjustments); + setEditor({ + activePresets: nextActive, + }); + } + } + setTimeout(() => { deleteItem(id); if (isFolder) { @@ -1142,6 +1248,13 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp onClick: async () => { const updated = overwritePreset(data?.id ?? null); if (updated) { + if (activePresetsRef.current.some((p) => p.preset.id === updated.id)) { + setEditor({ + activePresets: activePresetsRef.current.map((p) => + p.preset.id === updated.id ? { ...p, preset: updated } : p, + ), + }); + } await generateSinglePreview(updated); } }, @@ -1316,8 +1429,8 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp onContextMenu={(e: any) => handleContextMenu(e, { preset })} preset={preset} previewUrl={previews[preset.id] || ''} - isActive={preset.id === activePresetId} - intensity={preset.id === activePresetId ? presetIntensity : 100} + isActive={activePresetMap.has(preset.id)} + intensity={activePresetMap.get(preset.id)?.intensity ?? 100} onIntensityChange={(val) => handleIntensityChange(preset, val)} onDragStateChange={handleDragStateChange} /> @@ -1347,9 +1460,10 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp onContextMenu={(e: any) => handleContextMenu(e, item)} preset={item.preset} previewUrl={(item.preset?.id ? previews[item.preset.id] : '') || ''} - isActive={item.preset?.id === activePresetId} - intensity={item.preset?.id === activePresetId ? presetIntensity : 100} + isActive={item.preset?.id ? activePresetMap.has(item.preset.id) : false} + intensity={item.preset?.id ? activePresetMap.get(item.preset.id)?.intensity ?? 100 : 100} onIntensityChange={(val) => handleIntensityChange(item.preset as Preset, val)} + onDragStateChange={handleDragStateChange} /> ))} diff --git a/src/store/useEditorStore.ts b/src/store/useEditorStore.ts index f05b5238c3..441df08fac 100644 --- a/src/store/useEditorStore.ts +++ b/src/store/useEditorStore.ts @@ -1,11 +1,16 @@ import { create } from 'zustand'; import { Adjustments, INITIAL_ADJUSTMENTS, MaskContainer, AiPatch } from '../utils/adjustments'; -import { SelectedImage, WaveformData, BrushSettings } from '../components/ui/AppProperties'; +import { SelectedImage, WaveformData, BrushSettings, Preset } from '../components/ui/AppProperties'; import { ChannelConfig } from '../components/adjustments/Curves'; import { ImageDimensions } from '../hooks/useImageRenderSize'; import { ToolType } from '../components/panel/right/Masks'; import { OverlayMode } from '../components/panel/right/CropPanel'; +export interface ActivePresetState { + preset: Preset; + intensity: number; +} + interface InteractivePatch { url: string; normX: number; @@ -79,6 +84,10 @@ interface EditorState { copiedMask: MaskContainer | null; copiedAdjustments: Adjustments | null; + // Presets State + activePresets: ActivePresetState[]; + presetBaseAdjustments: Adjustments | null; + // Actions setEditor: (updater: Partial | ((state: EditorState) => Partial)) => void; pushHistory: (newAdjustments: Adjustments) => void; @@ -130,6 +139,9 @@ export const useEditorStore = create((set) => ({ brushSettings: { size: 50, feather: 50, tool: ToolType.Brush }, copiedAdjustments: null, + activePresets: [], + presetBaseAdjustments: null, + isGeneratingAiMask: false, isAIConnectorConnected: false, isGeneratingAi: false, @@ -170,6 +182,8 @@ export const useEditorStore = create((set) => ({ history: [initialState], historyIndex: 0, adjustments: initialState, + activePresets: [], + presetBaseAdjustments: null, }), goToHistoryIndex: (index) => From 77917d8c605245e71989236ab4313753977cf764 Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Sat, 26 Sep 2026 16:26:34 -0500 Subject: [PATCH 02/16] feat: implement usePresets hook for managing user presets and folders --- src/hooks/usePresets.ts | 53 +++++++++++++++-------------------------- 1 file changed, 19 insertions(+), 34 deletions(-) diff --git a/src/hooks/usePresets.ts b/src/hooks/usePresets.ts index 91cbbf3d87..25f3680442 100644 --- a/src/hooks/usePresets.ts +++ b/src/hooks/usePresets.ts @@ -193,7 +193,7 @@ export function usePresets(currentAdjustments: Adjustments) { includeMasks: boolean, includeCropTransform: boolean, presetType: 'tool' | 'style', - ) => { + ): Preset | null => { let existingPreset: Preset | null = null; for (const item of presets) { @@ -243,17 +243,17 @@ export function usePresets(currentAdjustments: Adjustments) { for (const k of GEOMETRY_KEYS) delete newAdjustments[k]; } - let updatedPreset: Preset | null = null; + const updatedPreset: Preset = { + ...existingPreset, + name, + adjustments: newAdjustments, + includeMasks, + includeCropTransform, + presetType, + }; + const updatedPresets = presets.map((item: UserPreset) => { if (item.preset?.id === id) { - updatedPreset = { - ...item.preset, - name, - adjustments: newAdjustments, - includeMasks, - includeCropTransform, - presetType, - }; return { preset: updatedPreset }; } if (item.folder) { @@ -261,14 +261,6 @@ export function usePresets(currentAdjustments: Adjustments) { const newChildren = item.folder.children.map((child: Preset) => { if (child.id === id) { found = true; - updatedPreset = { - ...child, - name, - adjustments: newAdjustments, - includeMasks, - includeCropTransform, - presetType, - }; return updatedPreset; } return child; @@ -285,7 +277,7 @@ export function usePresets(currentAdjustments: Adjustments) { return updatedPreset; }; - const overwritePreset = (id: string | null) => { + const overwritePreset = (id: string | null): Preset | null => { let existingPreset: Preset | null = null; for (const item of presets) { @@ -337,16 +329,16 @@ export function usePresets(currentAdjustments: Adjustments) { } } - let updatedPreset: Preset | null = null; + const updatedPreset: Preset = { + ...existingPreset, + adjustments: presetAdjustments, + includeMasks, + includeCropTransform, + presetType, + }; + const updatedPresets = presets.map((item: UserPreset) => { if (item.preset?.id === id) { - updatedPreset = { - ...item.preset, - adjustments: presetAdjustments, - includeMasks, - includeCropTransform, - presetType, - }; return { preset: updatedPreset }; } if (item.folder) { @@ -354,13 +346,6 @@ export function usePresets(currentAdjustments: Adjustments) { const newChildren = item.folder.children.map((child: Preset) => { if (child.id === id) { found = true; - updatedPreset = { - ...child, - adjustments: presetAdjustments, - includeMasks, - includeCropTransform, - presetType, - }; return updatedPreset; } return child; From 9415492cab8cd76e184ff04ab82d3c0bf934dd2f Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Sat, 26 Sep 2026 16:58:32 -0500 Subject: [PATCH 03/16] style: format presets panel and clean up assignment --- src/components/panel/right/PresetsPanel.tsx | 22 +++++++++------------ 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/components/panel/right/PresetsPanel.tsx b/src/components/panel/right/PresetsPanel.tsx index 8914c54cf2..679f6862d2 100644 --- a/src/components/panel/right/PresetsPanel.tsx +++ b/src/components/panel/right/PresetsPanel.tsx @@ -969,9 +969,7 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp const handleIntensityChange = useCallback( (preset: Preset, intensity: number) => { const base = presetBaseAdjustmentsRef.current ?? adjustmentsRef.current; - const nextActive = activePresetsRef.current.map((p) => - p.preset.id === preset.id ? { ...p, intensity } : p, - ); + const nextActive = activePresetsRef.current.map((p) => (p.preset.id === preset.id ? { ...p, intensity } : p)); const nextAdjustments = applyPresetStack(base, nextActive); setAdjustments(nextAdjustments); setEditor({ @@ -1038,15 +1036,13 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp } const currentActive = activePresetsRef.current; - let nextActive = currentActive; - - if (isFolder) { - const folderItem = presets.find((item: any) => item.folder && item.folder.id === id); - const childIds = new Set((folderItem?.folder?.children || []).map((c: any) => c.id)); - nextActive = currentActive.filter((p) => !childIds.has(p.preset.id)); - } else { - nextActive = currentActive.filter((p) => p.preset.id !== id); - } + const nextActive = isFolder + ? (() => { + const folderItem = presets.find((item: any) => item.folder && item.folder.id === id); + const childIds = new Set((folderItem?.folder?.children || []).map((c: any) => c.id)); + return currentActive.filter((p) => !childIds.has(p.preset.id)); + })() + : currentActive.filter((p) => p.preset.id !== id); if (nextActive.length !== currentActive.length) { if (nextActive.length === 0) { @@ -1461,7 +1457,7 @@ export default function PresetsPanel({ onNavigateToCommunity }: PresetsPanelProp preset={item.preset} previewUrl={(item.preset?.id ? previews[item.preset.id] : '') || ''} isActive={item.preset?.id ? activePresetMap.has(item.preset.id) : false} - intensity={item.preset?.id ? activePresetMap.get(item.preset.id)?.intensity ?? 100 : 100} + intensity={item.preset?.id ? (activePresetMap.get(item.preset.id)?.intensity ?? 100) : 100} onIntensityChange={(val) => handleIntensityChange(item.preset as Preset, val)} onDragStateChange={handleDragStateChange} /> From 31b362f8b81c0e4d8be9f25c7824976adcffe983 Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Mon, 5 Oct 2026 14:55:22 -0500 Subject: [PATCH 04/16] feat: implement effects adjustment panel and Tauri backend processing commands --- src-tauri/src/adjustment_utils.rs | 3 +- src-tauri/src/ai_commands.rs | 28 ++ src-tauri/src/ai_processing.rs | 145 +++++++++ src-tauri/src/cache_utils.rs | 36 +++ src-tauri/src/file_management.rs | 3 +- src-tauri/src/lib.rs | 8 +- src-tauri/src/relight.rs | 290 ++++++++++++++++++ src/components/adjustments/Effects.tsx | 212 ++++++++++++- src/components/panel/Editor.tsx | 5 +- src/components/panel/editor/ImageCanvas.tsx | 51 +++ .../panel/editor/overlays/RelightHandle.tsx | 133 ++++++++ src/i18n/locales/ca.json | 14 + src/i18n/locales/de.json | 14 + src/i18n/locales/en.json | 14 + src/i18n/locales/es.json | 14 + src/i18n/locales/fr.json | 14 + src/i18n/locales/it.json | 14 + src/i18n/locales/ja.json | 14 + src/i18n/locales/ko.json | 14 + src/i18n/locales/pl.json | 14 + src/i18n/locales/pt.json | 14 + src/i18n/locales/ru.json | 14 + src/i18n/locales/zh-CN.json | 14 + src/i18n/locales/zh-TW.json | 14 + src/utils/adjustments.ts | 56 ++++ src/utils/relightUtils.ts | 91 ++++++ 26 files changed, 1237 insertions(+), 6 deletions(-) create mode 100644 src-tauri/src/relight.rs create mode 100644 src/components/panel/editor/overlays/RelightHandle.tsx create mode 100644 src/utils/relightUtils.ts diff --git a/src-tauri/src/adjustment_utils.rs b/src-tauri/src/adjustment_utils.rs index bacabf872d..0d6c8d1cd2 100644 --- a/src-tauri/src/adjustment_utils.rs +++ b/src-tauri/src/adjustment_utils.rs @@ -120,7 +120,8 @@ pub fn apply_all_transformations<'a, I: IntoCowImage<'a>>( let image = image.into_cow(); let warped_image = apply_geometry_warp(image, adjustments); - let blurred_image = crate::lens_blur::apply_lens_blur(warped_image, adjustments); + let relit_image = crate::relight::apply_relight(warped_image, adjustments); + let blurred_image = crate::lens_blur::apply_lens_blur(relit_image, adjustments); let (cropped_image, unscaled_crop_offset) = apply_spatial_transformations(blurred_image, adjustments); diff --git a/src-tauri/src/ai_commands.rs b/src-tauri/src/ai_commands.rs index 1673ce8681..493902767b 100644 --- a/src-tauri/src/ai_commands.rs +++ b/src-tauri/src/ai_commands.rs @@ -301,6 +301,34 @@ pub async fn generate_full_image_depth_map( Ok(format!("data:image/png;base64,{}", base64_str)) } +#[tauri::command] +pub async fn generate_full_image_normal_map( + js_adjustments: serde_json::Value, + state: tauri::State<'_, AppState>, + app_handle: tauri::AppHandle, +) -> Result { + let normal_model = crate::ai_processing::get_or_init_normal_model( + &app_handle, + &state.ai_state, + &state.ai_init_lock, + ) + .await + .map_err(|e| e.to_string())?; + + let warped_image = crate::get_cached_full_warped_image(&state, &js_adjustments)?; + + let normal_img = crate::ai_processing::run_normal_model(warped_image.as_ref(), &normal_model) + .map_err(|e| e.to_string())?; + + let mut buf = std::io::Cursor::new(Vec::new()); + normal_img + .write_to(&mut buf, image::ImageFormat::Png) + .map_err(|e| e.to_string())?; + let base64_str = base64::engine::general_purpose::STANDARD.encode(buf.get_ref()); + + Ok(format!("data:image/png;base64,{}", base64_str)) +} + #[allow(clippy::too_many_arguments)] #[tauri::command] pub async fn generate_ai_subject_mask( diff --git a/src-tauri/src/ai_processing.rs b/src-tauri/src/ai_processing.rs index 504e7263dc..a69ec98c1d 100644 --- a/src-tauri/src/ai_processing.rs +++ b/src-tauri/src/ai_processing.rs @@ -60,6 +60,12 @@ const DEPTH_FILENAME: &str = "depth_anything_v2_vits.onnx"; const DEPTH_INPUT_SIZE: u32 = 518; const DEPTH_SHA256: &str = "d2b11a11c1d4a12b47608fa65a17ee9a4c605b55ee1730c8e3b526304f2562be"; +const NORMAL_URL: &str = "https://huggingface.co/Ruicheng/moge-2-vits-normal-onnx/resolve/e50ffda41565591092adea54c6ac83d6212e1e23/model.onnx?download=true"; +const NORMAL_FILENAME: &str = "moge2_vits_normal.onnx"; +const NORMAL_SHA256: &str = "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0"; +const NORMAL_MAX_SIZE: u32 = 1024; +const NORMAL_NUM_TOKENS: i64 = 3600; + pub struct AiModels { pub sam_encoder: Mutex, pub sam_decoder: Mutex, @@ -92,6 +98,7 @@ pub struct AiState { pub denoise_model: Option>>, pub clip_models: Option>, pub lama_model: Option>>, + pub normal_model: Option>>, pub embeddings: Option, pub depth_map: Option, } @@ -550,6 +557,7 @@ pub async fn get_or_init_ai_models( denoise_model: None, clip_models: None, lama_model: None, + normal_model: None, embeddings: None, depth_map: None, }); @@ -610,6 +618,7 @@ pub async fn get_or_init_denoise_model( denoise_model: Some(denoise_model.clone()), clip_models: None, lama_model: None, + normal_model: None, embeddings: None, depth_map: None, }); @@ -682,6 +691,7 @@ pub async fn get_or_init_clip_models( denoise_model: None, clip_models: Some(clip_models.clone()), lama_model: None, + normal_model: None, embeddings: None, depth_map: None, }); @@ -742,6 +752,7 @@ pub async fn get_or_init_lama_model( denoise_model: None, clip_models: None, lama_model: Some(lama_model.clone()), + normal_model: None, embeddings: None, depth_map: None, }); @@ -750,6 +761,67 @@ pub async fn get_or_init_lama_model( Ok(lama_model) } +pub async fn get_or_init_normal_model( + app_handle: &tauri::AppHandle, + ai_state_mutex: &Mutex>, + ai_init_lock: &TokioMutex<()>, +) -> Result>> { + if let Some(normal_model) = ai_state_mutex + .lock() + .unwrap() + .as_ref() + .and_then(|state| state.normal_model.clone()) + { + return Ok(normal_model); + } + + let _guard = ai_init_lock.lock().await; + + if let Some(normal_model) = ai_state_mutex + .lock() + .unwrap() + .as_ref() + .and_then(|state| state.normal_model.clone()) + { + return Ok(normal_model); + } + + let models_dir = get_models_dir(app_handle)?; + download_and_verify_model( + app_handle, + &models_dir, + NORMAL_FILENAME, + NORMAL_URL, + NORMAL_SHA256, + "Normal Map Model", + ) + .await?; + + let _ = ort::init().with_name("AI-Normal").commit(); + let model_path = models_dir.join(NORMAL_FILENAME); + let session = Session::builder()?.commit_from_file(model_path)?; + let normal_model = Arc::new(Mutex::new(session)); + + crate::register_exit_handler(); + + let mut ai_state_lock = ai_state_mutex.lock().unwrap(); + if let Some(state) = ai_state_lock.as_mut() { + state.normal_model = Some(normal_model.clone()); + } else { + *ai_state_lock = Some(AiState { + models: None, + denoise_model: None, + clip_models: None, + lama_model: None, + normal_model: Some(normal_model.clone()), + embeddings: None, + depth_map: None, + }); + } + + Ok(normal_model) +} + #[derive(Clone, Copy)] struct TileParams { cs: usize, @@ -1684,6 +1756,79 @@ pub fn run_depth_anything_model( Ok(depth_map) } +/// Runs MoGe-2 and returns a camera-space normal map: +/// RGB = normal (X right, Y up, Z towards the viewer) packed as n * 0.5 + 0.5, +/// A = surface validity (0 where there is no geometry, e.g. sky). +pub fn run_normal_model( + image: &DynamicImage, + normal_session: &Mutex, +) -> Result { + let resized_image = if image.width().max(image.height()) > NORMAL_MAX_SIZE { + image.resize(NORMAL_MAX_SIZE, NORMAL_MAX_SIZE, FilterType::Triangle) + } else { + image.clone() + }; + let (resized_w, resized_h) = resized_image.dimensions(); + let resized_rgb = resized_image.into_rgb8(); + let raw_pixels = resized_rgb.as_raw(); + + let rw = resized_w as usize; + let rh = resized_h as usize; + + let mut input_tensor: Array = Array::zeros((1, 3, rh, rw)); + for y in 0..rh { + for x in 0..rw { + let idx = (y * rw + x) * 3; + input_tensor[[0, 0, y, x]] = raw_pixels[idx] as f32 / 255.0; + input_tensor[[0, 1, y, x]] = raw_pixels[idx + 1] as f32 / 255.0; + input_tensor[[0, 2, y, x]] = raw_pixels[idx + 2] as f32 / 255.0; + } + } + + let t_image = Tensor::from_array(input_tensor.into_dyn())?; + let t_num_tokens = Tensor::from_array(ndarray::arr0(NORMAL_NUM_TOKENS))?; + + let mut session = normal_session.lock().unwrap(); + let outputs = session.run(ort::inputs![ + "image" => t_image, + "num_tokens" => t_num_tokens, + ])?; + + let normal_tensor = outputs["normal"].try_extract_array::()?; + let mask_tensor = outputs["mask"].try_extract_array::()?; + + if normal_tensor.shape() != [1, rh, rw, 3] || mask_tensor.shape() != [1, rh, rw] { + return Err(anyhow::anyhow!( + "Unexpected normal model output shape: {:?} / {:?}", + normal_tensor.shape(), + mask_tensor.shape() + )); + } + + let normals = normal_tensor.as_standard_layout(); + let normals = normals.as_slice().unwrap(); + let mask = mask_tensor.as_standard_layout(); + let mask = mask.as_slice().unwrap(); + + let pack = |v: f32| ((v * 0.5 + 0.5).clamp(0.0, 1.0) * 255.0).round() as u8; + + let mut normal_data = Vec::with_capacity(rw * rh * 4); + for (n, m) in normals.chunks_exact(3).zip(mask) { + // The model outputs OpenCV camera space (Y down, Z away from the viewer). + let (nx, ny, nz) = (n[0], -n[1], -n[2]); + let len = (nx * nx + ny * ny + nz * nz).sqrt(); + if len.is_finite() && len > 1e-6 { + normal_data.extend_from_slice(&[pack(nx / len), pack(ny / len), pack(nz / len)]); + } else { + normal_data.extend_from_slice(&[128, 128, 255]); + } + normal_data.push((m.clamp(0.0, 1.0) * 255.0).round() as u8); + } + + RgbaImage::from_raw(resized_w, resized_h, normal_data) + .ok_or_else(|| anyhow::anyhow!("Failed to create normal map from model output")) +} + #[derive(Serialize, Deserialize, Debug, Clone, Default)] #[serde(rename_all = "camelCase")] pub struct AiSubjectMaskParameters { diff --git a/src-tauri/src/cache_utils.rs b/src-tauri/src/cache_utils.rs index b099105b67..f16bde67b2 100644 --- a/src-tauri/src/cache_utils.rs +++ b/src-tauri/src/cache_utils.rs @@ -77,9 +77,43 @@ pub fn calculate_patched_warped_hash(adjustments: &serde_json::Value) -> u64 { } } + hash_relight(adjustments, effects_visible, &mut hasher); + hasher.finish() } +fn hash_relight( + adjustments: &serde_json::Value, + effects_visible: bool, + hasher: &mut DefaultHasher, +) { + let relight_enabled = + effects_visible && adjustments["relightEnabled"].as_bool().unwrap_or(false); + relight_enabled.hash(hasher); + + if relight_enabled { + let relight_keys = [ + "relightIntensity", + "relightX", + "relightY", + "relightHeight", + "relightAmbient", + "relightSoftness", + "relightSpecular", + "relightRange", + "relightColor", + "relightNormalMap", + ]; + + for key in relight_keys { + if let Some(val) = adjustments.get(key) { + key.hash(hasher); + val.to_string().hash(hasher); + } + } + } +} + pub fn calculate_thumbnail_base_hash(adjustments: &serde_json::Value) -> u64 { let mut hasher = DefaultHasher::new(); @@ -166,6 +200,8 @@ pub fn calculate_transform_hash(adjustments: &serde_json::Value) -> u64 { } } + hash_relight(adjustments, effects_visible, &mut hasher); + if let Some(crop_val) = adjustments.get("crop") && !crop_val.is_null() { diff --git a/src-tauri/src/file_management.rs b/src-tauri/src/file_management.rs index 2a87459333..eff7762631 100644 --- a/src-tauri/src/file_management.rs +++ b/src-tauri/src/file_management.rs @@ -1575,7 +1575,8 @@ pub fn generate_thumbnail_data( let warped_image = apply_geometry_warp(Cow::Borrowed(&composite_image), &meta.adjustments); - let blurred_image = crate::lens_blur::apply_lens_blur(warped_image, &meta.adjustments); + let relit_image = crate::relight::apply_relight(warped_image, &meta.adjustments); + let blurred_image = crate::lens_blur::apply_lens_blur(relit_image, &meta.adjustments); let orientation_steps = meta.adjustments["orientationSteps"].as_u64().unwrap_or(0) as u8; diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 0ef3107514..31c203ccab 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -39,6 +39,7 @@ mod panorama_stitching; mod panorama_utils; mod preset_converter; mod raw_processing; +mod relight; mod tagging; mod tagging_utils; mod window_customizer; @@ -250,7 +251,8 @@ fn compute_patched_and_warped( }; let warped = apply_geometry_warp(patched_image, adjustments); - let blurred = crate::lens_blur::apply_lens_blur(warped, adjustments); + let relit = crate::relight::apply_relight(warped, adjustments); + let blurred = crate::lens_blur::apply_lens_blur(relit, adjustments); Ok(Arc::new(blurred.into_owned())) } @@ -819,8 +821,9 @@ async fn generate_uncropped_preview( Cow::Borrowed(loaded_image.image.as_ref()) }; + let relit_image = crate::relight::apply_relight(patched_image, &adjustments_clone); let blurred_image = - crate::lens_blur::apply_lens_blur(patched_image, &adjustments_clone); + crate::lens_blur::apply_lens_blur(relit_image, &adjustments_clone); let settings = load_settings(app_handle.clone()).unwrap_or_default(); let target_dim = (settings.editor_preview_resolution.unwrap_or(1920) as f32) as u32; @@ -2177,6 +2180,7 @@ pub fn run() { ai_commands::check_ai_connector_status, ai_commands::test_ai_connector_connection, ai_commands::generate_full_image_depth_map, + ai_commands::generate_full_image_normal_map, ai_commands::cancel_ai_task, apple_raw::is_raw9_available, inpainting::invoke_generative_replace_with_mask_def, diff --git a/src-tauri/src/relight.rs b/src-tauri/src/relight.rs new file mode 100644 index 0000000000..482b27664f --- /dev/null +++ b/src-tauri/src/relight.rs @@ -0,0 +1,290 @@ +use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; +use image::{DynamicImage, GenericImageView, RgbaImage}; +use rayon::prelude::*; +use std::borrow::Cow; + +const SPECULAR_POWER: f32 = 24.0; + +struct RelightParams { + // Light position as a fraction of the image width / height. + light_x: f32, + light_y: f32, + // Distance of the light from the image plane, in units of the longest image side. + light_z: f32, + // Falloff radius, in units of the longest image side. + radius: f32, + light_color: [f32; 3], + intensity: f32, + ambient: f32, + wrap: f32, + specular: f32, +} + +pub fn apply_relight<'a>( + image: Cow<'a, DynamicImage>, + adjustments: &serde_json::Value, +) -> Cow<'a, DynamicImage> { + let effects_visible = adjustments + .get("sectionVisibility") + .and_then(|v| v.get("effects")) + .and_then(|s| s.as_bool()) + .unwrap_or(true); + + if !adjustments["relightEnabled"].as_bool().unwrap_or(false) || !effects_visible { + return image; + } + + let normal_b64 = adjustments["relightNormalMap"].as_str().unwrap_or(""); + if normal_b64.is_empty() { + return image; + } + + let b64_data = match normal_b64.find(',') { + Some(idx) => &normal_b64[idx + 1..], + None => normal_b64, + }; + let decoded = match BASE64.decode(b64_data) { + Ok(b) => b, + Err(_) => return image, + }; + let normal_map = match image::load_from_memory(&decoded) { + Ok(img) => img.into_rgba8(), + Err(_) => return image, + }; + if normal_map.width() < 2 || normal_map.height() < 2 { + return image; + } + + let (w, h) = image.dimensions(); + if w < 2 || h < 2 { + return image; + } + + let params = read_params(adjustments); + + let start = std::time::Instant::now(); + let mut out = image.as_ref().to_rgb32f(); + relight_in_place(&mut out, &normal_map, ¶ms); + + log::info!("relight ({}x{}) took {:.2?}", w, h, start.elapsed()); + + Cow::Owned(DynamicImage::ImageRgb32F(out)) +} + +fn srgb_to_linear(v: f32) -> f32 { + if v <= 0.04045 { + v / 12.92 + } else { + ((v + 0.055) / 1.055).powf(2.4) + } +} + +/// Parses "#rrggbb" into a linear light tint that keeps roughly the same brightness as white. +fn parse_light_color(hex: &str) -> [f32; 3] { + let hex = hex.trim().trim_start_matches('#'); + if hex.len() != 6 || !hex.is_ascii() { + return [1.0; 3]; + } + + let mut color = [1.0f32; 3]; + for (i, c) in color.iter_mut().enumerate() { + match u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16) { + Ok(v) => *c = srgb_to_linear(v as f32 / 255.0), + Err(_) => return [1.0; 3], + } + } + + let luma = 0.2126 * color[0] + 0.7152 * color[1] + 0.0722 * color[2]; + let norm = luma.max(0.3); + color.map(|c| c / norm) +} + +fn read_params(adjustments: &serde_json::Value) -> RelightParams { + let num = |key: &str, default: f64| adjustments[key].as_f64().unwrap_or(default) as f32; + + let range = (num("relightRange", 60.0) / 100.0).clamp(0.0, 1.0); + + RelightParams { + light_x: num("relightX", 0.3), + light_y: num("relightY", 0.3), + light_z: 0.02 + (num("relightHeight", 30.0) / 100.0).clamp(0.0, 1.0), + radius: 0.1 + 1.9 * range * range, + light_color: parse_light_color(adjustments["relightColor"].as_str().unwrap_or("#ffffff")), + intensity: (num("relightIntensity", 50.0) / 100.0).clamp(0.0, 1.0) * 2.0, + ambient: (num("relightAmbient", 80.0) / 100.0).clamp(0.0, 1.0), + wrap: (num("relightSoftness", 30.0) / 100.0).clamp(0.0, 1.0), + specular: (num("relightSpecular", 0.0) / 100.0).clamp(0.0, 1.0), + } +} + +fn relight_in_place(image: &mut image::Rgb32FImage, normal_map: &RgbaImage, p: &RelightParams) { + let w = image.width() as usize; + let h = image.height() as usize; + let nw = normal_map.width() as usize; + let nh = normal_map.height() as usize; + let normals = normal_map.as_raw(); + + let x_scale = nw as f32 / w as f32; + let y_scale = nh as f32 / h as f32; + + let inv_long_side = 1.0 / w.max(h) as f32; + let light_px = p.light_x * w as f32; + let light_py = p.light_y * h as f32; + let inv_radius_sq = 1.0 / (p.radius * p.radius); + + image + .as_mut() + .par_chunks_exact_mut(w * 3) + .enumerate() + .for_each(|(y, row)| { + let fy = ((y as f32 + 0.5) * y_scale - 0.5).clamp(0.0, (nh - 1) as f32); + let y0 = fy as usize; + let y1 = (y0 + 1).min(nh - 1); + let ty = fy - y0 as f32; + + // Image rows grow downwards while the normal map uses Y up. + let ly = -(light_py - (y as f32 + 0.5)) * inv_long_side; + + for (x, px) in row.chunks_exact_mut(3).enumerate() { + let fx = ((x as f32 + 0.5) * x_scale - 0.5).clamp(0.0, (nw - 1) as f32); + let x0 = fx as usize; + let x1 = (x0 + 1).min(nw - 1); + let tx = fx - x0 as f32; + + let mut s = [0.0f32; 4]; + for (sx, sy, weight) in [ + (x0, y0, (1.0 - tx) * (1.0 - ty)), + (x1, y0, tx * (1.0 - ty)), + (x0, y1, (1.0 - tx) * ty), + (x1, y1, tx * ty), + ] { + let i = (sy * nw + sx) * 4; + for c in 0..4 { + s[c] += normals[i + c] as f32 * weight; + } + } + + let coverage = s[3] / 255.0; + if coverage <= 0.0 { + continue; + } + + let nx = s[0] / 127.5 - 1.0; + let ny = s[1] / 127.5 - 1.0; + let nz = s[2] / 127.5 - 1.0; + let inv_len = 1.0 / (nx * nx + ny * ny + nz * nz).sqrt().max(1e-6); + let (nx, ny, nz) = (nx * inv_len, ny * inv_len, nz * inv_len); + + let lx = (light_px - (x as f32 + 0.5)) * inv_long_side; + let planar_dist_sq = lx * lx + ly * ly; + let inv_l_len = 1.0 / (planar_dist_sq + p.light_z * p.light_z).sqrt(); + let (dx, dy, dz) = (lx * inv_l_len, ly * inv_l_len, p.light_z * inv_l_len); + + let attenuation = 1.0 / (1.0 + planar_dist_sq * inv_radius_sq); + + let n_dot_l = nx * dx + ny * dy + nz * dz; + let diffuse = ((n_dot_l + p.wrap) / (1.0 + p.wrap)).max(0.0); + + let spec = if p.specular > 0.0 && n_dot_l > 0.0 { + // Half vector between the light and the viewer (0, 0, 1). + let hz = dz + 1.0; + let inv_h_len = 1.0 / (dx * dx + dy * dy + hz * hz).sqrt().max(1e-6); + let n_dot_h = (nx * dx + ny * dy + nz * hz) * inv_h_len; + n_dot_h.max(0.0).powf(SPECULAR_POWER) * p.specular + } else { + 0.0 + }; + + let lit = p.intensity * attenuation * (diffuse + spec); + for c in 0..3 { + let gain = p.ambient + lit * p.light_color[c]; + px[c] *= 1.0 + (gain - 1.0) * coverage; + } + } + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use image::{Rgb, Rgb32FImage, Rgba}; + use serde_json::json; + + fn flat_normal_map(normal: [u8; 3], alpha: u8) -> RgbaImage { + RgbaImage::from_pixel(4, 4, Rgba([normal[0], normal[1], normal[2], alpha])) + } + + fn relit_pixel(adjustments: serde_json::Value, normal_map: &RgbaImage, x: u32) -> [f32; 3] { + let mut image = Rgb32FImage::from_pixel(8, 8, Rgb([0.5, 0.5, 0.5])); + relight_in_place(&mut image, normal_map, &read_params(&adjustments)); + image.get_pixel(x, 3).0 + } + + #[test] + fn surface_facing_the_light_gets_brighter_than_one_facing_away() { + let adjustments = json!({ + "relightX": 1.0, "relightY": 0.5, "relightHeight": 0, "relightSoftness": 0, "relightRange": 100, + }); + let facing_right = relit_pixel( + adjustments.clone(), + &flat_normal_map([255, 128, 128], 255), + 3, + ); + let facing_left = relit_pixel(adjustments, &flat_normal_map([0, 128, 128], 255), 3); + + assert!(facing_right[1] > 0.5); + assert!(facing_left[1] < 0.5); + } + + #[test] + fn light_above_the_pixel_brightens_upward_facing_surfaces() { + let adjustments = json!({ + "relightX": 0.5, "relightY": 0.0, "relightHeight": 0, "relightSoftness": 0, "relightRange": 100, + }); + let facing_up = relit_pixel( + adjustments.clone(), + &flat_normal_map([128, 255, 128], 255), + 3, + ); + let facing_down = relit_pixel(adjustments, &flat_normal_map([128, 0, 128], 255), 3); + + assert!(facing_up[1] > facing_down[1]); + } + + #[test] + fn light_falls_off_with_distance() { + let adjustments = + json!({ "relightX": 0.0, "relightY": 0.5, "relightHeight": 50, "relightRange": 20 }); + let normal_map = flat_normal_map([128, 128, 255], 255); + let near = relit_pixel(adjustments.clone(), &normal_map, 0); + let far = relit_pixel(adjustments, &normal_map, 7); + + assert!(near[1] > far[1]); + } + + #[test] + fn light_color_tints_the_lit_surface() { + let adjustments = json!({ "relightX": 0.5, "relightY": 0.5, "relightColor": "#ff8000" }); + let lit = relit_pixel(adjustments, &flat_normal_map([128, 128, 255], 255), 3); + + assert!(lit[0] > lit[1] && lit[1] > lit[2]); + } + + #[test] + fn invalid_light_color_falls_back_to_white() { + assert_eq!(parse_light_color("not-a-color"), [1.0; 3]); + assert!( + parse_light_color("#ffffff") + .iter() + .all(|c| (c - 1.0).abs() < 1e-4) + ); + } + + #[test] + fn pixels_without_geometry_are_left_untouched() { + let adjustments = json!({ "relightIntensity": 100, "relightAmbient": 0 }); + let lit = relit_pixel(adjustments, &flat_normal_map([128, 128, 255], 0), 3); + + assert_eq!(lit, [0.5; 3]); + } +} diff --git a/src/components/adjustments/Effects.tsx b/src/components/adjustments/Effects.tsx index a096065f05..2dd80afb17 100644 --- a/src/components/adjustments/Effects.tsx +++ b/src/components/adjustments/Effects.tsx @@ -7,7 +7,14 @@ import { motion } from 'framer-motion'; import clsx from 'clsx'; import Slider from '../ui/Slider'; import Switch from '../ui/Switch'; -import { Adjustments, Effect, CreativeAdjustment, getAdjustmentToolOrder, getHiddenAdjustmentTools } from '../../utils/adjustments'; +import Button from '../ui/Button'; +import { + Adjustments, + Effect, + CreativeAdjustment, + getAdjustmentToolOrder, + getHiddenAdjustmentTools, +} from '../../utils/adjustments'; import LUTControl from '../ui/LUTControl'; import { AppSettings } from '../ui/AppProperties'; import Text from '../ui/Text'; @@ -15,6 +22,7 @@ import AdjustmentSubSection from './AdjustmentSubSection'; import { TextVariants } from '../../types/typography'; import { DepthRangePicker } from '../ui/DepthRangePicker'; import { useProcessStore } from '../../store/useProcessStore'; +import { RELIGHT_TEMPERATURE_PRESETS, kelvinToHex } from '../../utils/relightUtils'; interface EffectsPanelProps { adjustments: Adjustments; @@ -148,6 +156,7 @@ export default function EffectsPanel({ }: EffectsPanelProps) { const { t } = useTranslation(); const [isGeneratingDepth, setIsGeneratingDepth] = useState(false); + const [isGeneratingNormal, setIsGeneratingNormal] = useState(false); const aiModelDownloadStatus = useProcessStore((state) => state.aiModelDownloadStatus); const handleGenerateLensBlurDepthMap = async () => { @@ -166,6 +175,22 @@ export default function EffectsPanel({ } }; + const handleGenerateRelightNormalMap = async () => { + setIsGeneratingNormal(true); + try { + const b64: string = await invoke('generate_full_image_normal_map', { jsAdjustments: adjustments }); + setAdjustments((prev: Partial) => ({ + ...prev, + relightNormalMap: b64, + })); + } catch (e: any) { + toast.error(`Failed to generate normal map: ${e}`); + setAdjustments((prev: Partial) => ({ ...prev, relightEnabled: false })); + } finally { + setIsGeneratingNormal(false); + } + }; + const handleAdjustmentChange = (key: string, value: any) => { const numericValue = typeof value === 'boolean' ? value : parseInt(value, 10); setAdjustments((prev: Partial) => ({ ...prev, [key]: numericValue })); @@ -194,6 +219,19 @@ export default function EffectsPanel({ } }; + const relightColor = (adjustments.relightColor || '#ffffff').toLowerCase(); + + const handleRelightColorChange = (color: string) => { + setAdjustments((prev: Partial) => ({ ...prev, relightColor: color.toLowerCase() })); + }; + + const handleRelightToggle = (enabled: boolean) => { + handleAdjustmentChange(Effect.RelightEnabled, enabled); + if (enabled && !adjustments.relightNormalMap) { + handleGenerateRelightNormalMap(); + } + }; + const hiddenTools = getHiddenAdjustmentTools(appSettings?.adjustmentLayout); const toolOrder = getAdjustmentToolOrder('effects', appSettings?.adjustmentLayout?.toolOrder); @@ -344,6 +382,178 @@ export default function EffectsPanel({ )} + {!hiddenTools.includes('relight') && ( + + + +
+
+
+ {isGeneratingNormal ? ( +
+
+ + + {aiModelDownloadStatus + ? t('editor.masks.settings.aiModelDownloading') + : t('editor.ai.generatingNormalMap')} + +
+ {aiModelDownloadStatus && ( + + {aiModelDownloadStatus} + + )} +
+ ) : ( + <> + handleAdjustmentChange(Effect.RelightIntensity, e.target.value)} + step={1} + value={adjustments.relightIntensity ?? 50} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightHeight, e.target.value)} + step={1} + value={adjustments.relightHeight ?? 30} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightRange, e.target.value)} + step={1} + value={adjustments.relightRange ?? 60} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightAmbient, e.target.value)} + step={1} + value={adjustments.relightAmbient ?? 80} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightSoftness, e.target.value)} + step={1} + value={adjustments.relightSoftness ?? 30} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightSpecular, e.target.value)} + step={1} + value={adjustments.relightSpecular ?? 0} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + +
+ + {t('adjustments.effects.relightColor')} + +
+ {RELIGHT_TEMPERATURE_PRESETS.map((kelvin) => { + const presetColor = kelvinToHex(kelvin); + return ( + + ); + })} +
+ +
+ + {adjustments.relightNormalMap && ( + {t('adjustments.effects.relightNormalMap')} + )} + + + + )} +
+
+
+ + )} + {!hiddenTools.includes('lut') && ( ) => { - wasPanningDisabledOnDown.current = isPanningDisabled; + // This runs in the capture phase, so the light handle cannot stop it by itself. + const isRelightHandle = !!(e.target as HTMLElement).closest('[data-relight-handle]'); + wasPanningDisabledOnDown.current = isPanningDisabled || isRelightHandle; + if (isRelightHandle) return; const isCropHandle = isCropping && e.button === 0 && !!(e.target as HTMLElement).closest('[data-ord]'); cropResizeStartRef.current = isCropHandle ? lastValidCropRef.current : null; diff --git a/src/components/panel/editor/ImageCanvas.tsx b/src/components/panel/editor/ImageCanvas.tsx index 8469a86bdd..0ef9cb8a79 100644 --- a/src/components/panel/editor/ImageCanvas.tsx +++ b/src/components/panel/editor/ImageCanvas.tsx @@ -12,10 +12,13 @@ import { RenderSize } from '../../../hooks/useImageRenderSize'; import { useOsPlatform } from '../../../hooks/useOsPlatform'; import { useTranslation } from 'react-i18next'; import { useEditorStore } from '../../../store/useEditorStore'; +import { useUIStore } from '../../../store/useUIStore'; import type { OverlayMode } from '../right/CropPanel'; import CompositionOverlays from './overlays/CompositionOverlays'; +import RelightHandle from './overlays/RelightHandle'; import { calculateStraightenAngle } from '../../../utils/cropUtils'; import { toast } from 'react-toastify'; +import { RelightGeometry } from '../../../utils/relightUtils'; interface CursorPreview { visible: boolean; @@ -1664,6 +1667,44 @@ const ImageCanvas = memo( : activeCrop.y : 0; + const isEffectsSectionOpen = useUIStore((state) => state.collapsibleSectionsState.effects); + const activeRotation = + liveRotation !== null && liveRotation !== undefined ? liveRotation : adjustments.rotation || 0; + const relightGeometry = useMemo( + () => ({ + imageWidth: selectedImage.width || 0, + imageHeight: selectedImage.height || 0, + orientationSteps: adjustments.orientationSteps || 0, + flipHorizontal: !!adjustments.flipHorizontal, + flipVertical: !!adjustments.flipVertical, + rotation: activeRotation, + cropX, + cropY, + }), + [ + selectedImage.width, + selectedImage.height, + adjustments.orientationSteps, + adjustments.flipHorizontal, + adjustments.flipVertical, + activeRotation, + cropX, + cropY, + ], + ); + const isRelightHandleVisible = + !!adjustments.relightEnabled && + !!adjustments.relightNormalMap && + adjustments.sectionVisibility?.effects !== false && + isEffectsSectionOpen && + !isCropping && + !isMasking && + !isAiEditing && + !isWbPickerActive && + !showOriginal && + relightGeometry.imageWidth > 0 && + imageRenderSize.width > 0; + const effectiveZoomScale = transformState.scale > 0 ? transformState.scale : 1; const brushStageSize = (brushSettings?.size ?? 0) / effectiveZoomScale; const brushImageSpaceSize = brushStageSize / (imageRenderSize.scale || 1); @@ -3154,6 +3195,16 @@ const ImageCanvas = memo(
+ {isRelightHandleVisible && ( + + )} + {!isDrawing.current && directPatchMarkers.map((m) => { const left = (m.cx - cropX) * imageRenderSize.scale + imageRenderSize.offsetX; diff --git a/src/components/panel/editor/overlays/RelightHandle.tsx b/src/components/panel/editor/overlays/RelightHandle.tsx new file mode 100644 index 0000000000..b6d322f700 --- /dev/null +++ b/src/components/panel/editor/overlays/RelightHandle.tsx @@ -0,0 +1,133 @@ +import { useCallback, useEffect, useRef, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { Adjustments, Coord } from '../../../../utils/adjustments'; +import { RenderSize } from '../../../../hooks/useImageRenderSize'; +import { RelightGeometry, relightDisplayToUv, relightUvToDisplay } from '../../../../utils/relightUtils'; + +interface RelightHandleProps { + adjustments: Adjustments; + geometry: RelightGeometry; + imageRenderSize: RenderSize; + inverseScale: number; + setAdjustments(fn: (prev: Adjustments) => Adjustments): void; +} + +const HANDLE_SIZE = 30; + +const clamp01 = (value: number) => Math.max(0, Math.min(1, value)); + +export default function RelightHandle({ + adjustments, + geometry, + imageRenderSize, + inverseScale, + setAdjustments, +}: RelightHandleProps) { + const { t } = useTranslation(); + const handleRef = useRef(null); + const frameRef = useRef(null); + const pendingUvRef = useRef(null); + const [dragUv, setDragUv] = useState(null); + + useEffect( + () => () => { + if (frameRef.current !== null) cancelAnimationFrame(frameRef.current); + }, + [], + ); + + const commitUv = useCallback( + (uv: Coord) => { + pendingUvRef.current = uv; + if (frameRef.current !== null) return; + frameRef.current = requestAnimationFrame(() => { + frameRef.current = null; + const next = pendingUvRef.current; + if (!next) return; + setAdjustments((prev: Adjustments) => ({ ...prev, relightX: next.x, relightY: next.y })); + }); + }, + [setAdjustments], + ); + + const pointerToUv = useCallback( + (clientX: number, clientY: number): Coord | null => { + const container = handleRef.current?.parentElement; + if (!container || !container.offsetWidth || !imageRenderSize.scale) return null; + + // The container sits inside the zoom/pan transform, so undo its on-screen scale. + const rect = container.getBoundingClientRect(); + const zoom = rect.width / container.offsetWidth || 1; + const localX = (clientX - rect.left) / zoom; + const localY = (clientY - rect.top) / zoom; + + const uv = relightDisplayToUv( + { + x: (localX - imageRenderSize.offsetX) / imageRenderSize.scale, + y: (localY - imageRenderSize.offsetY) / imageRenderSize.scale, + }, + geometry, + ); + return { x: clamp01(uv.x), y: clamp01(uv.y) }; + }, + [geometry, imageRenderSize], + ); + + const handlePointerDown = (e: React.PointerEvent) => { + e.stopPropagation(); + e.preventDefault(); + e.currentTarget.setPointerCapture(e.pointerId); + setDragUv({ x: adjustments.relightX ?? 0.3, y: adjustments.relightY ?? 0.3 }); + }; + + const handlePointerMove = (e: React.PointerEvent) => { + if (!dragUv) return; + e.stopPropagation(); + const uv = pointerToUv(e.clientX, e.clientY); + if (!uv) return; + setDragUv(uv); + commitUv(uv); + }; + + const handlePointerUp = (e: React.PointerEvent) => { + if (!dragUv) return; + e.stopPropagation(); + if (e.currentTarget.hasPointerCapture(e.pointerId)) { + e.currentTarget.releasePointerCapture(e.pointerId); + } + setDragUv(null); + }; + + const uv = dragUv ?? { x: adjustments.relightX ?? 0.3, y: adjustments.relightY ?? 0.3 }; + const display = relightUvToDisplay(uv, geometry); + const color = adjustments.relightColor || '#ffffff'; + + return ( +
e.stopPropagation()} + onTouchStart={(e) => e.stopPropagation()} + onClick={(e) => e.stopPropagation()} + style={{ + left: display.x * imageRenderSize.scale + imageRenderSize.offsetX, + top: display.y * imageRenderSize.scale + imageRenderSize.offsetY, + width: HANDLE_SIZE, + height: HANDLE_SIZE, + transform: `translate(-50%, -50%) scale(${inverseScale})`, + transformOrigin: 'center', + cursor: dragUv ? 'grabbing' : 'grab', + touchAction: 'none', + background: `radial-gradient(circle at 35% 30%, #ffffff 0%, ${color} 38%, color-mix(in srgb, ${color} 45%, #000000) 100%)`, + border: '1.5px solid rgba(255, 255, 255, 0.9)', + boxShadow: `0 0 0 1px rgba(0, 0, 0, 0.45), 0 0 14px 4px color-mix(in srgb, ${color} 55%, transparent), 0 2px 6px rgba(0, 0, 0, 0.5)`, + }} + /> + ); +} diff --git a/src/i18n/locales/ca.json b/src/i18n/locales/ca.json index befa7b8573..04b04396f7 100644 --- a/src/i18n/locales/ca.json +++ b/src/i18n/locales/ca.json @@ -127,6 +127,19 @@ "lightFlares": "Esclats de llum", "lut": "LUT", "midpoint": "Punt mitjà", + "relight": "Reil·luminar", + "relightAmbient": "Ambient", + "relightColor": "Color de la llum", + "relightCustomColor": "Color personalitzat", + "relightDragLight": "Arrossega per moure la llum", + "relightGenerateNormalMap": "Genera el mapa de normals", + "relightHeight": "Alçada", + "relightIntensity": "Intensitat", + "relightNormalMap": "Mapa de normals", + "relightRange": "Abast", + "relightRegenerateNormalMap": "Regenera el mapa de normals", + "relightSoftness": "Suavitat", + "relightSpecular": "Especular", "roughness": "Rugositat", "roundness": "Rodonesa", "size": "Mida", @@ -371,6 +384,7 @@ "editSettingsTitle": "Configuració d'edició", "editsTitle": "Edicions", "generatingDepthMap": "S'està calculant el mapa de profunditat d'IA...", + "generatingNormalMap": "S'està calculant el mapa de normals d'IA...", "generativeEditTitle": "Edició generativa", "inpaintingTitle": "Inpainting", "liquify": { diff --git a/src/i18n/locales/de.json b/src/i18n/locales/de.json index a05e62e2bc..96f18b648c 100644 --- a/src/i18n/locales/de.json +++ b/src/i18n/locales/de.json @@ -127,6 +127,19 @@ "lightFlares": "Blendenflecke", "lut": "LUT", "midpoint": "Mittelpunkt", + "relight": "Neu beleuchten", + "relightAmbient": "Umgebungslicht", + "relightColor": "Lichtfarbe", + "relightCustomColor": "Eigene Farbe", + "relightDragLight": "Ziehen, um das Licht zu bewegen", + "relightGenerateNormalMap": "Normal Map erzeugen", + "relightHeight": "Höhe", + "relightIntensity": "Intensität", + "relightNormalMap": "Normal Map", + "relightRange": "Reichweite", + "relightRegenerateNormalMap": "Normal Map neu erzeugen", + "relightSoftness": "Weichheit", + "relightSpecular": "Glanzlicht", "roughness": "Rauheit", "roundness": "Rundheit", "size": "Größe", @@ -371,6 +384,7 @@ "editSettingsTitle": "Bearbeitungseinstellungen", "editsTitle": "Bearbeitungen", "generatingDepthMap": "KI-Tiefenkarte wird berechnet...", + "generatingNormalMap": "KI-Normal-Map wird berechnet...", "generativeEditTitle": "Generative Bearbeitung", "inpaintingTitle": "Reparieren", "liquify": { diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index 9af1931ce1..61c086d95d 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -127,6 +127,19 @@ "lightFlares": "Light Flares", "lut": "LUT", "midpoint": "Midpoint", + "relight": "Relight", + "relightAmbient": "Ambient", + "relightColor": "Light Color", + "relightCustomColor": "Custom Color", + "relightDragLight": "Drag to move the light", + "relightGenerateNormalMap": "Generate Normal Map", + "relightHeight": "Height", + "relightIntensity": "Intensity", + "relightNormalMap": "Normal Map", + "relightRange": "Range", + "relightRegenerateNormalMap": "Regenerate Normal Map", + "relightSoftness": "Softness", + "relightSpecular": "Specular", "roughness": "Roughness", "roundness": "Roundness", "size": "Size", @@ -371,6 +384,7 @@ "editSettingsTitle": "Edit Settings", "editsTitle": "Edits", "generatingDepthMap": "Calculating AI Depth Map...", + "generatingNormalMap": "Calculating AI Normal Map...", "generativeEditTitle": "Generative Edit", "inpaintingTitle": "Inpainting", "liquify": { diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index 3d33088cbc..409c4ca440 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -127,6 +127,19 @@ "lightFlares": "Destellos de luz", "lut": "LUT", "midpoint": "Punto medio", + "relight": "Reiluminar", + "relightAmbient": "Ambiente", + "relightColor": "Color de la luz", + "relightCustomColor": "Color personalizado", + "relightDragLight": "Arrastra para mover la luz", + "relightGenerateNormalMap": "Generar mapa de normales", + "relightHeight": "Altura", + "relightIntensity": "Intensidad", + "relightNormalMap": "Mapa de normales", + "relightRange": "Alcance", + "relightRegenerateNormalMap": "Regenerar mapa de normales", + "relightSoftness": "Suavidad", + "relightSpecular": "Especular", "roughness": "Rugosidad", "roundness": "Redondez", "size": "Tamaño", @@ -387,6 +400,7 @@ "editSettingsTitle": "Ajustes de edición", "editsTitle": "Ediciones", "generatingDepthMap": "Calculando mapa de profundidad IA...", + "generatingNormalMap": "Calculando mapa de normales IA...", "generativeEditTitle": "Edición generativa", "inpaintingTitle": "Inpainting", "liquify": { diff --git a/src/i18n/locales/fr.json b/src/i18n/locales/fr.json index 57c7646980..6bc18d4420 100644 --- a/src/i18n/locales/fr.json +++ b/src/i18n/locales/fr.json @@ -127,6 +127,19 @@ "lightFlares": "Reflets", "lut": "LUT", "midpoint": "Point médian", + "relight": "Rééclairage", + "relightAmbient": "Ambiance", + "relightColor": "Couleur de la lumière", + "relightCustomColor": "Couleur personnalisée", + "relightDragLight": "Faites glisser pour déplacer la lumière", + "relightGenerateNormalMap": "Générer la carte de normales", + "relightHeight": "Hauteur", + "relightIntensity": "Intensité", + "relightNormalMap": "Carte de normales", + "relightRange": "Portée", + "relightRegenerateNormalMap": "Régénérer la carte de normales", + "relightSoftness": "Douceur", + "relightSpecular": "Spéculaire", "roughness": "Rugosité", "roundness": "Arrondi", "size": "Taille", @@ -387,6 +400,7 @@ "editSettingsTitle": "Paramètres de retouche", "editsTitle": "Retouches", "generatingDepthMap": "Calcul de la carte de profondeur IA...", + "generatingNormalMap": "Calcul de la carte de normales IA...", "generativeEditTitle": "Édition générative", "inpaintingTitle": "Inpainting", "liquify": { diff --git a/src/i18n/locales/it.json b/src/i18n/locales/it.json index ff0abe7434..6dd20314ba 100644 --- a/src/i18n/locales/it.json +++ b/src/i18n/locales/it.json @@ -127,6 +127,19 @@ "lightFlares": "Riflessi di Luce", "lut": "LUT", "midpoint": "Punto Medio", + "relight": "Reilluminazione", + "relightAmbient": "Ambiente", + "relightColor": "Colore della luce", + "relightCustomColor": "Colore personalizzato", + "relightDragLight": "Trascina per spostare la luce", + "relightGenerateNormalMap": "Genera mappa delle normali", + "relightHeight": "Altezza", + "relightIntensity": "Intensità", + "relightNormalMap": "Mappa delle normali", + "relightRange": "Portata", + "relightRegenerateNormalMap": "Rigenera mappa delle normali", + "relightSoftness": "Morbidezza", + "relightSpecular": "Speculare", "roughness": "Ruvidezza", "roundness": "Rotondità", "size": "Dimensione", @@ -387,6 +400,7 @@ "editSettingsTitle": "Impostazioni Modifica", "editsTitle": "Modifiche", "generatingDepthMap": "Calcolo mappa di profondità IA...", + "generatingNormalMap": "Calcolo mappa delle normali IA...", "generativeEditTitle": "Modifica generativa", "inpaintingTitle": "Inpainting", "liquify": { diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index 8ec890be1d..865c54a97e 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -127,6 +127,19 @@ "lightFlares": "ライトフレア", "lut": "LUT", "midpoint": "中間点", + "relight": "リライト", + "relightAmbient": "環境光", + "relightColor": "ライトの色", + "relightCustomColor": "カスタムカラー", + "relightDragLight": "ドラッグしてライトを移動", + "relightGenerateNormalMap": "法線マップを生成", + "relightHeight": "高さ", + "relightIntensity": "強度", + "relightNormalMap": "法線マップ", + "relightRange": "範囲", + "relightRegenerateNormalMap": "法線マップを再生成", + "relightSoftness": "柔らかさ", + "relightSpecular": "スペキュラ", "roughness": "粗さ", "roundness": "丸み", "size": "サイズ", @@ -371,6 +384,7 @@ "editSettingsTitle": "編集設定", "editsTitle": "編集一覧", "generatingDepthMap": "AI深度マップを計算中...", + "generatingNormalMap": "AI法線マップを計算中...", "generativeEditTitle": "生成編集", "inpaintingTitle": "インペインティング", "liquify": { diff --git a/src/i18n/locales/ko.json b/src/i18n/locales/ko.json index 2f33596513..a015755ace 100644 --- a/src/i18n/locales/ko.json +++ b/src/i18n/locales/ko.json @@ -127,6 +127,19 @@ "lightFlares": "빛 번짐", "lut": "LUT", "midpoint": "중간점", + "relight": "재조명", + "relightAmbient": "주변광", + "relightColor": "조명 색상", + "relightCustomColor": "사용자 지정 색상", + "relightDragLight": "드래그하여 조명 이동", + "relightGenerateNormalMap": "노멀 맵 생성", + "relightHeight": "높이", + "relightIntensity": "강도", + "relightNormalMap": "노멀 맵", + "relightRange": "범위", + "relightRegenerateNormalMap": "노멀 맵 다시 생성", + "relightSoftness": "부드러움", + "relightSpecular": "스페큘러", "roughness": "거칠기", "roundness": "둥글기", "size": "크기", @@ -371,6 +384,7 @@ "editSettingsTitle": "편집 설정", "editsTitle": "편집", "generatingDepthMap": "AI 심도 맵 계산 중...", + "generatingNormalMap": "AI 노멀 맵 계산 중...", "generativeEditTitle": "생성형 편집", "inpaintingTitle": "인페인팅", "liquify": { diff --git a/src/i18n/locales/pl.json b/src/i18n/locales/pl.json index 61dc217842..27b56705aa 100644 --- a/src/i18n/locales/pl.json +++ b/src/i18n/locales/pl.json @@ -127,6 +127,19 @@ "lightFlares": "Flary świetlne", "lut": "LUT", "midpoint": "Punkt środkowy", + "relight": "Ponowne oświetlenie", + "relightAmbient": "Światło otoczenia", + "relightColor": "Kolor światła", + "relightCustomColor": "Własny kolor", + "relightDragLight": "Przeciągnij, aby przesunąć światło", + "relightGenerateNormalMap": "Generuj mapę normalnych", + "relightHeight": "Wysokość", + "relightIntensity": "Intensywność", + "relightNormalMap": "Mapa normalnych", + "relightRange": "Zasięg", + "relightRegenerateNormalMap": "Generuj ponownie mapę normalnych", + "relightSoftness": "Miękkość", + "relightSpecular": "Odbicia", "roughness": "Szorstkość", "roundness": "Zaokrąglenie", "size": "Rozmiar", @@ -403,6 +416,7 @@ "editSettingsTitle": "Ustawienia edycji", "editsTitle": "Edycje", "generatingDepthMap": "Obliczanie mapy głębi AI...", + "generatingNormalMap": "Obliczanie mapy normalnych AI...", "generativeEditTitle": "Edycja generatywna", "inpaintingTitle": "Inpainting", "liquify": { diff --git a/src/i18n/locales/pt.json b/src/i18n/locales/pt.json index a792f3c55c..348edcfa74 100644 --- a/src/i18n/locales/pt.json +++ b/src/i18n/locales/pt.json @@ -127,6 +127,19 @@ "lightFlares": "Efeitos de Luz", "lut": "LUT", "midpoint": "Ponto Médio", + "relight": "Reiluminar", + "relightAmbient": "Ambiente", + "relightColor": "Cor da luz", + "relightCustomColor": "Cor personalizada", + "relightDragLight": "Arraste para mover a luz", + "relightGenerateNormalMap": "Gerar mapa de normais", + "relightHeight": "Altura", + "relightIntensity": "Intensidade", + "relightNormalMap": "Mapa de normais", + "relightRange": "Alcance", + "relightRegenerateNormalMap": "Regenerar mapa de normais", + "relightSoftness": "Suavidade", + "relightSpecular": "Especular", "roughness": "Rugosidade", "roundness": "Arredondamento", "size": "Tamanho", @@ -387,6 +400,7 @@ "editSettingsTitle": "Configurações de Edição", "editsTitle": "Edições", "generatingDepthMap": "A calcular mapa de profundidade IA...", + "generatingNormalMap": "A calcular mapa de normais IA...", "generativeEditTitle": "Edição Generativa", "inpaintingTitle": "Preenchimento", "liquify": { diff --git a/src/i18n/locales/ru.json b/src/i18n/locales/ru.json index 507056dcdb..622535d613 100644 --- a/src/i18n/locales/ru.json +++ b/src/i18n/locales/ru.json @@ -127,6 +127,19 @@ "lightFlares": "Блики", "lut": "LUT", "midpoint": "Средняя точка", + "relight": "Переосвещение", + "relightAmbient": "Окружающий свет", + "relightColor": "Цвет света", + "relightCustomColor": "Свой цвет", + "relightDragLight": "Перетащите, чтобы переместить свет", + "relightGenerateNormalMap": "Создать карту нормалей", + "relightHeight": "Высота", + "relightIntensity": "Интенсивность", + "relightNormalMap": "Карта нормалей", + "relightRange": "Дальность", + "relightRegenerateNormalMap": "Пересоздать карту нормалей", + "relightSoftness": "Мягкость", + "relightSpecular": "Блики", "roughness": "Шероховатость", "roundness": "Округлость", "size": "Размер", @@ -403,6 +416,7 @@ "editSettingsTitle": "Настройки редактирования", "editsTitle": "Правки", "generatingDepthMap": "Расчет карты глубины ИИ...", + "generatingNormalMap": "Расчет карты нормалей ИИ...", "generativeEditTitle": "Генеративное редактирование", "inpaintingTitle": "Восстановление", "liquify": { diff --git a/src/i18n/locales/zh-CN.json b/src/i18n/locales/zh-CN.json index d66e2c925f..1daac1eb23 100644 --- a/src/i18n/locales/zh-CN.json +++ b/src/i18n/locales/zh-CN.json @@ -127,6 +127,19 @@ "lightFlares": "光线耀斑", "lut": "LUT", "midpoint": "中点", + "relight": "重新打光", + "relightAmbient": "环境光", + "relightColor": "灯光颜色", + "relightCustomColor": "自定义颜色", + "relightDragLight": "拖动以移动灯光", + "relightGenerateNormalMap": "生成法线贴图", + "relightHeight": "高度", + "relightIntensity": "强度", + "relightNormalMap": "法线贴图", + "relightRange": "范围", + "relightRegenerateNormalMap": "重新生成法线贴图", + "relightSoftness": "柔和度", + "relightSpecular": "高光", "roughness": "粗糙度", "roundness": "圆度", "size": "大小", @@ -371,6 +384,7 @@ "editSettingsTitle": "编辑设置", "editsTitle": "编辑", "generatingDepthMap": "正在计算 AI 深度图...", + "generatingNormalMap": "正在计算 AI 法线贴图...", "generativeEditTitle": "生成式编辑", "inpaintingTitle": "图像修复", "liquify": { diff --git a/src/i18n/locales/zh-TW.json b/src/i18n/locales/zh-TW.json index c6e2f99e83..61db256dc5 100644 --- a/src/i18n/locales/zh-TW.json +++ b/src/i18n/locales/zh-TW.json @@ -127,6 +127,19 @@ "lightFlares": "光線耀斑", "lut": "LUT", "midpoint": "中點", + "relight": "重新打光", + "relightAmbient": "環境光", + "relightColor": "燈光顏色", + "relightCustomColor": "自訂顏色", + "relightDragLight": "拖曳以移動燈光", + "relightGenerateNormalMap": "產生法線貼圖", + "relightHeight": "高度", + "relightIntensity": "強度", + "relightNormalMap": "法線貼圖", + "relightRange": "範圍", + "relightRegenerateNormalMap": "重新產生法線貼圖", + "relightSoftness": "柔和度", + "relightSpecular": "高光", "roughness": "粗糙度", "roundness": "圓度", "size": "大小", @@ -371,6 +384,7 @@ "editSettingsTitle": "編輯設定", "editsTitle": "編輯", "generatingDepthMap": "正在計算 AI 深度圖...", + "generatingNormalMap": "正在計算 AI 法線貼圖...", "generativeEditTitle": "生成式編輯", "inpaintingTitle": "影像修復", "liquify": { diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index 9b8d5f5c0e..a61b646580 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -96,6 +96,17 @@ export enum Effect { LensBlurMaxDepth = 'lensBlurMaxDepth', LensBlurMinFade = 'lensBlurMinFade', LensBlurMaxFade = 'lensBlurMaxFade', + RelightEnabled = 'relightEnabled', + RelightIntensity = 'relightIntensity', + RelightX = 'relightX', + RelightY = 'relightY', + RelightHeight = 'relightHeight', + RelightRange = 'relightRange', + RelightAmbient = 'relightAmbient', + RelightSoftness = 'relightSoftness', + RelightSpecular = 'relightSpecular', + RelightColor = 'relightColor', + RelightNormalMap = 'relightNormalMap', } export enum CreativeAdjustment { @@ -243,6 +254,17 @@ export interface Adjustments { lutIsSceneReferred?: boolean; masks: Array; orientationSteps: number; + relightAmbient: number; + relightColor: string; + relightEnabled: boolean; + relightHeight: number; + relightIntensity: number; + relightNormalMap: string | null; + relightRange: number; + relightSoftness: number; + relightSpecular: number; + relightX: number; + relightY: number; rotation: number; saturation: number; sectionVisibility: SectionVisibility; @@ -581,6 +603,17 @@ export const INITIAL_ADJUSTMENTS: Adjustments = { lutIsSceneReferred: false, masks: [], orientationSteps: 0, + relightAmbient: 80, + relightColor: '#ffffff', + relightEnabled: false, + relightHeight: 30, + relightIntensity: 50, + relightNormalMap: null, + relightRange: 60, + relightSoftness: 30, + relightSpecular: 0, + relightX: 0.3, + relightY: 0.3, rotation: 0, saturation: 0, sectionVisibility: { @@ -728,6 +761,17 @@ export const normalizeLoadedAdjustments = (loadedAdjustments: Adjustments): any lensBlurMaxFade: loadedAdjustments.lensBlurMaxFade ?? INITIAL_ADJUSTMENTS.lensBlurMaxFade, lensBlurMinDepth: loadedAdjustments.lensBlurMinDepth ?? INITIAL_ADJUSTMENTS.lensBlurMinDepth, lensBlurMinFade: loadedAdjustments.lensBlurMinFade ?? INITIAL_ADJUSTMENTS.lensBlurMinFade, + relightAmbient: loadedAdjustments.relightAmbient ?? INITIAL_ADJUSTMENTS.relightAmbient, + relightColor: loadedAdjustments.relightColor ?? INITIAL_ADJUSTMENTS.relightColor, + relightEnabled: loadedAdjustments.relightEnabled ?? INITIAL_ADJUSTMENTS.relightEnabled, + relightHeight: loadedAdjustments.relightHeight ?? INITIAL_ADJUSTMENTS.relightHeight, + relightIntensity: loadedAdjustments.relightIntensity ?? INITIAL_ADJUSTMENTS.relightIntensity, + relightNormalMap: loadedAdjustments.relightNormalMap ?? INITIAL_ADJUSTMENTS.relightNormalMap, + relightRange: loadedAdjustments.relightRange ?? INITIAL_ADJUSTMENTS.relightRange, + relightSoftness: loadedAdjustments.relightSoftness ?? INITIAL_ADJUSTMENTS.relightSoftness, + relightSpecular: loadedAdjustments.relightSpecular ?? INITIAL_ADJUSTMENTS.relightSpecular, + relightX: loadedAdjustments.relightX ?? INITIAL_ADJUSTMENTS.relightX, + relightY: loadedAdjustments.relightY ?? INITIAL_ADJUSTMENTS.relightY, lensCorrectionMode: loadedAdjustments.lensCorrectionMode || 'manual', lensMaker: loadedAdjustments.lensMaker ?? INITIAL_ADJUSTMENTS.lensMaker, lensModel: loadedAdjustments.lensModel ?? INITIAL_ADJUSTMENTS.lensModel, @@ -952,6 +996,17 @@ export const ADJUSTMENT_SECTIONS: Sections = { Effect.LensBlurMaxDepth, Effect.LensBlurMinFade, Effect.LensBlurMaxFade, + Effect.RelightEnabled, + Effect.RelightIntensity, + Effect.RelightX, + Effect.RelightY, + Effect.RelightHeight, + Effect.RelightRange, + Effect.RelightAmbient, + Effect.RelightSoftness, + Effect.RelightSpecular, + Effect.RelightColor, + Effect.RelightNormalMap, ], }; @@ -999,6 +1054,7 @@ export const ADJUSTMENT_SECTION_TOOLS: Record + Math.round(Math.max(0, Math.min(255, value))) + .toString(16) + .padStart(2, '0'); + +// Tanner Helland's black body approximation, good enough for picking a light tint. +export const kelvinToHex = (kelvin: number): string => { + const temp = Math.max(1000, Math.min(40000, kelvin)) / 100; + + const red = temp <= 66 ? 255 : 329.698727446 * Math.pow(temp - 60, -0.1332047592); + const green = + temp <= 66 ? 99.4708025861 * Math.log(temp) - 161.1195681661 : 288.1221695283 * Math.pow(temp - 60, -0.0755148492); + const blue = temp >= 66 ? 255 : temp <= 19 ? 0 : 138.5177312231 * Math.log(temp - 10) - 305.0447927307; + + return `#${toHexChannel(red)}${toHexChannel(green)}${toHexChannel(blue)}`; +}; + +export interface RelightGeometry { + imageWidth: number; + imageHeight: number; + orientationSteps: number; + flipHorizontal: boolean; + flipVertical: boolean; + rotation: number; + cropX: number; + cropY: number; +} + +const orientPoint = (x: number, y: number, w: number, h: number, steps: number): Coord => { + const s = ((steps % 4) + 4) % 4; + if (s === 0) return { x, y }; + if (s === 1) return { x: h - y, y: x }; + if (s === 2) return { x: w - x, y: h - y }; + return { x: y, y: w - x }; +}; + +const getOrientedSize = (g: RelightGeometry) => { + const swapped = (((g.orientationSteps % 4) + 4) % 4) % 2 !== 0; + return { + width: swapped ? g.imageHeight : g.imageWidth, + height: swapped ? g.imageWidth : g.imageHeight, + }; +}; + +const rotateAboutCenter = (x: number, y: number, width: number, height: number, degrees: number): Coord => { + if (Math.abs(degrees) < 1e-4) return { x, y }; + const rad = (degrees * Math.PI) / 180; + const cos = Math.cos(rad); + const sin = Math.sin(rad); + const cx = width / 2; + const cy = height / 2; + const dx = x - cx; + const dy = y - cy; + return { x: cx + dx * cos - dy * sin, y: cy + dx * sin + dy * cos }; +}; + +/** + * The light position is stored as a fraction of the image before orientation, flip, + * rotation and crop, because that is where the relight pass runs. These helpers convert + * between that space and pixels of the cropped image shown in the editor. + */ +export const relightUvToDisplay = (uv: Coord, g: RelightGeometry): Coord => { + const oriented = getOrientedSize(g); + let { x, y } = orientPoint( + uv.x * g.imageWidth, + uv.y * g.imageHeight, + g.imageWidth, + g.imageHeight, + g.orientationSteps, + ); + if (g.flipHorizontal) x = oriented.width - x; + if (g.flipVertical) y = oriented.height - y; + const rotated = rotateAboutCenter(x, y, oriented.width, oriented.height, g.rotation); + return { x: rotated.x - g.cropX, y: rotated.y - g.cropY }; +}; + +export const relightDisplayToUv = (point: Coord, g: RelightGeometry): Coord => { + const oriented = getOrientedSize(g); + let { x, y } = rotateAboutCenter(point.x + g.cropX, point.y + g.cropY, oriented.width, oriented.height, -g.rotation); + if (g.flipHorizontal) x = oriented.width - x; + if (g.flipVertical) y = oriented.height - y; + const original = orientPoint(x, y, oriented.width, oriented.height, 4 - (((g.orientationSteps % 4) + 4) % 4)); + return { + x: g.imageWidth > 0 ? original.x / g.imageWidth : 0, + y: g.imageHeight > 0 ? original.y / g.imageHeight : 0, + }; +}; From 76ead1ff8598db7806719852c120063a3fa46590 Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Mon, 5 Oct 2026 15:52:59 -0500 Subject: [PATCH 05/16] feat: add lens blur and relight image processing features with localized UI components --- src-tauri/src/ai_commands.rs | 35 +- src-tauri/src/ai_processing.rs | 143 ++++- src-tauri/src/cache_utils.rs | 4 +- src-tauri/src/lens_blur.rs | 57 +- src-tauri/src/lib.rs | 2 +- src-tauri/src/relight.rs | 528 +++++++++++++++--- src/components/adjustments/Effects.tsx | 19 +- .../panel/editor/overlays/RelightHandle.tsx | 29 +- src/i18n/locales/ca.json | 4 +- src/i18n/locales/de.json | 4 +- src/i18n/locales/en.json | 4 +- src/i18n/locales/es.json | 4 +- src/i18n/locales/fr.json | 4 +- src/i18n/locales/it.json | 4 +- src/i18n/locales/ja.json | 4 +- src/i18n/locales/ko.json | 4 +- src/i18n/locales/pl.json | 4 +- src/i18n/locales/pt.json | 4 +- src/i18n/locales/ru.json | 4 +- src/i18n/locales/zh-CN.json | 4 +- src/i18n/locales/zh-TW.json | 4 +- src/utils/adjustments.ts | 20 +- 22 files changed, 750 insertions(+), 139 deletions(-) diff --git a/src-tauri/src/ai_commands.rs b/src-tauri/src/ai_commands.rs index 493902767b..54d4ccbf7b 100644 --- a/src-tauri/src/ai_commands.rs +++ b/src-tauri/src/ai_commands.rs @@ -301,12 +301,29 @@ pub async fn generate_full_image_depth_map( Ok(format!("data:image/png;base64,{}", base64_str)) } +#[derive(serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct RelightMapsPayload { + normal_map: String, + depth_map: String, + depth_scale: f32, +} + +fn encode_png_data_url(image: &image::DynamicImage) -> Result { + let mut buf = Cursor::new(Vec::new()); + image + .write_to(&mut buf, ImageFormat::Png) + .map_err(|e| e.to_string())?; + let base64_str = general_purpose::STANDARD.encode(buf.get_ref()); + Ok(format!("data:image/png;base64,{}", base64_str)) +} + #[tauri::command] -pub async fn generate_full_image_normal_map( +pub async fn generate_relight_maps( js_adjustments: serde_json::Value, state: tauri::State<'_, AppState>, app_handle: tauri::AppHandle, -) -> Result { +) -> Result { let normal_model = crate::ai_processing::get_or_init_normal_model( &app_handle, &state.ai_state, @@ -317,16 +334,14 @@ pub async fn generate_full_image_normal_map( let warped_image = crate::get_cached_full_warped_image(&state, &js_adjustments)?; - let normal_img = crate::ai_processing::run_normal_model(warped_image.as_ref(), &normal_model) + let maps = crate::ai_processing::run_normal_model(warped_image.as_ref(), &normal_model) .map_err(|e| e.to_string())?; - let mut buf = std::io::Cursor::new(Vec::new()); - normal_img - .write_to(&mut buf, image::ImageFormat::Png) - .map_err(|e| e.to_string())?; - let base64_str = base64::engine::general_purpose::STANDARD.encode(buf.get_ref()); - - Ok(format!("data:image/png;base64,{}", base64_str)) + Ok(RelightMapsPayload { + normal_map: encode_png_data_url(&image::DynamicImage::ImageRgba8(maps.normal))?, + depth_map: encode_png_data_url(&image::DynamicImage::ImageLuma16(maps.depth))?, + depth_scale: maps.depth_scale, + }) } #[allow(clippy::too_many_arguments)] diff --git a/src-tauri/src/ai_processing.rs b/src-tauri/src/ai_processing.rs index a69ec98c1d..b2016324fc 100644 --- a/src-tauri/src/ai_processing.rs +++ b/src-tauri/src/ai_processing.rs @@ -1756,13 +1756,83 @@ pub fn run_depth_anything_model( Ok(depth_map) } -/// Runs MoGe-2 and returns a camera-space normal map: -/// RGB = normal (X right, Y up, Z towards the viewer) packed as n * 0.5 + 0.5, -/// A = surface validity (0 where there is no geometry, e.g. sky). +pub struct RelightMaps { + /// RGB = camera-space normal (X right, Y up, Z towards the viewer) packed as n * 0.5 + 0.5, + /// A = surface validity (0 where there is no geometry, e.g. sky). + pub normal: RgbaImage, + /// Distance from the nearest surface, 0 = nearest and 65535 = farthest. + pub depth: ImageBuffer, Vec>, + /// Depth covered by the full range of `depth`, in units of the longest image side. + pub depth_scale: f32, +} + +/// Turns the model's camera-space point map into a depth map whose scale matches +/// the image plane, so a light can be placed in front of or behind the subject. +fn build_depth_map(points: &[f32], mask: &[f32], w: usize, h: usize) -> (Vec, f32) { + let long_side = w.max(h) as f32; + + let mut depths = Vec::with_capacity(w * h); + // Least squares fit of X = s * Z * x_norm, where s is the sensor size over the focal length. + let mut fit_num = 0.0f64; + let mut fit_den = 0.0f64; + + for (i, (point, m)) in points.chunks_exact(3).zip(mask).enumerate() { + let z = point[2]; + if *m < 0.5 || !z.is_finite() || z <= 0.0 || !point[0].is_finite() || !point[1].is_finite() + { + continue; + } + depths.push(z); + + let xn = ((i % w) as f32 + 0.5 - w as f32 * 0.5) / long_side; + let yn = ((i / w) as f32 + 0.5 - h as f32 * 0.5) / long_side; + fit_num += (point[0] * z * xn + point[1] * z * yn) as f64; + fit_den += ((z * xn) * (z * xn) + (z * yn) * (z * yn)) as f64; + } + + if depths.len() < 16 { + return (vec![0; w * h], 0.0); + } + + depths.sort_unstable_by(|a, b| a.total_cmp(b)); + let percentile = |p: f32| depths[((depths.len() - 1) as f32 * p) as usize]; + let near = percentile(0.005); + let far = percentile(0.995); + let reference = percentile(0.5); + let range = far - near; + + if range <= 1e-6 || reference <= 1e-6 { + return (vec![0; w * h], 0.0); + } + + let sensor_over_focal = (fit_num / fit_den.max(1e-12)) as f32; + let sensor_over_focal = if sensor_over_focal.is_finite() && sensor_over_focal > 0.05 { + sensor_over_focal + } else { + 1.0 + }; + + let depth_data = points + .chunks_exact(3) + .zip(mask) + .map(|(point, m)| { + let z = point[2]; + if *m < 0.5 || !z.is_finite() { + u16::MAX + } else { + (((z - near) / range).clamp(0.0, 1.0) * 65535.0).round() as u16 + } + }) + .collect(); + + (depth_data, range / (sensor_over_focal * reference)) +} + +/// Runs MoGe-2 and returns the normal and depth maps used for relighting. pub fn run_normal_model( image: &DynamicImage, normal_session: &Mutex, -) -> Result { +) -> Result { let resized_image = if image.width().max(image.height()) > NORMAL_MAX_SIZE { image.resize(NORMAL_MAX_SIZE, NORMAL_MAX_SIZE, FilterType::Triangle) } else { @@ -1796,12 +1866,17 @@ pub fn run_normal_model( let normal_tensor = outputs["normal"].try_extract_array::()?; let mask_tensor = outputs["mask"].try_extract_array::()?; + let points_tensor = outputs["points"].try_extract_array::()?; - if normal_tensor.shape() != [1, rh, rw, 3] || mask_tensor.shape() != [1, rh, rw] { + if normal_tensor.shape() != [1, rh, rw, 3] + || mask_tensor.shape() != [1, rh, rw] + || points_tensor.shape() != [1, rh, rw, 3] + { return Err(anyhow::anyhow!( - "Unexpected normal model output shape: {:?} / {:?}", + "Unexpected normal model output shape: {:?} / {:?} / {:?}", normal_tensor.shape(), - mask_tensor.shape() + mask_tensor.shape(), + points_tensor.shape() )); } @@ -1825,8 +1900,19 @@ pub fn run_normal_model( normal_data.push((m.clamp(0.0, 1.0) * 255.0).round() as u8); } - RgbaImage::from_raw(resized_w, resized_h, normal_data) - .ok_or_else(|| anyhow::anyhow!("Failed to create normal map from model output")) + let normal = RgbaImage::from_raw(resized_w, resized_h, normal_data) + .ok_or_else(|| anyhow::anyhow!("Failed to create normal map from model output"))?; + + let points = points_tensor.as_standard_layout(); + let (depth_data, depth_scale) = build_depth_map(points.as_slice().unwrap(), mask, rw, rh); + let depth = ImageBuffer::from_raw(resized_w, resized_h, depth_data) + .ok_or_else(|| anyhow::anyhow!("Failed to create depth map from model output"))?; + + Ok(RelightMaps { + normal, + depth, + depth_scale, + }) } #[derive(Serialize, Deserialize, Debug, Clone, Default)] @@ -1902,3 +1988,42 @@ pub struct AiDepthMaskParameters { #[serde(default)] pub orientation_steps: Option, } + +#[cfg(test)] +mod relight_map_tests { + use super::build_depth_map; + + #[test] + fn depth_map_spans_nearest_to_farthest_and_matches_image_scale() { + // A 32x32 view with a 90 degree field of view (sensor / focal = 2) of a plane + // tilted so that depth grows from 2 on the left to 4 on the right. + let (w, h) = (32usize, 32usize); + let mut points = Vec::new(); + for y in 0..h { + for x in 0..w { + let z = 2.0 + 2.0 * x as f32 / (w - 1) as f32; + let xn = (x as f32 + 0.5 - 16.0) / 32.0; + let yn = (y as f32 + 0.5 - 16.0) / 32.0; + points.extend_from_slice(&[2.0 * z * xn, 2.0 * z * yn, z]); + } + } + let mask = vec![1.0f32; w * h]; + + let (depth, scale) = build_depth_map(&points, &mask, w, h); + + assert!(depth[0] < 1000); + assert!(depth[w - 1] > 64000); + assert!(depth[w / 2] > depth[w / 4]); + // Depth range ~2 over an image that is 2 * 3 = 6 units wide at the median depth. + assert!((scale - 2.0 / 6.0).abs() < 0.03, "scale was {scale}"); + } + + #[test] + fn image_without_geometry_has_no_depth() { + let points = [[0.0f32, 0.0, 1.0]; 64].concat(); + let (depth, scale) = build_depth_map(&points, &[0.0; 64], 8, 8); + + assert_eq!(scale, 0.0); + assert!(depth.iter().all(|d| *d == 0)); + } +} diff --git a/src-tauri/src/cache_utils.rs b/src-tauri/src/cache_utils.rs index f16bde67b2..b31974cf08 100644 --- a/src-tauri/src/cache_utils.rs +++ b/src-tauri/src/cache_utils.rs @@ -96,13 +96,15 @@ fn hash_relight( "relightIntensity", "relightX", "relightY", - "relightHeight", + "relightDepth", "relightAmbient", "relightSoftness", "relightSpecular", "relightRange", "relightColor", "relightNormalMap", + "relightDepthMap", + "relightDepthScale", ]; for key in relight_keys { diff --git a/src-tauri/src/lens_blur.rs b/src-tauri/src/lens_blur.rs index 6a411752c7..ca8e4b6afe 100644 --- a/src-tauri/src/lens_blur.rs +++ b/src-tauri/src/lens_blur.rs @@ -134,15 +134,16 @@ fn depth_to_signed_coc( } #[allow(clippy::too_many_arguments)] -fn build_coc_field( - src: &Rgb32FImage, - depth: &image::GrayImage, - min_depth: f32, - max_depth: f32, - min_fade: f32, - max_fade: f32, - max_radius: f32, -) -> Vec { +/// Edge guide for joint upsampling: a lightly smoothed luma plane at full resolution +/// plus a copy downscaled so that its longest side is at most 1024 px. +pub(crate) struct LumaGuide { + pub luma_full: Vec, + pub guide: Vec, + pub gw: usize, + pub gh: usize, +} + +pub(crate) fn build_luma_guide(src: &Rgb32FImage) -> LumaGuide { let w = src.width() as usize; let h = src.height() as usize; let raw = src.as_raw(); @@ -196,6 +197,33 @@ fn build_coc_field( } }); + LumaGuide { + luma_full, + guide, + gw, + gh, + } +} + +fn build_coc_field( + src: &Rgb32FImage, + depth: &image::GrayImage, + min_depth: f32, + max_depth: f32, + min_fade: f32, + max_fade: f32, + max_radius: f32, +) -> Vec { + let w = src.width() as usize; + let h = src.height() as usize; + + let LumaGuide { + luma_full, + guide, + gw, + gh, + } = build_luma_guide(src); + let depth_small = dof_depth_to_f32(depth, gw, gh); let radius = ((gw.max(gh) as f32) * 0.015).round().max(2.0) as usize; let packed = build_guided_model(&guide, &depth_small, gw, gh, radius); @@ -243,7 +271,14 @@ fn build_coc_field( coc } -fn build_guided_model(guide: &[f32], p: &[f32], w: usize, h: usize, radius: usize) -> Vec { +/// Returns, per pixel, the local linear model `a * guide + b` and the `[lo, hi]` band it must stay in. +pub(crate) fn build_guided_model( + guide: &[f32], + p: &[f32], + w: usize, + h: usize, + radius: usize, +) -> Vec { let n = w * h; let eps = 4.0e-3f32; @@ -881,7 +916,7 @@ fn blur_layer_bokeh( out } -fn dof_box_filter(buf: &mut [f32], w: usize, h: usize, ch: usize, radius: usize) { +pub(crate) fn dof_box_filter(buf: &mut [f32], w: usize, h: usize, ch: usize, radius: usize) { if radius == 0 || w < 2 || h < 2 { return; } diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 31c203ccab..1d94907125 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2180,7 +2180,7 @@ pub fn run() { ai_commands::check_ai_connector_status, ai_commands::test_ai_connector_connection, ai_commands::generate_full_image_depth_map, - ai_commands::generate_full_image_normal_map, + ai_commands::generate_relight_maps, ai_commands::cancel_ai_task, apple_raw::is_raw9_available, inpainting::invoke_generative_replace_with_mask_def, diff --git a/src-tauri/src/relight.rs b/src-tauri/src/relight.rs index 482b27664f..67ed213a42 100644 --- a/src-tauri/src/relight.rs +++ b/src-tauri/src/relight.rs @@ -1,16 +1,40 @@ use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; -use image::{DynamicImage, GenericImageView, RgbaImage}; +use image::{DynamicImage, GenericImageView, ImageBuffer, Luma, Rgb32FImage, RgbaImage}; use rayon::prelude::*; use std::borrow::Cow; +use std::sync::{Arc, Mutex}; + +use crate::lens_blur::{LumaGuide, build_guided_model, build_luma_guide, dof_box_filter}; const SPECULAR_POWER: f32 = 24.0; +// How far in front of the nearest surface the light can be pulled, in units of the longest image side. +const MAX_FRONT_DISTANCE: f32 = 1.5; + +// Below this size there is nothing to gain from snapping the maps to the image edges. +const MIN_REFINE_SIZE: u32 = 64; + +type DepthMap = ImageBuffer, Vec>; + +/// Normal, coverage and depth for every pixel of the image being relit. +struct SurfaceMaps { + normals: Vec<[i8; 3]>, + coverage: Vec, + depth: Vec, +} + +// Building the maps only depends on the photo and the AI maps, never on the light, so the +// last result is kept around while the light is being moved. +static SURFACE_MAPS_CACHE: Mutex)>> = Mutex::new(None); struct RelightParams { // Light position as a fraction of the image width / height. light_x: f32, light_y: f32, - // Distance of the light from the image plane, in units of the longest image side. + // Light depth in units of the longest image side: 0 is the nearest surface, + // positive values move towards the viewer and negative values go into the scene. light_z: f32, + // Depth covered by the full range of the depth map, in the same units. + depth_scale: f32, // Falloff radius, in units of the longest image side. radius: f32, light_color: [f32; 3], @@ -39,38 +63,325 @@ pub fn apply_relight<'a>( return image; } - let b64_data = match normal_b64.find(',') { - Some(idx) => &normal_b64[idx + 1..], - None => normal_b64, - }; - let decoded = match BASE64.decode(b64_data) { - Ok(b) => b, - Err(_) => return image, - }; - let normal_map = match image::load_from_memory(&decoded) { - Ok(img) => img.into_rgba8(), - Err(_) => return image, - }; - if normal_map.width() < 2 || normal_map.height() < 2 { - return image; - } - let (w, h) = image.dimensions(); if w < 2 || h < 2 { return image; } - let params = read_params(adjustments); - let start = std::time::Instant::now(); let mut out = image.as_ref().to_rgb32f(); - relight_in_place(&mut out, &normal_map, ¶ms); + + let depth_b64 = adjustments["relightDepthMap"].as_str().unwrap_or(""); + let depth_scale = adjustments["relightDepthScale"].as_f64().unwrap_or(0.0) as f32; + let cache_key = surface_maps_key(&out, normal_b64, depth_b64, depth_scale); + + let cached = SURFACE_MAPS_CACHE + .lock() + .unwrap_or_else(|e| e.into_inner()) + .as_ref() + .filter(|(key, _)| *key == cache_key) + .map(|(_, maps)| Arc::clone(maps)); + + let maps = match cached { + Some(maps) => maps, + None => { + let Some(normal_map) = decode_data_url(normal_b64).map(|img| img.into_rgba8()) else { + return image; + }; + if normal_map.width() < 2 || normal_map.height() < 2 { + return image; + } + let depth_map = decode_data_url(depth_b64) + .map(|img| img.into_luma16()) + .filter(|depth| depth.dimensions() == normal_map.dimensions()); + + let refine = w >= MIN_REFINE_SIZE && h >= MIN_REFINE_SIZE; + let maps = Arc::new(build_surface_maps( + &out, + &normal_map, + depth_map.as_ref(), + depth_scale, + refine, + )); + log::info!( + "relight maps ({}x{}) built in {:.2?}", + w, + h, + start.elapsed() + ); + + *SURFACE_MAPS_CACHE.lock().unwrap_or_else(|e| e.into_inner()) = + Some((cache_key, Arc::clone(&maps))); + maps + } + }; + + let params = read_params(adjustments, !depth_b64.is_empty()); + relight_in_place(&mut out, &maps, ¶ms); log::info!("relight ({}x{}) took {:.2?}", w, h, start.elapsed()); Cow::Owned(DynamicImage::ImageRgb32F(out)) } +fn decode_data_url(data_url: &str) -> Option { + if data_url.is_empty() { + return None; + } + let b64_data = match data_url.find(',') { + Some(idx) => &data_url[idx + 1..], + None => data_url, + }; + let decoded = BASE64.decode(b64_data).ok()?; + image::load_from_memory(&decoded).ok() +} + +fn surface_maps_key( + image: &Rgb32FImage, + normal_b64: &str, + depth_b64: &str, + depth_scale: f32, +) -> blake3::Hash { + let mut hasher = blake3::Hasher::new(); + hasher.update(&image.width().to_le_bytes()); + hasher.update(&image.height().to_le_bytes()); + hasher.update(&depth_scale.to_le_bytes()); + hasher.update(&(normal_b64.len() as u64).to_le_bytes()); + hasher.update(normal_b64.as_bytes()); + hasher.update(depth_b64.as_bytes()); + + // A sparse sample is enough to notice that the photo itself changed. + let raw = image.as_raw(); + let step = (raw.len() / 4096).max(1); + for value in raw.iter().step_by(step) { + hasher.update(&value.to_le_bytes()); + } + + hasher.finalize() +} + +fn smoothstep(e0: f32, e1: f32, x: f32) -> f32 { + let t = ((x - e0) / (e1 - e0)).clamp(0.0, 1.0); + t * t * (3.0 - 2.0 * t) +} + +/// Position of a destination pixel inside a smaller source grid, for bilinear sampling. +struct Tap { + i0: usize, + i1: usize, + t: f32, +} + +fn taps(dst: usize, src: usize) -> Vec { + let scale = src as f32 / dst as f32; + (0..dst) + .map(|d| { + let f = ((d as f32 + 0.5) * scale - 0.5).clamp(0.0, (src - 1) as f32); + let i0 = f as usize; + Tap { + i0, + i1: (i0 + 1).min(src - 1), + t: f - i0 as f32, + } + }) + .collect() +} + +fn resample_plane(src: &[f32], sw: usize, sh: usize, dw: usize, dh: usize) -> Vec { + if sw == dw && sh == dh { + return src.to_vec(); + } + let x_taps = taps(dw, sw); + let y_taps = taps(dh, sh); + let mut out = vec![0.0f32; dw * dh]; + out.par_chunks_exact_mut(dw) + .enumerate() + .for_each(|(y, row)| { + let ty = &y_taps[y]; + for (x, value) in row.iter_mut().enumerate() { + let tx = &x_taps[x]; + let top = src[ty.i0 * sw + tx.i0] * (1.0 - tx.t) + src[ty.i0 * sw + tx.i1] * tx.t; + let bot = src[ty.i1 * sw + tx.i0] * (1.0 - tx.t) + src[ty.i1 * sw + tx.i1] * tx.t; + *value = top + (bot - top) * ty.t; + } + }); + out +} + +/// Local standard deviation of a plane, used to find where it has a hard edge. +fn local_deviation(plane: &[f32], w: usize, h: usize, radius: usize) -> Vec { + let mut mean = plane.to_vec(); + dof_box_filter(&mut mean, w, h, 1, radius); + let mut mean_sq: Vec = plane.iter().map(|v| v * v).collect(); + dof_box_filter(&mut mean_sq, w, h, 1, radius); + mean.iter() + .zip(&mean_sq) + .map(|(m, sq)| (sq - m * m).max(0.0).sqrt()) + .collect() +} + +/// Scales the AI maps up to the image. With `refine`, silhouettes (hair, object outlines) +/// are snapped to the edges of the photo instead of being left as a soft blur, while the +/// inside of each surface keeps the detail the model predicted. +fn build_surface_maps( + image: &Rgb32FImage, + normal_map: &RgbaImage, + depth_map: Option<&DepthMap>, + depth_scale: f32, + refine: bool, +) -> SurfaceMaps { + let w = image.width() as usize; + let h = image.height() as usize; + let nw = normal_map.width() as usize; + let nh = normal_map.height() as usize; + + // Planes 0-2: normal packed to 0..1, 3: coverage, 4: depth. + let normal_raw = normal_map.as_raw(); + let mut planes: Vec> = (0..4) + .map(|c| { + normal_raw + .iter() + .skip(c) + .step_by(4) + .map(|v| *v as f32 / 255.0) + .collect() + }) + .collect(); + planes.push(match depth_map { + Some(depth) => depth.as_raw().iter().map(|v| *v as f32 / 65535.0).collect(), + None => vec![0.0; nw * nh], + }); + + let (gw, gh, models, edge_weight, luma_full) = if refine { + let LumaGuide { + luma_full, + guide, + gw, + gh, + } = build_luma_guide(image); + + planes = planes + .iter() + .map(|plane| resample_plane(plane, nw, nh, gw, gh)) + .collect(); + + let radius = ((gw.max(gh) as f32) * 0.015).round().max(2.0) as usize; + + // Only trust the photo's edges where the geometry really breaks: a depth jump or the + // border of the valid area. Elsewhere the filter would just blur the normals. + let coverage_dev = local_deviation(&planes[3], gw, gh, radius); + let has_depth = depth_map.is_some() && depth_scale > 0.0; + let other_dev = if has_depth { + local_deviation(&planes[4], gw, gh, radius) + } else { + local_deviation(&planes[2], gw, gh, radius) + }; + let mut edge_weight: Vec = coverage_dev + .iter() + .zip(&other_dev) + .map(|(cov, other)| { + let geometry = if has_depth { + smoothstep(0.01, 0.06, other * depth_scale) + } else { + smoothstep(0.05, 0.2, *other) + }; + smoothstep(0.05, 0.25, *cov).max(geometry) + }) + .collect(); + dof_box_filter(&mut edge_weight, gw, gh, 1, (radius / 2).max(1)); + + let models: Vec> = planes + .par_iter() + .map(|plane| build_guided_model(&guide, plane, gw, gh, radius)) + .collect(); + + (gw, gh, models, edge_weight, luma_full) + } else { + (nw, nh, Vec::new(), Vec::new(), Vec::new()) + }; + + let x_taps = taps(w, gw); + let y_taps = taps(h, gh); + + let mut normals = vec![[0i8; 3]; w * h]; + let mut coverage = vec![0u8; w * h]; + let mut depth = vec![0u16; w * h]; + + normals + .par_chunks_exact_mut(w) + .zip(coverage.par_chunks_exact_mut(w)) + .zip(depth.par_chunks_exact_mut(w)) + .enumerate() + .for_each(|(y, ((normal_row, coverage_row), depth_row))| { + let ty = &y_taps[y]; + let (r0, r1) = (ty.i0 * gw, ty.i1 * gw); + + for x in 0..w { + let tx = &x_taps[x]; + let (i00, i10, i01, i11) = (r0 + tx.i0, r0 + tx.i1, r1 + tx.i0, r1 + tx.i1); + let bilerp = |v00: f32, v10: f32, v01: f32, v11: f32| { + let top = v00 + (v10 - v00) * tx.t; + let bot = v01 + (v11 - v01) * tx.t; + top + (bot - top) * ty.t + }; + + let mut values = [0.0f32; 5]; + for (c, value) in values.iter_mut().enumerate() { + let plane = &planes[c]; + *value = bilerp(plane[i00], plane[i10], plane[i01], plane[i11]); + } + + if refine { + let weight = bilerp( + edge_weight[i00], + edge_weight[i10], + edge_weight[i01], + edge_weight[i11], + ); + if weight > 1e-3 { + let luma = luma_full[y * w + x]; + for (c, value) in values.iter_mut().enumerate() { + let model = &models[c]; + let m = |k: usize| { + bilerp( + model[i00 * 4 + k], + model[i10 * 4 + k], + model[i01 * 4 + k], + model[i11 * 4 + k], + ) + }; + let mut guided = (m(0) * luma + m(1)).clamp(m(2), m(3)).clamp(0.0, 1.0); + if c == 3 { + // Coverage is a matte: push it back towards solid / empty so the + // photo's texture does not leak into it as a halo. + guided = smoothstep(0.25, 0.75, guided); + } + *value += (guided - *value) * weight; + } + } + } + + let nx = values[0] * 2.0 - 1.0; + let ny = values[1] * 2.0 - 1.0; + let nz = values[2] * 2.0 - 1.0; + let scale = 127.0 / (nx * nx + ny * ny + nz * nz).sqrt().max(1e-6); + normal_row[x] = [ + (nx * scale).round() as i8, + (ny * scale).round() as i8, + (nz * scale).round() as i8, + ]; + coverage_row[x] = (values[3].clamp(0.0, 1.0) * 255.0).round() as u8; + depth_row[x] = (values[4].clamp(0.0, 1.0) * 65535.0).round() as u16; + } + }); + + SurfaceMaps { + normals, + coverage, + depth, + } +} + fn srgb_to_linear(v: f32) -> f32 { if v <= 0.04045 { v / 12.92 @@ -99,15 +410,31 @@ fn parse_light_color(hex: &str) -> [f32; 3] { color.map(|c| c / norm) } -fn read_params(adjustments: &serde_json::Value) -> RelightParams { +fn read_params(adjustments: &serde_json::Value, has_depth_map: bool) -> RelightParams { let num = |key: &str, default: f64| adjustments[key].as_f64().unwrap_or(default) as f32; let range = (num("relightRange", 60.0) / 100.0).clamp(0.0, 1.0); + let depth_scale = if has_depth_map { + num("relightDepthScale", 0.0).max(0.0) + } else { + 0.0 + }; + + // Forwards the light moves linearly away from the scene. Backwards it eases in, so + // small moves place it just behind the subject even when the background is far away. + let depth = (num("relightDepth", 30.0) / 100.0).clamp(-1.0, 1.0); + let light_z = if depth >= 0.0 { + 0.02 + depth * MAX_FRONT_DISTANCE + } else { + 0.02 - depth * depth * depth_scale.clamp(0.5, 10.0) + }; + RelightParams { light_x: num("relightX", 0.3), light_y: num("relightY", 0.3), - light_z: 0.02 + (num("relightHeight", 30.0) / 100.0).clamp(0.0, 1.0), + light_z, + depth_scale, radius: 0.1 + 1.9 * range * range, light_color: parse_light_color(adjustments["relightColor"].as_str().unwrap_or("#ffffff")), intensity: (num("relightIntensity", 50.0) / 100.0).clamp(0.0, 1.0) * 2.0, @@ -117,16 +444,11 @@ fn read_params(adjustments: &serde_json::Value) -> RelightParams { } } -fn relight_in_place(image: &mut image::Rgb32FImage, normal_map: &RgbaImage, p: &RelightParams) { +fn relight_in_place(image: &mut Rgb32FImage, maps: &SurfaceMaps, p: &RelightParams) { let w = image.width() as usize; let h = image.height() as usize; - let nw = normal_map.width() as usize; - let nh = normal_map.height() as usize; - let normals = normal_map.as_raw(); - - let x_scale = nw as f32 / w as f32; - let y_scale = nh as f32 / h as f32; + let depth_to_z = -p.depth_scale / 65535.0; let inv_long_side = 1.0 / w.max(h) as f32; let light_px = p.light_x * w as f32; let light_py = p.light_y * h as f32; @@ -137,50 +459,31 @@ fn relight_in_place(image: &mut image::Rgb32FImage, normal_map: &RgbaImage, p: & .par_chunks_exact_mut(w * 3) .enumerate() .for_each(|(y, row)| { - let fy = ((y as f32 + 0.5) * y_scale - 0.5).clamp(0.0, (nh - 1) as f32); - let y0 = fy as usize; - let y1 = (y0 + 1).min(nh - 1); - let ty = fy - y0 as f32; - // Image rows grow downwards while the normal map uses Y up. let ly = -(light_py - (y as f32 + 0.5)) * inv_long_side; for (x, px) in row.chunks_exact_mut(3).enumerate() { - let fx = ((x as f32 + 0.5) * x_scale - 0.5).clamp(0.0, (nw - 1) as f32); - let x0 = fx as usize; - let x1 = (x0 + 1).min(nw - 1); - let tx = fx - x0 as f32; - - let mut s = [0.0f32; 4]; - for (sx, sy, weight) in [ - (x0, y0, (1.0 - tx) * (1.0 - ty)), - (x1, y0, tx * (1.0 - ty)), - (x0, y1, (1.0 - tx) * ty), - (x1, y1, tx * ty), - ] { - let i = (sy * nw + sx) * 4; - for c in 0..4 { - s[c] += normals[i + c] as f32 * weight; - } - } + let i = y * w + x; - let coverage = s[3] / 255.0; + let coverage = maps.coverage[i] as f32 / 255.0; if coverage <= 0.0 { continue; } - let nx = s[0] / 127.5 - 1.0; - let ny = s[1] / 127.5 - 1.0; - let nz = s[2] / 127.5 - 1.0; - let inv_len = 1.0 / (nx * nx + ny * ny + nz * nz).sqrt().max(1e-6); - let (nx, ny, nz) = (nx * inv_len, ny * inv_len, nz * inv_len); + let n = maps.normals[i]; + let (nx, ny, nz) = ( + n[0] as f32 / 127.0, + n[1] as f32 / 127.0, + n[2] as f32 / 127.0, + ); let lx = (light_px - (x as f32 + 0.5)) * inv_long_side; - let planar_dist_sq = lx * lx + ly * ly; - let inv_l_len = 1.0 / (planar_dist_sq + p.light_z * p.light_z).sqrt(); - let (dx, dy, dz) = (lx * inv_l_len, ly * inv_l_len, p.light_z * inv_l_len); + let lz = p.light_z - maps.depth[i] as f32 * depth_to_z; + let dist_sq = lx * lx + ly * ly + lz * lz; + let inv_l_len = 1.0 / dist_sq.sqrt().max(1e-6); + let (dx, dy, dz) = (lx * inv_l_len, ly * inv_l_len, lz * inv_l_len); - let attenuation = 1.0 / (1.0 + planar_dist_sq * inv_radius_sq); + let attenuation = 1.0 / (1.0 + dist_sq * inv_radius_sq); let n_dot_l = nx * dx + ny * dy + nz * dz; let diffuse = ((n_dot_l + p.wrap) / (1.0 + p.wrap)).max(0.0); @@ -215,15 +518,26 @@ mod tests { } fn relit_pixel(adjustments: serde_json::Value, normal_map: &RgbaImage, x: u32) -> [f32; 3] { + relit_pixel_with_depth(adjustments, normal_map, None, x) + } + + fn relit_pixel_with_depth( + adjustments: serde_json::Value, + normal_map: &RgbaImage, + depth_map: Option<&DepthMap>, + x: u32, + ) -> [f32; 3] { let mut image = Rgb32FImage::from_pixel(8, 8, Rgb([0.5, 0.5, 0.5])); - relight_in_place(&mut image, normal_map, &read_params(&adjustments)); + let params = read_params(&adjustments, depth_map.is_some()); + let maps = build_surface_maps(&image, normal_map, depth_map, params.depth_scale, false); + relight_in_place(&mut image, &maps, ¶ms); image.get_pixel(x, 3).0 } #[test] fn surface_facing_the_light_gets_brighter_than_one_facing_away() { let adjustments = json!({ - "relightX": 1.0, "relightY": 0.5, "relightHeight": 0, "relightSoftness": 0, "relightRange": 100, + "relightX": 1.0, "relightY": 0.5, "relightDepth": 0, "relightSoftness": 0, "relightRange": 100, }); let facing_right = relit_pixel( adjustments.clone(), @@ -239,7 +553,7 @@ mod tests { #[test] fn light_above_the_pixel_brightens_upward_facing_surfaces() { let adjustments = json!({ - "relightX": 0.5, "relightY": 0.0, "relightHeight": 0, "relightSoftness": 0, "relightRange": 100, + "relightX": 0.5, "relightY": 0.0, "relightDepth": 0, "relightSoftness": 0, "relightRange": 100, }); let facing_up = relit_pixel( adjustments.clone(), @@ -254,7 +568,7 @@ mod tests { #[test] fn light_falls_off_with_distance() { let adjustments = - json!({ "relightX": 0.0, "relightY": 0.5, "relightHeight": 50, "relightRange": 20 }); + json!({ "relightX": 0.0, "relightY": 0.5, "relightDepth": 50, "relightRange": 20 }); let normal_map = flat_normal_map([128, 128, 255], 255); let near = relit_pixel(adjustments.clone(), &normal_map, 0); let far = relit_pixel(adjustments, &normal_map, 7); @@ -262,6 +576,88 @@ mod tests { assert!(near[1] > far[1]); } + #[test] + fn light_behind_a_surface_does_not_light_its_front() { + let normal_map = flat_normal_map([128, 128, 255], 255); + let depth_map = DepthMap::from_pixel(4, 4, Luma([0])); + let adjustments = |depth: i32| { + json!({ + "relightX": 0.5, "relightY": 0.5, "relightDepth": depth, "relightSoftness": 0, + "relightAmbient": 100, "relightDepthScale": 2.0, + }) + }; + + let in_front = relit_pixel_with_depth(adjustments(40), &normal_map, Some(&depth_map), 3); + let behind = relit_pixel_with_depth(adjustments(-40), &normal_map, Some(&depth_map), 3); + + assert!(in_front[1] > 0.5); + assert_eq!(behind[1], 0.5); + } + + #[test] + fn light_behind_the_subject_still_reaches_the_background() { + let normal_map = flat_normal_map([128, 128, 255], 255); + let adjustments = json!({ + "relightX": 0.5, "relightY": 0.5, "relightDepth": -40, "relightSoftness": 0, + "relightAmbient": 100, "relightDepthScale": 2.0, + }); + + let subject = DepthMap::from_pixel(4, 4, Luma([0])); + let background = DepthMap::from_pixel(4, 4, Luma([u16::MAX])); + let on_subject = + relit_pixel_with_depth(adjustments.clone(), &normal_map, Some(&subject), 3); + let on_background = relit_pixel_with_depth(adjustments, &normal_map, Some(&background), 3); + + assert_eq!(on_subject[1], 0.5); + assert!(on_background[1] > 0.5); + } + + #[test] + fn refinement_snaps_a_soft_depth_edge_to_the_photo_edge() { + // Photo: dark subject on the left, bright background on the right, hard edge at x = 2048. + let image = Rgb32FImage::from_fn(4096, 96, |x, _| { + if x < 2048 { + Rgb([0.05, 0.05, 0.05]) + } else { + Rgb([0.8, 0.8, 0.8]) + } + }); + // Maps: the same split at 1/16 of the resolution, so the edge comes out blurry. + let normal_map = RgbaImage::from_pixel(256, 6, Rgba([128, 128, 255, 255])); + let depth_map = + DepthMap::from_fn(256, 6, |x, _| Luma([if x < 128 { 0 } else { u16::MAX }])); + + // Depth change over the few pixels around the photo edge. + let depth_step = |refine: bool| { + let maps = build_surface_maps(&image, &normal_map, Some(&depth_map), 2.0, refine); + let row = 48 * 4096; + maps.depth[row + 2051] as f32 - maps.depth[row + 2045] as f32 + }; + + let soft = depth_step(false); + let refined = depth_step(true); + + assert!(soft > 0.0); + assert!(refined > soft * 1.5, "refined {refined} vs bilinear {soft}"); + } + + #[test] + fn refinement_keeps_normals_away_from_edges() { + let image = Rgb32FImage::from_fn(128, 128, |x, y| { + let v = ((x / 8 + y / 8) % 2) as f32 * 0.5 + 0.2; + Rgb([v, v, v]) + }); + // A textured photo over a flat, evenly tilted surface: there is no geometric edge to snap to. + let normal_map = RgbaImage::from_pixel(16, 16, Rgba([200, 128, 220, 255])); + let depth_map = DepthMap::from_pixel(16, 16, Luma([1000])); + + let plain = build_surface_maps(&image, &normal_map, Some(&depth_map), 2.0, false); + let refined = build_surface_maps(&image, &normal_map, Some(&depth_map), 2.0, true); + + assert_eq!(plain.normals, refined.normals); + assert_eq!(plain.depth, refined.depth); + } + #[test] fn light_color_tints_the_lit_surface() { let adjustments = json!({ "relightX": 0.5, "relightY": 0.5, "relightColor": "#ff8000" }); diff --git a/src/components/adjustments/Effects.tsx b/src/components/adjustments/Effects.tsx index 2dd80afb17..546c8eed74 100644 --- a/src/components/adjustments/Effects.tsx +++ b/src/components/adjustments/Effects.tsx @@ -178,10 +178,14 @@ export default function EffectsPanel({ const handleGenerateRelightNormalMap = async () => { setIsGeneratingNormal(true); try { - const b64: string = await invoke('generate_full_image_normal_map', { jsAdjustments: adjustments }); + const maps: { normalMap: string; depthMap: string; depthScale: number } = await invoke('generate_relight_maps', { + jsAdjustments: adjustments, + }); setAdjustments((prev: Partial) => ({ ...prev, - relightNormalMap: b64, + relightNormalMap: maps.normalMap, + relightDepthMap: maps.depthMap, + relightDepthScale: maps.depthScale, })); } catch (e: any) { toast.error(`Failed to generate normal map: ${e}`); @@ -227,7 +231,7 @@ export default function EffectsPanel({ const handleRelightToggle = (enabled: boolean) => { handleAdjustmentChange(Effect.RelightEnabled, enabled); - if (enabled && !adjustments.relightNormalMap) { + if (enabled && (!adjustments.relightNormalMap || !adjustments.relightDepthMap)) { handleGenerateRelightNormalMap(); } }; @@ -432,15 +436,14 @@ export default function EffectsPanel({ /> handleAdjustmentChange(Effect.RelightHeight, e.target.value)} + onChange={(e: any) => handleAdjustmentChange(Effect.RelightDepth, e.target.value)} step={1} - value={adjustments.relightHeight ?? 30} + value={adjustments.relightDepth ?? 30} onDragStateChange={onDragStateChange} - fillOrigin="min" /> Math.max(0, Math.min(1, value)); @@ -36,6 +37,25 @@ export default function RelightHandle({ [], ); + // The editor zooms with a native wheel listener, so this has to be native too to get in first. + useEffect(() => { + const handle = handleRef.current; + if (!handle) return; + + const handleWheel = (e: WheelEvent) => { + e.preventDefault(); + e.stopPropagation(); + const step = e.deltaY < 0 ? WHEEL_DEPTH_STEP : -WHEEL_DEPTH_STEP; + setAdjustments((prev: Adjustments) => ({ + ...prev, + relightDepth: Math.max(-100, Math.min(100, (prev.relightDepth ?? 30) + step)), + })); + }; + + handle.addEventListener('wheel', handleWheel, { passive: false }); + return () => handle.removeEventListener('wheel', handleWheel); + }, [setAdjustments]); + const commitUv = useCallback( (uv: Coord) => { pendingUvRef.current = uv; @@ -101,6 +121,10 @@ export default function RelightHandle({ const uv = dragUv ?? { x: adjustments.relightX ?? 0.3, y: adjustments.relightY ?? 0.3 }; const display = relightUvToDisplay(uv, geometry); const color = adjustments.relightColor || '#ffffff'; + const depth = Math.max(-100, Math.min(100, adjustments.relightDepth ?? 30)); + // Closer lights look bigger; a light behind the subject gets a dashed outline. + const depthScale = 0.7 + 0.6 * ((depth + 100) / 200); + const isBehind = depth < 0; return (
diff --git a/src/i18n/locales/ca.json b/src/i18n/locales/ca.json index 04b04396f7..fd273dbce6 100644 --- a/src/i18n/locales/ca.json +++ b/src/i18n/locales/ca.json @@ -131,9 +131,9 @@ "relightAmbient": "Ambient", "relightColor": "Color de la llum", "relightCustomColor": "Color personalitzat", - "relightDragLight": "Arrossega per moure la llum", + "relightDepth": "Profunditat", + "relightDragLight": "Arrossega per moure la llum, fes servir la roda per acostar-la o allunyar-la", "relightGenerateNormalMap": "Genera el mapa de normals", - "relightHeight": "Alçada", "relightIntensity": "Intensitat", "relightNormalMap": "Mapa de normals", "relightRange": "Abast", diff --git a/src/i18n/locales/de.json b/src/i18n/locales/de.json index 96f18b648c..63fd49fa2c 100644 --- a/src/i18n/locales/de.json +++ b/src/i18n/locales/de.json @@ -131,9 +131,9 @@ "relightAmbient": "Umgebungslicht", "relightColor": "Lichtfarbe", "relightCustomColor": "Eigene Farbe", - "relightDragLight": "Ziehen, um das Licht zu bewegen", + "relightDepth": "Tiefe", + "relightDragLight": "Ziehen, um das Licht zu bewegen, scrollen, um es näher oder weiter weg zu setzen", "relightGenerateNormalMap": "Normal Map erzeugen", - "relightHeight": "Höhe", "relightIntensity": "Intensität", "relightNormalMap": "Normal Map", "relightRange": "Reichweite", diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index 61c086d95d..9a3cf78311 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -131,9 +131,9 @@ "relightAmbient": "Ambient", "relightColor": "Light Color", "relightCustomColor": "Custom Color", - "relightDragLight": "Drag to move the light", + "relightDepth": "Depth", + "relightDragLight": "Drag to move the light, scroll to move it closer or farther", "relightGenerateNormalMap": "Generate Normal Map", - "relightHeight": "Height", "relightIntensity": "Intensity", "relightNormalMap": "Normal Map", "relightRange": "Range", diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index 409c4ca440..34d3a837b5 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -131,9 +131,9 @@ "relightAmbient": "Ambiente", "relightColor": "Color de la luz", "relightCustomColor": "Color personalizado", - "relightDragLight": "Arrastra para mover la luz", + "relightDepth": "Profundidad", + "relightDragLight": "Arrastra para mover la luz, usa la rueda para acercarla o alejarla", "relightGenerateNormalMap": "Generar mapa de normales", - "relightHeight": "Altura", "relightIntensity": "Intensidad", "relightNormalMap": "Mapa de normales", "relightRange": "Alcance", diff --git a/src/i18n/locales/fr.json b/src/i18n/locales/fr.json index 6bc18d4420..40590442c5 100644 --- a/src/i18n/locales/fr.json +++ b/src/i18n/locales/fr.json @@ -131,9 +131,9 @@ "relightAmbient": "Ambiance", "relightColor": "Couleur de la lumière", "relightCustomColor": "Couleur personnalisée", - "relightDragLight": "Faites glisser pour déplacer la lumière", + "relightDepth": "Profondeur", + "relightDragLight": "Faites glisser pour déplacer la lumière, utilisez la molette pour la rapprocher ou l'éloigner", "relightGenerateNormalMap": "Générer la carte de normales", - "relightHeight": "Hauteur", "relightIntensity": "Intensité", "relightNormalMap": "Carte de normales", "relightRange": "Portée", diff --git a/src/i18n/locales/it.json b/src/i18n/locales/it.json index 6dd20314ba..6cd4288d59 100644 --- a/src/i18n/locales/it.json +++ b/src/i18n/locales/it.json @@ -131,9 +131,9 @@ "relightAmbient": "Ambiente", "relightColor": "Colore della luce", "relightCustomColor": "Colore personalizzato", - "relightDragLight": "Trascina per spostare la luce", + "relightDepth": "Profondità", + "relightDragLight": "Trascina per spostare la luce, usa la rotellina per avvicinarla o allontanarla", "relightGenerateNormalMap": "Genera mappa delle normali", - "relightHeight": "Altezza", "relightIntensity": "Intensità", "relightNormalMap": "Mappa delle normali", "relightRange": "Portata", diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index 865c54a97e..52f967425f 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -131,9 +131,9 @@ "relightAmbient": "環境光", "relightColor": "ライトの色", "relightCustomColor": "カスタムカラー", - "relightDragLight": "ドラッグしてライトを移動", + "relightDepth": "奥行き", + "relightDragLight": "ドラッグでライトを移動、スクロールで前後に移動", "relightGenerateNormalMap": "法線マップを生成", - "relightHeight": "高さ", "relightIntensity": "強度", "relightNormalMap": "法線マップ", "relightRange": "範囲", diff --git a/src/i18n/locales/ko.json b/src/i18n/locales/ko.json index a015755ace..7eb16f2cf5 100644 --- a/src/i18n/locales/ko.json +++ b/src/i18n/locales/ko.json @@ -131,9 +131,9 @@ "relightAmbient": "주변광", "relightColor": "조명 색상", "relightCustomColor": "사용자 지정 색상", - "relightDragLight": "드래그하여 조명 이동", + "relightDepth": "깊이", + "relightDragLight": "드래그하여 조명 이동, 스크롤하여 앞뒤로 이동", "relightGenerateNormalMap": "노멀 맵 생성", - "relightHeight": "높이", "relightIntensity": "강도", "relightNormalMap": "노멀 맵", "relightRange": "범위", diff --git a/src/i18n/locales/pl.json b/src/i18n/locales/pl.json index 27b56705aa..b069a3b67a 100644 --- a/src/i18n/locales/pl.json +++ b/src/i18n/locales/pl.json @@ -131,9 +131,9 @@ "relightAmbient": "Światło otoczenia", "relightColor": "Kolor światła", "relightCustomColor": "Własny kolor", - "relightDragLight": "Przeciągnij, aby przesunąć światło", + "relightDepth": "Głębia", + "relightDragLight": "Przeciągnij, aby przesunąć światło, przewiń, aby je przybliżyć lub oddalić", "relightGenerateNormalMap": "Generuj mapę normalnych", - "relightHeight": "Wysokość", "relightIntensity": "Intensywność", "relightNormalMap": "Mapa normalnych", "relightRange": "Zasięg", diff --git a/src/i18n/locales/pt.json b/src/i18n/locales/pt.json index 348edcfa74..a8f89fbaf0 100644 --- a/src/i18n/locales/pt.json +++ b/src/i18n/locales/pt.json @@ -131,9 +131,9 @@ "relightAmbient": "Ambiente", "relightColor": "Cor da luz", "relightCustomColor": "Cor personalizada", - "relightDragLight": "Arraste para mover a luz", + "relightDepth": "Profundidade", + "relightDragLight": "Arraste para mover a luz, use a roda para a aproximar ou afastar", "relightGenerateNormalMap": "Gerar mapa de normais", - "relightHeight": "Altura", "relightIntensity": "Intensidade", "relightNormalMap": "Mapa de normais", "relightRange": "Alcance", diff --git a/src/i18n/locales/ru.json b/src/i18n/locales/ru.json index 622535d613..bff9c0d5a3 100644 --- a/src/i18n/locales/ru.json +++ b/src/i18n/locales/ru.json @@ -131,9 +131,9 @@ "relightAmbient": "Окружающий свет", "relightColor": "Цвет света", "relightCustomColor": "Свой цвет", - "relightDragLight": "Перетащите, чтобы переместить свет", + "relightDepth": "Глубина", + "relightDragLight": "Перетащите, чтобы переместить свет, прокрутите, чтобы приблизить или отдалить его", "relightGenerateNormalMap": "Создать карту нормалей", - "relightHeight": "Высота", "relightIntensity": "Интенсивность", "relightNormalMap": "Карта нормалей", "relightRange": "Дальность", diff --git a/src/i18n/locales/zh-CN.json b/src/i18n/locales/zh-CN.json index 1daac1eb23..e95f08a8d6 100644 --- a/src/i18n/locales/zh-CN.json +++ b/src/i18n/locales/zh-CN.json @@ -131,9 +131,9 @@ "relightAmbient": "环境光", "relightColor": "灯光颜色", "relightCustomColor": "自定义颜色", - "relightDragLight": "拖动以移动灯光", + "relightDepth": "深度", + "relightDragLight": "拖动以移动灯光，滚动以前后移动", "relightGenerateNormalMap": "生成法线贴图", - "relightHeight": "高度", "relightIntensity": "强度", "relightNormalMap": "法线贴图", "relightRange": "范围", diff --git a/src/i18n/locales/zh-TW.json b/src/i18n/locales/zh-TW.json index 61db256dc5..80df6f03c0 100644 --- a/src/i18n/locales/zh-TW.json +++ b/src/i18n/locales/zh-TW.json @@ -131,9 +131,9 @@ "relightAmbient": "環境光", "relightColor": "燈光顏色", "relightCustomColor": "自訂顏色", - "relightDragLight": "拖曳以移動燈光", + "relightDepth": "深度", + "relightDragLight": "拖曳以移動燈光，捲動以前後移動", "relightGenerateNormalMap": "產生法線貼圖", - "relightHeight": "高度", "relightIntensity": "強度", "relightNormalMap": "法線貼圖", "relightRange": "範圍", diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index a61b646580..deab44a812 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -100,13 +100,15 @@ export enum Effect { RelightIntensity = 'relightIntensity', RelightX = 'relightX', RelightY = 'relightY', - RelightHeight = 'relightHeight', + RelightDepth = 'relightDepth', RelightRange = 'relightRange', RelightAmbient = 'relightAmbient', RelightSoftness = 'relightSoftness', RelightSpecular = 'relightSpecular', RelightColor = 'relightColor', RelightNormalMap = 'relightNormalMap', + RelightDepthMap = 'relightDepthMap', + RelightDepthScale = 'relightDepthScale', } export enum CreativeAdjustment { @@ -256,8 +258,10 @@ export interface Adjustments { orientationSteps: number; relightAmbient: number; relightColor: string; + relightDepth: number; + relightDepthMap: string | null; + relightDepthScale: number; relightEnabled: boolean; - relightHeight: number; relightIntensity: number; relightNormalMap: string | null; relightRange: number; @@ -605,8 +609,10 @@ export const INITIAL_ADJUSTMENTS: Adjustments = { orientationSteps: 0, relightAmbient: 80, relightColor: '#ffffff', + relightDepth: 30, + relightDepthMap: null, + relightDepthScale: 0, relightEnabled: false, - relightHeight: 30, relightIntensity: 50, relightNormalMap: null, relightRange: 60, @@ -763,8 +769,10 @@ export const normalizeLoadedAdjustments = (loadedAdjustments: Adjustments): any lensBlurMinFade: loadedAdjustments.lensBlurMinFade ?? INITIAL_ADJUSTMENTS.lensBlurMinFade, relightAmbient: loadedAdjustments.relightAmbient ?? INITIAL_ADJUSTMENTS.relightAmbient, relightColor: loadedAdjustments.relightColor ?? INITIAL_ADJUSTMENTS.relightColor, + relightDepth: loadedAdjustments.relightDepth ?? INITIAL_ADJUSTMENTS.relightDepth, + relightDepthMap: loadedAdjustments.relightDepthMap ?? INITIAL_ADJUSTMENTS.relightDepthMap, + relightDepthScale: loadedAdjustments.relightDepthScale ?? INITIAL_ADJUSTMENTS.relightDepthScale, relightEnabled: loadedAdjustments.relightEnabled ?? INITIAL_ADJUSTMENTS.relightEnabled, - relightHeight: loadedAdjustments.relightHeight ?? INITIAL_ADJUSTMENTS.relightHeight, relightIntensity: loadedAdjustments.relightIntensity ?? INITIAL_ADJUSTMENTS.relightIntensity, relightNormalMap: loadedAdjustments.relightNormalMap ?? INITIAL_ADJUSTMENTS.relightNormalMap, relightRange: loadedAdjustments.relightRange ?? INITIAL_ADJUSTMENTS.relightRange, @@ -1000,13 +1008,15 @@ export const ADJUSTMENT_SECTIONS: Sections = { Effect.RelightIntensity, Effect.RelightX, Effect.RelightY, - Effect.RelightHeight, + Effect.RelightDepth, Effect.RelightRange, Effect.RelightAmbient, Effect.RelightSoftness, Effect.RelightSpecular, Effect.RelightColor, Effect.RelightNormalMap, + Effect.RelightDepthMap, + Effect.RelightDepthScale, ], }; From 5bd99165cd4a2004f191158b8f4e4684caf2202d Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Mon, 5 Oct 2026 16:27:21 -0500 Subject: [PATCH 06/16] feat: add localization files and relight effect components --- src-tauri/src/cache_utils.rs | 5 +- src-tauri/src/relight.rs | 185 +++++++++++++----- src/components/adjustments/Effects.tsx | 178 ++++++++++++----- src/components/panel/editor/ImageCanvas.tsx | 47 ++++- .../panel/editor/overlays/RelightHandle.tsx | 44 +++-- src/i18n/locales/ca.json | 3 + src/i18n/locales/de.json | 3 + src/i18n/locales/en.json | 3 + src/i18n/locales/es.json | 3 + src/i18n/locales/fr.json | 3 + src/i18n/locales/it.json | 3 + src/i18n/locales/ja.json | 3 + src/i18n/locales/ko.json | 3 + src/i18n/locales/pl.json | 3 + src/i18n/locales/pt.json | 3 + src/i18n/locales/ru.json | 3 + src/i18n/locales/zh-CN.json | 3 + src/i18n/locales/zh-TW.json | 3 + src/store/useEditorStore.ts | 2 + src/utils/adjustments.ts | 81 +++++--- 20 files changed, 423 insertions(+), 158 deletions(-) diff --git a/src-tauri/src/cache_utils.rs b/src-tauri/src/cache_utils.rs index b31974cf08..6f86bafe44 100644 --- a/src-tauri/src/cache_utils.rs +++ b/src-tauri/src/cache_utils.rs @@ -93,15 +93,16 @@ fn hash_relight( if relight_enabled { let relight_keys = [ + "relightLights", "relightIntensity", "relightX", "relightY", "relightDepth", + "relightRange", + "relightColor", "relightAmbient", "relightSoftness", "relightSpecular", - "relightRange", - "relightColor", "relightNormalMap", "relightDepthMap", "relightDepthScale", diff --git a/src-tauri/src/relight.rs b/src-tauri/src/relight.rs index 67ed213a42..1532c1ea17 100644 --- a/src-tauri/src/relight.rs +++ b/src-tauri/src/relight.rs @@ -12,6 +12,7 @@ const MAX_FRONT_DISTANCE: f32 = 1.5; // Below this size there is nothing to gain from snapping the maps to the image edges. const MIN_REFINE_SIZE: u32 = 64; +const MAX_LIGHTS: usize = 8; type DepthMap = ImageBuffer, Vec>; @@ -26,19 +27,23 @@ struct SurfaceMaps { // last result is kept around while the light is being moved. static SURFACE_MAPS_CACHE: Mutex)>> = Mutex::new(None); -struct RelightParams { - // Light position as a fraction of the image width / height. - light_x: f32, - light_y: f32, - // Light depth in units of the longest image side: 0 is the nearest surface, - // positive values move towards the viewer and negative values go into the scene. - light_z: f32, - // Depth covered by the full range of the depth map, in the same units. - depth_scale: f32, +struct Light { + // Position as a fraction of the image width / height. + x: f32, + y: f32, + // Depth in units of the longest image side: 0 is the nearest surface, positive values + // move towards the viewer and negative values go into the scene. + z: f32, // Falloff radius, in units of the longest image side. radius: f32, - light_color: [f32; 3], + color: [f32; 3], intensity: f32, +} + +struct RelightParams { + lights: Vec, + // Depth covered by the full range of the depth map, in units of the longest image side. + depth_scale: f32, ambient: f32, wrap: f32, specular: f32, @@ -410,34 +415,62 @@ fn parse_light_color(hex: &str) -> [f32; 3] { color.map(|c| c / norm) } +fn read_light(light: &serde_json::Value, depth_scale: f32) -> Light { + let num = |key: &str, default: f64| light[key].as_f64().unwrap_or(default) as f32; + + let range = (num("range", 60.0) / 100.0).clamp(0.0, 1.0); + + // Forwards the light moves linearly away from the scene. Backwards it eases in, so + // small moves place it just behind the subject even when the background is far away. + let depth = (num("depth", 30.0) / 100.0).clamp(-1.0, 1.0); + let z = if depth >= 0.0 { + 0.02 + depth * MAX_FRONT_DISTANCE + } else { + 0.02 - depth * depth * depth_scale.clamp(0.5, 10.0) + }; + + Light { + x: num("x", 0.3), + y: num("y", 0.3), + z, + radius: 0.1 + 1.9 * range * range, + color: parse_light_color(light["color"].as_str().unwrap_or("#ffffff")), + intensity: (num("intensity", 50.0) / 100.0).clamp(0.0, 1.0) * 2.0, + } +} + fn read_params(adjustments: &serde_json::Value, has_depth_map: bool) -> RelightParams { let num = |key: &str, default: f64| adjustments[key].as_f64().unwrap_or(default) as f32; - let range = (num("relightRange", 60.0) / 100.0).clamp(0.0, 1.0); - let depth_scale = if has_depth_map { num("relightDepthScale", 0.0).max(0.0) } else { 0.0 }; - // Forwards the light moves linearly away from the scene. Backwards it eases in, so - // small moves place it just behind the subject even when the background is far away. - let depth = (num("relightDepth", 30.0) / 100.0).clamp(-1.0, 1.0); - let light_z = if depth >= 0.0 { - 0.02 + depth * MAX_FRONT_DISTANCE - } else { - 0.02 - depth * depth * depth_scale.clamp(0.5, 10.0) + let lights = match adjustments["relightLights"].as_array() { + Some(lights) => lights + .iter() + .take(MAX_LIGHTS) + .map(|light| read_light(light, depth_scale)) + .collect(), + // Edits saved before multiple lights existed describe a single light at the top level. + None => vec![read_light( + &serde_json::json!({ + "x": adjustments["relightX"], + "y": adjustments["relightY"], + "depth": adjustments["relightDepth"], + "range": adjustments["relightRange"], + "color": adjustments["relightColor"], + "intensity": adjustments["relightIntensity"], + }), + depth_scale, + )], }; RelightParams { - light_x: num("relightX", 0.3), - light_y: num("relightY", 0.3), - light_z, + lights, depth_scale, - radius: 0.1 + 1.9 * range * range, - light_color: parse_light_color(adjustments["relightColor"].as_str().unwrap_or("#ffffff")), - intensity: (num("relightIntensity", 50.0) / 100.0).clamp(0.0, 1.0) * 2.0, ambient: (num("relightAmbient", 80.0) / 100.0).clamp(0.0, 1.0), wrap: (num("relightSoftness", 30.0) / 100.0).clamp(0.0, 1.0), specular: (num("relightSpecular", 0.0) / 100.0).clamp(0.0, 1.0), @@ -450,18 +483,28 @@ fn relight_in_place(image: &mut Rgb32FImage, maps: &SurfaceMaps, p: &RelightPara let depth_to_z = -p.depth_scale / 65535.0; let inv_long_side = 1.0 / w.max(h) as f32; - let light_px = p.light_x * w as f32; - let light_py = p.light_y * h as f32; - let inv_radius_sq = 1.0 / (p.radius * p.radius); + + // Per light: position in pixels, depth and inverse squared falloff radius. + let lights: Vec<(f32, f32, f32, f32, &Light)> = p + .lights + .iter() + .filter(|light| light.intensity > 0.0) + .map(|light| { + ( + light.x * w as f32, + light.y * h as f32, + light.z, + 1.0 / (light.radius * light.radius), + light, + ) + }) + .collect(); image .as_mut() .par_chunks_exact_mut(w * 3) .enumerate() .for_each(|(y, row)| { - // Image rows grow downwards while the normal map uses Y up. - let ly = -(light_py - (y as f32 + 0.5)) * inv_long_side; - for (x, px) in row.chunks_exact_mut(3).enumerate() { let i = y * w + x; @@ -476,31 +519,41 @@ fn relight_in_place(image: &mut Rgb32FImage, maps: &SurfaceMaps, p: &RelightPara n[1] as f32 / 127.0, n[2] as f32 / 127.0, ); + let surface_z = maps.depth[i] as f32 * depth_to_z; + + let mut lit = [0.0f32; 3]; + for (light_px, light_py, light_z, inv_radius_sq, light) in &lights { + let lx = (light_px - (x as f32 + 0.5)) * inv_long_side; + // Image rows grow downwards while the normal map uses Y up. + let ly = -(light_py - (y as f32 + 0.5)) * inv_long_side; + let lz = light_z - surface_z; + let dist_sq = lx * lx + ly * ly + lz * lz; + let inv_l_len = 1.0 / dist_sq.sqrt().max(1e-6); + let (dx, dy, dz) = (lx * inv_l_len, ly * inv_l_len, lz * inv_l_len); + + let attenuation = 1.0 / (1.0 + dist_sq * inv_radius_sq); + + let n_dot_l = nx * dx + ny * dy + nz * dz; + let diffuse = ((n_dot_l + p.wrap) / (1.0 + p.wrap)).max(0.0); + + let spec = if p.specular > 0.0 && n_dot_l > 0.0 { + // Half vector between the light and the viewer (0, 0, 1). + let hz = dz + 1.0; + let inv_h_len = 1.0 / (dx * dx + dy * dy + hz * hz).sqrt().max(1e-6); + let n_dot_h = (nx * dx + ny * dy + nz * hz) * inv_h_len; + n_dot_h.max(0.0).powf(SPECULAR_POWER) * p.specular + } else { + 0.0 + }; + + let amount = light.intensity * attenuation * (diffuse + spec); + for c in 0..3 { + lit[c] += amount * light.color[c]; + } + } - let lx = (light_px - (x as f32 + 0.5)) * inv_long_side; - let lz = p.light_z - maps.depth[i] as f32 * depth_to_z; - let dist_sq = lx * lx + ly * ly + lz * lz; - let inv_l_len = 1.0 / dist_sq.sqrt().max(1e-6); - let (dx, dy, dz) = (lx * inv_l_len, ly * inv_l_len, lz * inv_l_len); - - let attenuation = 1.0 / (1.0 + dist_sq * inv_radius_sq); - - let n_dot_l = nx * dx + ny * dy + nz * dz; - let diffuse = ((n_dot_l + p.wrap) / (1.0 + p.wrap)).max(0.0); - - let spec = if p.specular > 0.0 && n_dot_l > 0.0 { - // Half vector between the light and the viewer (0, 0, 1). - let hz = dz + 1.0; - let inv_h_len = 1.0 / (dx * dx + dy * dy + hz * hz).sqrt().max(1e-6); - let n_dot_h = (nx * dx + ny * dy + nz * hz) * inv_h_len; - n_dot_h.max(0.0).powf(SPECULAR_POWER) * p.specular - } else { - 0.0 - }; - - let lit = p.intensity * attenuation * (diffuse + spec); for c in 0..3 { - let gain = p.ambient + lit * p.light_color[c]; + let gain = p.ambient + lit[c]; px[c] *= 1.0 + (gain - 1.0) * coverage; } } @@ -658,6 +711,30 @@ mod tests { assert_eq!(plain.depth, refined.depth); } + #[test] + fn several_lights_add_up_with_their_own_colors() { + let normal_map = flat_normal_map([128, 128, 255], 255); + let red = json!({ "x": 0.0, "y": 0.5, "color": "#ff0000", "range": 20 }); + let blue = json!({ "x": 1.0, "y": 0.5, "color": "#0000ff", "range": 20 }); + let both = json!({ "relightLights": [red, blue] }); + + let near_red = relit_pixel(both.clone(), &normal_map, 0); + let near_blue = relit_pixel(both.clone(), &normal_map, 7); + assert!(near_red[0] > near_red[2]); + assert!(near_blue[2] > near_blue[0]); + + let only_red = relit_pixel(json!({ "relightLights": [red] }), &normal_map, 0); + assert!(near_red[2] > only_red[2]); + } + + #[test] + fn no_lights_leaves_only_the_ambient_level() { + let adjustments = json!({ "relightLights": [], "relightAmbient": 50 }); + let lit = relit_pixel(adjustments, &flat_normal_map([128, 128, 255], 255), 3); + + assert_eq!(lit, [0.25; 3]); + } + #[test] fn light_color_tints_the_lit_surface() { let adjustments = json!({ "relightX": 0.5, "relightY": 0.5, "relightColor": "#ff8000" }); diff --git a/src/components/adjustments/Effects.tsx b/src/components/adjustments/Effects.tsx index 546c8eed74..3ce2592294 100644 --- a/src/components/adjustments/Effects.tsx +++ b/src/components/adjustments/Effects.tsx @@ -2,7 +2,7 @@ import { useState, useEffect, useRef, useMemo } from 'react'; import { useTranslation } from 'react-i18next'; import { invoke } from '@tauri-apps/api/core'; import { toast } from 'react-toastify'; -import { Loader2, Circle, Hexagon, Octagon, Aperture } from 'lucide-react'; +import { Loader2, Circle, Hexagon, Octagon, Aperture, Plus, Trash2 } from 'lucide-react'; import { motion } from 'framer-motion'; import clsx from 'clsx'; import Slider from '../ui/Slider'; @@ -14,6 +14,9 @@ import { CreativeAdjustment, getAdjustmentToolOrder, getHiddenAdjustmentTools, + RelightLight, + DEFAULT_RELIGHT_LIGHT, + MAX_RELIGHT_LIGHTS, } from '../../utils/adjustments'; import LUTControl from '../ui/LUTControl'; import { AppSettings } from '../ui/AppProperties'; @@ -23,6 +26,10 @@ import { TextVariants } from '../../types/typography'; import { DepthRangePicker } from '../ui/DepthRangePicker'; import { useProcessStore } from '../../store/useProcessStore'; import { RELIGHT_TEMPERATURE_PRESETS, kelvinToHex } from '../../utils/relightUtils'; +import { useEditorStore } from '../../store/useEditorStore'; +import { v4 as uuidv4 } from 'uuid'; + +const NEW_RELIGHT_LIGHT_TEMPERATURES = [5600, 3200, 8000, 2700, 10000, 4000]; interface EffectsPanelProps { adjustments: Adjustments; @@ -223,10 +230,51 @@ export default function EffectsPanel({ } }; - const relightColor = (adjustments.relightColor || '#ffffff').toLowerCase(); + const relightLights = adjustments.relightLights ?? []; + const activeRelightLightId = useEditorStore((state) => state.activeRelightLightId); + const setEditor = useEditorStore((state) => state.setEditor); + const activeRelightLight = relightLights.find((light) => light.id === activeRelightLightId) ?? relightLights[0]; + const relightColor = (activeRelightLight?.color || '#ffffff').toLowerCase(); + + const updateActiveRelightLight = (changes: Partial) => { + if (!activeRelightLight) return; + const id = activeRelightLight.id; + setAdjustments((prev: Partial) => ({ + ...prev, + relightLights: (prev.relightLights ?? []).map((light) => (light.id === id ? { ...light, ...changes } : light)), + })); + }; const handleRelightColorChange = (color: string) => { - setAdjustments((prev: Partial) => ({ ...prev, relightColor: color.toLowerCase() })); + updateActiveRelightLight({ color: color.toLowerCase() }); + }; + + const handleAddRelightLight = () => { + if (relightLights.length >= MAX_RELIGHT_LIGHTS) return; + // Each new light starts somewhere else and with another tint, so it is easy to tell apart. + const index = relightLights.length; + const light: RelightLight = { + ...DEFAULT_RELIGHT_LIGHT, + id: uuidv4(), + x: [0.7, 0.5, 0.3, 0.7, 0.5][index % 5], + y: [0.3, 0.7, 0.7, 0.7, 0.2][index % 5], + color: kelvinToHex(NEW_RELIGHT_LIGHT_TEMPERATURES[index % NEW_RELIGHT_LIGHT_TEMPERATURES.length]), + }; + setAdjustments((prev: Partial) => ({ + ...prev, + relightLights: [...(prev.relightLights ?? []), light], + })); + setEditor({ activeRelightLightId: light.id }); + }; + + const handleRemoveRelightLight = () => { + if (!activeRelightLight || relightLights.length <= 1) return; + const id = activeRelightLight.id; + setAdjustments((prev: Partial) => ({ + ...prev, + relightLights: (prev.relightLights ?? []).filter((light) => light.id !== id), + })); + setEditor({ activeRelightLightId: null }); }; const handleRelightToggle = (enabled: boolean) => { @@ -423,14 +471,54 @@ export default function EffectsPanel({
) : ( <> +
+ + {t('adjustments.effects.relightLights')} + +
+ {relightLights.map((light) => ( + + )} + {relightLights.length > 1 && ( + + )} +
+
+ handleAdjustmentChange(Effect.RelightIntensity, e.target.value)} + onChange={(e: any) => updateActiveRelightLight({ intensity: parseInt(e.target.value, 10) })} step={1} - value={adjustments.relightIntensity ?? 50} + value={activeRelightLight?.intensity ?? 50} onDragStateChange={onDragStateChange} fillOrigin="min" /> @@ -440,9 +528,9 @@ export default function EffectsPanel({ max={100} min={-100} defaultValue={30} - onChange={(e: any) => handleAdjustmentChange(Effect.RelightDepth, e.target.value)} + onChange={(e: any) => updateActiveRelightLight({ depth: parseInt(e.target.value, 10) })} step={1} - value={adjustments.relightDepth ?? 30} + value={activeRelightLight?.depth ?? 30} onDragStateChange={onDragStateChange} /> @@ -451,45 +539,9 @@ export default function EffectsPanel({ max={100} min={0} defaultValue={60} - onChange={(e: any) => handleAdjustmentChange(Effect.RelightRange, e.target.value)} - step={1} - value={adjustments.relightRange ?? 60} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - handleAdjustmentChange(Effect.RelightAmbient, e.target.value)} - step={1} - value={adjustments.relightAmbient ?? 80} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - handleAdjustmentChange(Effect.RelightSoftness, e.target.value)} - step={1} - value={adjustments.relightSoftness ?? 30} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - handleAdjustmentChange(Effect.RelightSpecular, e.target.value)} + onChange={(e: any) => updateActiveRelightLight({ range: parseInt(e.target.value, 10) })} step={1} - value={adjustments.relightSpecular ?? 0} + value={activeRelightLight?.range ?? 60} onDragStateChange={onDragStateChange} fillOrigin="min" /> @@ -535,6 +587,42 @@ export default function EffectsPanel({
+ handleAdjustmentChange(Effect.RelightAmbient, e.target.value)} + step={1} + value={adjustments.relightAmbient ?? 80} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightSoftness, e.target.value)} + step={1} + value={adjustments.relightSoftness ?? 30} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleAdjustmentChange(Effect.RelightSpecular, e.target.value)} + step={1} + value={adjustments.relightSpecular ?? 0} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + {adjustments.relightNormalMap && ( state.activeRelightLightId); + const setEditor = useEditorStore((state) => state.setEditor); + const handleRelightLightSelect = useCallback((id: string) => setEditor({ activeRelightLightId: id }), [setEditor]); + const handleRelightLightChange = useCallback( + (id: string, changes: Partial) => { + setAdjustments((prev: Adjustments) => ({ + ...prev, + relightLights: (prev.relightLights ?? []).map((light) => + light.id === id ? { ...light, ...changes } : light, + ), + })); + }, + [setAdjustments], + ); const isRelightHandleVisible = !!adjustments.relightEnabled && !!adjustments.relightNormalMap && @@ -3195,15 +3218,19 @@ const ImageCanvas = memo(
- {isRelightHandleVisible && ( - - )} + {isRelightHandleVisible && + relightLights.map((light) => ( + 1 && (activeRelightLightId ?? relightLights[0].id) === light.id} + geometry={relightGeometry} + imageRenderSize={imageRenderSize} + inverseScale={1 / maxSafeScale} + onSelect={handleRelightLightSelect} + onChange={handleRelightLightChange} + /> + ))} {!isDrawing.current && directPatchMarkers.map((m) => { diff --git a/src/components/panel/editor/overlays/RelightHandle.tsx b/src/components/panel/editor/overlays/RelightHandle.tsx index 8619f8bf5f..6f8d31667e 100644 --- a/src/components/panel/editor/overlays/RelightHandle.tsx +++ b/src/components/panel/editor/overlays/RelightHandle.tsx @@ -1,15 +1,17 @@ import { useCallback, useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; -import { Adjustments, Coord } from '../../../../utils/adjustments'; +import { Coord, RelightLight } from '../../../../utils/adjustments'; import { RenderSize } from '../../../../hooks/useImageRenderSize'; import { RelightGeometry, relightDisplayToUv, relightUvToDisplay } from '../../../../utils/relightUtils'; interface RelightHandleProps { - adjustments: Adjustments; + light: RelightLight; + isActive: boolean; geometry: RelightGeometry; imageRenderSize: RenderSize; inverseScale: number; - setAdjustments(fn: (prev: Adjustments) => Adjustments): void; + onSelect(id: string): void; + onChange(id: string, changes: Partial): void; } const HANDLE_SIZE = 30; @@ -18,18 +20,23 @@ const WHEEL_DEPTH_STEP = 4; const clamp01 = (value: number) => Math.max(0, Math.min(1, value)); export default function RelightHandle({ - adjustments, + light, + isActive, geometry, imageRenderSize, inverseScale, - setAdjustments, + onSelect, + onChange, }: RelightHandleProps) { const { t } = useTranslation(); const handleRef = useRef(null); const frameRef = useRef(null); const pendingUvRef = useRef(null); + const depthRef = useRef(light.depth); const [dragUv, setDragUv] = useState(null); + depthRef.current = light.depth; + useEffect( () => () => { if (frameRef.current !== null) cancelAnimationFrame(frameRef.current); @@ -46,15 +53,15 @@ export default function RelightHandle({ e.preventDefault(); e.stopPropagation(); const step = e.deltaY < 0 ? WHEEL_DEPTH_STEP : -WHEEL_DEPTH_STEP; - setAdjustments((prev: Adjustments) => ({ - ...prev, - relightDepth: Math.max(-100, Math.min(100, (prev.relightDepth ?? 30) + step)), - })); + const depth = Math.max(-100, Math.min(100, depthRef.current + step)); + depthRef.current = depth; + onSelect(light.id); + onChange(light.id, { depth }); }; handle.addEventListener('wheel', handleWheel, { passive: false }); return () => handle.removeEventListener('wheel', handleWheel); - }, [setAdjustments]); + }, [light.id, onSelect, onChange]); const commitUv = useCallback( (uv: Coord) => { @@ -64,10 +71,10 @@ export default function RelightHandle({ frameRef.current = null; const next = pendingUvRef.current; if (!next) return; - setAdjustments((prev: Adjustments) => ({ ...prev, relightX: next.x, relightY: next.y })); + onChange(light.id, { x: next.x, y: next.y }); }); }, - [setAdjustments], + [light.id, onChange], ); const pointerToUv = useCallback( @@ -97,7 +104,8 @@ export default function RelightHandle({ e.stopPropagation(); e.preventDefault(); e.currentTarget.setPointerCapture(e.pointerId); - setDragUv({ x: adjustments.relightX ?? 0.3, y: adjustments.relightY ?? 0.3 }); + onSelect(light.id); + setDragUv({ x: light.x, y: light.y }); }; const handlePointerMove = (e: React.PointerEvent) => { @@ -118,13 +126,14 @@ export default function RelightHandle({ setDragUv(null); }; - const uv = dragUv ?? { x: adjustments.relightX ?? 0.3, y: adjustments.relightY ?? 0.3 }; + const uv = dragUv ?? { x: light.x, y: light.y }; const display = relightUvToDisplay(uv, geometry); - const color = adjustments.relightColor || '#ffffff'; - const depth = Math.max(-100, Math.min(100, adjustments.relightDepth ?? 30)); + const color = light.color || '#ffffff'; + const depth = Math.max(-100, Math.min(100, light.depth)); // Closer lights look bigger; a light behind the subject gets a dashed outline. const depthScale = 0.7 + 0.6 * ((depth + 100) / 200); const isBehind = depth < 0; + const outline = isActive ? '0 0 0 2.5px var(--color-accent, #ffffff), ' : ''; return (
); diff --git a/src/i18n/locales/ca.json b/src/i18n/locales/ca.json index fd273dbce6..4a93844cfd 100644 --- a/src/i18n/locales/ca.json +++ b/src/i18n/locales/ca.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Punt mitjà", "relight": "Reil·luminar", + "relightAddLight": "Afegeix una llum", "relightAmbient": "Ambient", "relightColor": "Color de la llum", "relightCustomColor": "Color personalitzat", @@ -135,9 +136,11 @@ "relightDragLight": "Arrossega per moure la llum, fes servir la roda per acostar-la o allunyar-la", "relightGenerateNormalMap": "Genera el mapa de normals", "relightIntensity": "Intensitat", + "relightLights": "Llums", "relightNormalMap": "Mapa de normals", "relightRange": "Abast", "relightRegenerateNormalMap": "Regenera el mapa de normals", + "relightRemoveLight": "Elimina la llum seleccionada", "relightSoftness": "Suavitat", "relightSpecular": "Especular", "roughness": "Rugositat", diff --git a/src/i18n/locales/de.json b/src/i18n/locales/de.json index 63fd49fa2c..fffcfd10b2 100644 --- a/src/i18n/locales/de.json +++ b/src/i18n/locales/de.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Mittelpunkt", "relight": "Neu beleuchten", + "relightAddLight": "Licht hinzufügen", "relightAmbient": "Umgebungslicht", "relightColor": "Lichtfarbe", "relightCustomColor": "Eigene Farbe", @@ -135,9 +136,11 @@ "relightDragLight": "Ziehen, um das Licht zu bewegen, scrollen, um es näher oder weiter weg zu setzen", "relightGenerateNormalMap": "Normal Map erzeugen", "relightIntensity": "Intensität", + "relightLights": "Lichter", "relightNormalMap": "Normal Map", "relightRange": "Reichweite", "relightRegenerateNormalMap": "Normal Map neu erzeugen", + "relightRemoveLight": "Ausgewähltes Licht entfernen", "relightSoftness": "Weichheit", "relightSpecular": "Glanzlicht", "roughness": "Rauheit", diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index 9a3cf78311..272e8e0a36 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Midpoint", "relight": "Relight", + "relightAddLight": "Add light", "relightAmbient": "Ambient", "relightColor": "Light Color", "relightCustomColor": "Custom Color", @@ -135,9 +136,11 @@ "relightDragLight": "Drag to move the light, scroll to move it closer or farther", "relightGenerateNormalMap": "Generate Normal Map", "relightIntensity": "Intensity", + "relightLights": "Lights", "relightNormalMap": "Normal Map", "relightRange": "Range", "relightRegenerateNormalMap": "Regenerate Normal Map", + "relightRemoveLight": "Remove selected light", "relightSoftness": "Softness", "relightSpecular": "Specular", "roughness": "Roughness", diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index 34d3a837b5..2edf67ed65 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Punto medio", "relight": "Reiluminar", + "relightAddLight": "Añadir luz", "relightAmbient": "Ambiente", "relightColor": "Color de la luz", "relightCustomColor": "Color personalizado", @@ -135,9 +136,11 @@ "relightDragLight": "Arrastra para mover la luz, usa la rueda para acercarla o alejarla", "relightGenerateNormalMap": "Generar mapa de normales", "relightIntensity": "Intensidad", + "relightLights": "Luces", "relightNormalMap": "Mapa de normales", "relightRange": "Alcance", "relightRegenerateNormalMap": "Regenerar mapa de normales", + "relightRemoveLight": "Quitar la luz seleccionada", "relightSoftness": "Suavidad", "relightSpecular": "Especular", "roughness": "Rugosidad", diff --git a/src/i18n/locales/fr.json b/src/i18n/locales/fr.json index 40590442c5..ccac250d2a 100644 --- a/src/i18n/locales/fr.json +++ b/src/i18n/locales/fr.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Point médian", "relight": "Rééclairage", + "relightAddLight": "Ajouter une lumière", "relightAmbient": "Ambiance", "relightColor": "Couleur de la lumière", "relightCustomColor": "Couleur personnalisée", @@ -135,9 +136,11 @@ "relightDragLight": "Faites glisser pour déplacer la lumière, utilisez la molette pour la rapprocher ou l'éloigner", "relightGenerateNormalMap": "Générer la carte de normales", "relightIntensity": "Intensité", + "relightLights": "Lumières", "relightNormalMap": "Carte de normales", "relightRange": "Portée", "relightRegenerateNormalMap": "Régénérer la carte de normales", + "relightRemoveLight": "Supprimer la lumière sélectionnée", "relightSoftness": "Douceur", "relightSpecular": "Spéculaire", "roughness": "Rugosité", diff --git a/src/i18n/locales/it.json b/src/i18n/locales/it.json index 6cd4288d59..3830853d85 100644 --- a/src/i18n/locales/it.json +++ b/src/i18n/locales/it.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Punto Medio", "relight": "Reilluminazione", + "relightAddLight": "Aggiungi luce", "relightAmbient": "Ambiente", "relightColor": "Colore della luce", "relightCustomColor": "Colore personalizzato", @@ -135,9 +136,11 @@ "relightDragLight": "Trascina per spostare la luce, usa la rotellina per avvicinarla o allontanarla", "relightGenerateNormalMap": "Genera mappa delle normali", "relightIntensity": "Intensità", + "relightLights": "Luci", "relightNormalMap": "Mappa delle normali", "relightRange": "Portata", "relightRegenerateNormalMap": "Rigenera mappa delle normali", + "relightRemoveLight": "Rimuovi la luce selezionata", "relightSoftness": "Morbidezza", "relightSpecular": "Speculare", "roughness": "Ruvidezza", diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index 52f967425f..ac4a462dc6 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "中間点", "relight": "リライト", + "relightAddLight": "ライトを追加", "relightAmbient": "環境光", "relightColor": "ライトの色", "relightCustomColor": "カスタムカラー", @@ -135,9 +136,11 @@ "relightDragLight": "ドラッグでライトを移動、スクロールで前後に移動", "relightGenerateNormalMap": "法線マップを生成", "relightIntensity": "強度", + "relightLights": "ライト", "relightNormalMap": "法線マップ", "relightRange": "範囲", "relightRegenerateNormalMap": "法線マップを再生成", + "relightRemoveLight": "選択したライトを削除", "relightSoftness": "柔らかさ", "relightSpecular": "スペキュラ", "roughness": "粗さ", diff --git a/src/i18n/locales/ko.json b/src/i18n/locales/ko.json index 7eb16f2cf5..29a0957280 100644 --- a/src/i18n/locales/ko.json +++ b/src/i18n/locales/ko.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "중간점", "relight": "재조명", + "relightAddLight": "조명 추가", "relightAmbient": "주변광", "relightColor": "조명 색상", "relightCustomColor": "사용자 지정 색상", @@ -135,9 +136,11 @@ "relightDragLight": "드래그하여 조명 이동, 스크롤하여 앞뒤로 이동", "relightGenerateNormalMap": "노멀 맵 생성", "relightIntensity": "강도", + "relightLights": "조명", "relightNormalMap": "노멀 맵", "relightRange": "범위", "relightRegenerateNormalMap": "노멀 맵 다시 생성", + "relightRemoveLight": "선택한 조명 제거", "relightSoftness": "부드러움", "relightSpecular": "스페큘러", "roughness": "거칠기", diff --git a/src/i18n/locales/pl.json b/src/i18n/locales/pl.json index b069a3b67a..ca4ca0eee5 100644 --- a/src/i18n/locales/pl.json +++ b/src/i18n/locales/pl.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Punkt środkowy", "relight": "Ponowne oświetlenie", + "relightAddLight": "Dodaj światło", "relightAmbient": "Światło otoczenia", "relightColor": "Kolor światła", "relightCustomColor": "Własny kolor", @@ -135,9 +136,11 @@ "relightDragLight": "Przeciągnij, aby przesunąć światło, przewiń, aby je przybliżyć lub oddalić", "relightGenerateNormalMap": "Generuj mapę normalnych", "relightIntensity": "Intensywność", + "relightLights": "Światła", "relightNormalMap": "Mapa normalnych", "relightRange": "Zasięg", "relightRegenerateNormalMap": "Generuj ponownie mapę normalnych", + "relightRemoveLight": "Usuń wybrane światło", "relightSoftness": "Miękkość", "relightSpecular": "Odbicia", "roughness": "Szorstkość", diff --git a/src/i18n/locales/pt.json b/src/i18n/locales/pt.json index a8f89fbaf0..cb11474b5b 100644 --- a/src/i18n/locales/pt.json +++ b/src/i18n/locales/pt.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Ponto Médio", "relight": "Reiluminar", + "relightAddLight": "Adicionar luz", "relightAmbient": "Ambiente", "relightColor": "Cor da luz", "relightCustomColor": "Cor personalizada", @@ -135,9 +136,11 @@ "relightDragLight": "Arraste para mover a luz, use a roda para a aproximar ou afastar", "relightGenerateNormalMap": "Gerar mapa de normais", "relightIntensity": "Intensidade", + "relightLights": "Luzes", "relightNormalMap": "Mapa de normais", "relightRange": "Alcance", "relightRegenerateNormalMap": "Regenerar mapa de normais", + "relightRemoveLight": "Remover a luz selecionada", "relightSoftness": "Suavidade", "relightSpecular": "Especular", "roughness": "Rugosidade", diff --git a/src/i18n/locales/ru.json b/src/i18n/locales/ru.json index bff9c0d5a3..afd3550181 100644 --- a/src/i18n/locales/ru.json +++ b/src/i18n/locales/ru.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "Средняя точка", "relight": "Переосвещение", + "relightAddLight": "Добавить свет", "relightAmbient": "Окружающий свет", "relightColor": "Цвет света", "relightCustomColor": "Свой цвет", @@ -135,9 +136,11 @@ "relightDragLight": "Перетащите, чтобы переместить свет, прокрутите, чтобы приблизить или отдалить его", "relightGenerateNormalMap": "Создать карту нормалей", "relightIntensity": "Интенсивность", + "relightLights": "Источники света", "relightNormalMap": "Карта нормалей", "relightRange": "Дальность", "relightRegenerateNormalMap": "Пересоздать карту нормалей", + "relightRemoveLight": "Удалить выбранный свет", "relightSoftness": "Мягкость", "relightSpecular": "Блики", "roughness": "Шероховатость", diff --git a/src/i18n/locales/zh-CN.json b/src/i18n/locales/zh-CN.json index e95f08a8d6..a6f7b73620 100644 --- a/src/i18n/locales/zh-CN.json +++ b/src/i18n/locales/zh-CN.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "中点", "relight": "重新打光", + "relightAddLight": "添加灯光", "relightAmbient": "环境光", "relightColor": "灯光颜色", "relightCustomColor": "自定义颜色", @@ -135,9 +136,11 @@ "relightDragLight": "拖动以移动灯光，滚动以前后移动", "relightGenerateNormalMap": "生成法线贴图", "relightIntensity": "强度", + "relightLights": "灯光", "relightNormalMap": "法线贴图", "relightRange": "范围", "relightRegenerateNormalMap": "重新生成法线贴图", + "relightRemoveLight": "移除所选灯光", "relightSoftness": "柔和度", "relightSpecular": "高光", "roughness": "粗糙度", diff --git a/src/i18n/locales/zh-TW.json b/src/i18n/locales/zh-TW.json index 80df6f03c0..4606bb2c65 100644 --- a/src/i18n/locales/zh-TW.json +++ b/src/i18n/locales/zh-TW.json @@ -128,6 +128,7 @@ "lut": "LUT", "midpoint": "中點", "relight": "重新打光", + "relightAddLight": "新增燈光", "relightAmbient": "環境光", "relightColor": "燈光顏色", "relightCustomColor": "自訂顏色", @@ -135,9 +136,11 @@ "relightDragLight": "拖曳以移動燈光，捲動以前後移動", "relightGenerateNormalMap": "產生法線貼圖", "relightIntensity": "強度", + "relightLights": "燈光", "relightNormalMap": "法線貼圖", "relightRange": "範圍", "relightRegenerateNormalMap": "重新產生法線貼圖", + "relightRemoveLight": "移除所選燈光", "relightSoftness": "柔和度", "relightSpecular": "高光", "roughness": "粗糙度", diff --git a/src/store/useEditorStore.ts b/src/store/useEditorStore.ts index 441df08fac..a6cf7da75c 100644 --- a/src/store/useEditorStore.ts +++ b/src/store/useEditorStore.ts @@ -51,6 +51,7 @@ interface EditorState { // Interaction State isSliderDragging: boolean; + activeRelightLightId: string | null; zoom: number; displaySize: ImageDimensions; previewSize: ImageDimensions; @@ -114,6 +115,7 @@ export const useEditorStore = create((set) => ({ waveformHeight: 220, isSliderDragging: false, + activeRelightLightId: null, interactivePatch: null, activeMaskContainerId: null, activeMaskId: null, diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index deab44a812..6bcf17880e 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -97,15 +97,10 @@ export enum Effect { LensBlurMinFade = 'lensBlurMinFade', LensBlurMaxFade = 'lensBlurMaxFade', RelightEnabled = 'relightEnabled', - RelightIntensity = 'relightIntensity', - RelightX = 'relightX', - RelightY = 'relightY', - RelightDepth = 'relightDepth', - RelightRange = 'relightRange', + RelightLights = 'relightLights', RelightAmbient = 'relightAmbient', RelightSoftness = 'relightSoftness', RelightSpecular = 'relightSpecular', - RelightColor = 'relightColor', RelightNormalMap = 'relightNormalMap', RelightDepthMap = 'relightDepthMap', RelightDepthScale = 'relightDepthScale', @@ -257,18 +252,13 @@ export interface Adjustments { masks: Array; orientationSteps: number; relightAmbient: number; - relightColor: string; - relightDepth: number; relightDepthMap: string | null; relightDepthScale: number; relightEnabled: boolean; - relightIntensity: number; + relightLights: Array; relightNormalMap: string | null; - relightRange: number; relightSoftness: number; relightSpecular: number; - relightX: number; - relightY: number; rotation: number; saturation: number; sectionVisibility: SectionVisibility; @@ -538,6 +528,52 @@ export const INITIAL_MASK_CONTAINER: MaskContainer = { visible: true, }; +export interface RelightLight { + id: string; + // Position as a fraction of the image before orientation, flip, rotation and crop. + x: number; + y: number; + // -100 (far behind the subject) to 100 (close to the camera). + depth: number; + intensity: number; + range: number; + color: string; +} + +export const MAX_RELIGHT_LIGHTS = 6; + +export const DEFAULT_RELIGHT_LIGHT: RelightLight = { + id: 'light-1', + x: 0.3, + y: 0.3, + depth: 30, + intensity: 50, + range: 60, + color: '#ffffff', +}; + +// Edits saved before multiple lights existed describe a single light at the top level. +const normalizeRelightLights = (loaded: any): Array => { + if (Array.isArray(loaded.relightLights)) { + return loaded.relightLights.map((light: Partial, index: number) => ({ + ...DEFAULT_RELIGHT_LIGHT, + ...light, + id: light.id || `light-${index + 1}`, + })); + } + return [ + { + ...DEFAULT_RELIGHT_LIGHT, + x: loaded.relightX ?? DEFAULT_RELIGHT_LIGHT.x, + y: loaded.relightY ?? DEFAULT_RELIGHT_LIGHT.y, + depth: loaded.relightDepth ?? DEFAULT_RELIGHT_LIGHT.depth, + intensity: loaded.relightIntensity ?? DEFAULT_RELIGHT_LIGHT.intensity, + range: loaded.relightRange ?? DEFAULT_RELIGHT_LIGHT.range, + color: loaded.relightColor ?? DEFAULT_RELIGHT_LIGHT.color, + }, + ]; +}; + export const INITIAL_ADJUSTMENTS: Adjustments = { aiPatches: [], aspectRatio: null, @@ -608,18 +644,13 @@ export const INITIAL_ADJUSTMENTS: Adjustments = { masks: [], orientationSteps: 0, relightAmbient: 80, - relightColor: '#ffffff', - relightDepth: 30, relightDepthMap: null, relightDepthScale: 0, relightEnabled: false, - relightIntensity: 50, + relightLights: [{ ...DEFAULT_RELIGHT_LIGHT }], relightNormalMap: null, - relightRange: 60, relightSoftness: 30, relightSpecular: 0, - relightX: 0.3, - relightY: 0.3, rotation: 0, saturation: 0, sectionVisibility: { @@ -768,18 +799,13 @@ export const normalizeLoadedAdjustments = (loadedAdjustments: Adjustments): any lensBlurMinDepth: loadedAdjustments.lensBlurMinDepth ?? INITIAL_ADJUSTMENTS.lensBlurMinDepth, lensBlurMinFade: loadedAdjustments.lensBlurMinFade ?? INITIAL_ADJUSTMENTS.lensBlurMinFade, relightAmbient: loadedAdjustments.relightAmbient ?? INITIAL_ADJUSTMENTS.relightAmbient, - relightColor: loadedAdjustments.relightColor ?? INITIAL_ADJUSTMENTS.relightColor, - relightDepth: loadedAdjustments.relightDepth ?? INITIAL_ADJUSTMENTS.relightDepth, relightDepthMap: loadedAdjustments.relightDepthMap ?? INITIAL_ADJUSTMENTS.relightDepthMap, relightDepthScale: loadedAdjustments.relightDepthScale ?? INITIAL_ADJUSTMENTS.relightDepthScale, relightEnabled: loadedAdjustments.relightEnabled ?? INITIAL_ADJUSTMENTS.relightEnabled, - relightIntensity: loadedAdjustments.relightIntensity ?? INITIAL_ADJUSTMENTS.relightIntensity, + relightLights: normalizeRelightLights(loadedAdjustments), relightNormalMap: loadedAdjustments.relightNormalMap ?? INITIAL_ADJUSTMENTS.relightNormalMap, - relightRange: loadedAdjustments.relightRange ?? INITIAL_ADJUSTMENTS.relightRange, relightSoftness: loadedAdjustments.relightSoftness ?? INITIAL_ADJUSTMENTS.relightSoftness, relightSpecular: loadedAdjustments.relightSpecular ?? INITIAL_ADJUSTMENTS.relightSpecular, - relightX: loadedAdjustments.relightX ?? INITIAL_ADJUSTMENTS.relightX, - relightY: loadedAdjustments.relightY ?? INITIAL_ADJUSTMENTS.relightY, lensCorrectionMode: loadedAdjustments.lensCorrectionMode || 'manual', lensMaker: loadedAdjustments.lensMaker ?? INITIAL_ADJUSTMENTS.lensMaker, lensModel: loadedAdjustments.lensModel ?? INITIAL_ADJUSTMENTS.lensModel, @@ -1005,15 +1031,10 @@ export const ADJUSTMENT_SECTIONS: Sections = { Effect.LensBlurMinFade, Effect.LensBlurMaxFade, Effect.RelightEnabled, - Effect.RelightIntensity, - Effect.RelightX, - Effect.RelightY, - Effect.RelightDepth, - Effect.RelightRange, + Effect.RelightLights, Effect.RelightAmbient, Effect.RelightSoftness, Effect.RelightSpecular, - Effect.RelightColor, Effect.RelightNormalMap, Effect.RelightDepthMap, Effect.RelightDepthScale, From 25641b2790d52cf89df8c5c5df21d69c8f8204eb Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Mon, 5 Oct 2026 16:55:25 -0500 Subject: [PATCH 07/16] feat: implement AI features including depth mapping, image segmentation, and localizations --- src-tauri/src/ai_commands.rs | 38 +++ src-tauri/src/ai_processing.rs | 125 ++++--- src-tauri/src/cache_utils.rs | 1 + src-tauri/src/lib.rs | 1 + src-tauri/src/relight.rs | 357 ++++++++++++++------ src/components/adjustments/Effects.tsx | 57 +++- src/components/panel/editor/ImageCanvas.tsx | 64 ++++ src/i18n/locales/ca.json | 6 + src/i18n/locales/de.json | 6 + src/i18n/locales/en.json | 6 + src/i18n/locales/es.json | 6 + src/i18n/locales/fr.json | 6 + src/i18n/locales/it.json | 6 + src/i18n/locales/ja.json | 6 + src/i18n/locales/ko.json | 6 + src/i18n/locales/pl.json | 6 + src/i18n/locales/pt.json | 6 + src/i18n/locales/ru.json | 6 + src/i18n/locales/zh-CN.json | 6 + src/i18n/locales/zh-TW.json | 6 + src/store/useEditorStore.ts | 2 + src/utils/adjustments.ts | 13 + 22 files changed, 589 insertions(+), 147 deletions(-) diff --git a/src-tauri/src/ai_commands.rs b/src-tauri/src/ai_commands.rs index 54d4ccbf7b..f6ccd5f689 100644 --- a/src-tauri/src/ai_commands.rs +++ b/src-tauri/src/ai_commands.rs @@ -321,6 +321,7 @@ fn encode_png_data_url(image: &image::DynamicImage) -> Result { #[tauri::command] pub async fn generate_relight_maps( js_adjustments: serde_json::Value, + quality: Option, state: tauri::State<'_, AppState>, app_handle: tauri::AppHandle, ) -> Result { @@ -328,6 +329,7 @@ pub async fn generate_relight_maps( &app_handle, &state.ai_state, &state.ai_init_lock, + quality.unwrap_or_default(), ) .await .map_err(|e| e.to_string())?; @@ -344,6 +346,42 @@ pub async fn generate_relight_maps( }) } +const RELIGHT_PREVIEW_MAX_SIZE: u32 = 2560; + +/// Normal map as the relight pass uses it, laid out like the edited image (orientation, +/// rotation and crop applied) so it can be shown on top of it. +#[tauri::command] +pub async fn generate_relight_normal_preview( + js_adjustments: serde_json::Value, + state: tauri::State<'_, AppState>, +) -> Result { + let warped_image = crate::get_cached_full_warped_image(&state, &js_adjustments)?; + + let normals = crate::relight::render_normal_preview(warped_image.as_ref(), &js_adjustments) + .ok_or_else(|| "No normal map available".to_string())?; + + let (transformed, _) = crate::adjustment_utils::apply_spatial_transformations( + image::DynamicImage::ImageRgb8(normals), + &js_adjustments, + ); + let preview = if transformed.width().max(transformed.height()) > RELIGHT_PREVIEW_MAX_SIZE { + transformed.resize( + RELIGHT_PREVIEW_MAX_SIZE, + RELIGHT_PREVIEW_MAX_SIZE, + image::imageops::FilterType::Triangle, + ) + } else { + transformed.into_owned() + }; + + let mut buf = Cursor::new(Vec::new()); + image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buf, 92) + .encode_image(&preview.to_rgb8()) + .map_err(|e| e.to_string())?; + let base64_str = general_purpose::STANDARD.encode(buf.get_ref()); + Ok(format!("data:image/jpeg;base64,{}", base64_str)) +} + #[allow(clippy::too_many_arguments)] #[tauri::command] pub async fn generate_ai_subject_mask( diff --git a/src-tauri/src/ai_processing.rs b/src-tauri/src/ai_processing.rs index b2016324fc..d18fafe10d 100644 --- a/src-tauri/src/ai_processing.rs +++ b/src-tauri/src/ai_processing.rs @@ -60,9 +60,42 @@ const DEPTH_FILENAME: &str = "depth_anything_v2_vits.onnx"; const DEPTH_INPUT_SIZE: u32 = 518; const DEPTH_SHA256: &str = "d2b11a11c1d4a12b47608fa65a17ee9a4c605b55ee1730c8e3b526304f2562be"; -const NORMAL_URL: &str = "https://huggingface.co/Ruicheng/moge-2-vits-normal-onnx/resolve/e50ffda41565591092adea54c6ac83d6212e1e23/model.onnx?download=true"; -const NORMAL_FILENAME: &str = "moge2_vits_normal.onnx"; -const NORMAL_SHA256: &str = "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0"; +/// MoGe-2 checkpoints used for the relight maps. Both take the same inputs and return the +/// same outputs; the large one resolves much finer geometry but is far heavier. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum NormalModelQuality { + #[default] + Standard, + High, +} + +struct NormalModelSpec { + url: &'static str, + filename: &'static str, + sha256: &'static str, + name: &'static str, +} + +impl NormalModelQuality { + fn spec(self) -> NormalModelSpec { + match self { + Self::Standard => NormalModelSpec { + url: "https://huggingface.co/Ruicheng/moge-2-vits-normal-onnx/resolve/e50ffda41565591092adea54c6ac83d6212e1e23/model.onnx?download=true", + filename: "moge2_vits_normal.onnx", + sha256: "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0", + name: "Normal Map Model", + }, + Self::High => NormalModelSpec { + url: "https://huggingface.co/Ruicheng/moge-2-vitl-normal-onnx/resolve/1c02388b024dd34da413e1a934d7a941785f1460/model.onnx?download=true", + filename: "moge2_vitl_normal.onnx", + sha256: "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751", + name: "Normal Map Model (High Quality, 1.3 GB)", + }, + } + } +} + const NORMAL_MAX_SIZE: u32 = 1024; const NORMAL_NUM_TOKENS: i64 = 3600; @@ -98,7 +131,8 @@ pub struct AiState { pub denoise_model: Option>>, pub clip_models: Option>, pub lama_model: Option>>, - pub normal_model: Option>>, + // Only one checkpoint is kept loaded: the large one alone takes well over 1 GB of RAM. + pub normal_model: Option<(NormalModelQuality, Arc>)>, pub embeddings: Option, pub depth_map: Option, } @@ -321,13 +355,8 @@ fn get_models_dir(app_handle: &tauri::AppHandle) -> Result { Ok(models_dir) } -fn persist_downloaded_asset(dest: &Path, bytes: &[u8]) -> Result<()> { - if bytes.is_empty() { - return Err(anyhow::anyhow!( - "Downloaded asset for {} was empty", - dest.display() - )); - } +async fn download_model(url: &str, dest: &Path) -> Result<()> { + let mut response = reqwest::get(url).await?.error_for_status()?; let parent = dest.parent().ok_or_else(|| { anyhow::anyhow!( @@ -343,10 +372,29 @@ fn persist_downloaded_asset(dest: &Path, bytes: &[u8]) -> Result<()> { .ok_or_else(|| anyhow::anyhow!("Invalid downloaded asset path {}", dest.display()))?; let tmp_path = dest.with_file_name(format!(".{}.download", file_name)); - { + // Written in chunks so that large models never have to fit in memory. + let mut total_bytes = 0usize; + let write_result: Result<()> = async { let mut file = fs::File::create(&tmp_path)?; - file.write_all(bytes)?; + while let Some(chunk) = response.chunk().await? { + file.write_all(&chunk)?; + total_bytes += chunk.len(); + } file.sync_all()?; + Ok(()) + } + .await; + + if let Err(error) = write_result { + let _ = fs::remove_file(&tmp_path); + return Err(error); + } + if total_bytes == 0 { + let _ = fs::remove_file(&tmp_path); + return Err(anyhow::anyhow!( + "Downloaded asset for {} was empty", + dest.display() + )); } fs::rename(&tmp_path, dest).or_else(|rename_error| -> std::io::Result<()> { @@ -361,12 +409,6 @@ fn persist_downloaded_asset(dest: &Path, bytes: &[u8]) -> Result<()> { Ok(()) } -async fn download_model(url: &str, dest: &Path) -> Result<()> { - let response = reqwest::get(url).await?.error_for_status()?; - let bytes = response.bytes().await?; - persist_downloaded_asset(dest, &bytes) -} - fn verify_sha256(path: &Path, expected_hash: &str) -> Result { if !path.exists() { return Ok(false); @@ -765,40 +807,47 @@ pub async fn get_or_init_normal_model( app_handle: &tauri::AppHandle, ai_state_mutex: &Mutex>, ai_init_lock: &TokioMutex<()>, + quality: NormalModelQuality, ) -> Result>> { - if let Some(normal_model) = ai_state_mutex - .lock() - .unwrap() - .as_ref() - .and_then(|state| state.normal_model.clone()) - { + let loaded_model = || { + ai_state_mutex + .lock() + .unwrap() + .as_ref() + .and_then(|state| state.normal_model.clone()) + .filter(|(loaded_quality, _)| *loaded_quality == quality) + .map(|(_, model)| model) + }; + + if let Some(normal_model) = loaded_model() { return Ok(normal_model); } let _guard = ai_init_lock.lock().await; - if let Some(normal_model) = ai_state_mutex - .lock() - .unwrap() - .as_ref() - .and_then(|state| state.normal_model.clone()) - { + if let Some(normal_model) = loaded_model() { return Ok(normal_model); } + let spec = quality.spec(); let models_dir = get_models_dir(app_handle)?; download_and_verify_model( app_handle, &models_dir, - NORMAL_FILENAME, - NORMAL_URL, - NORMAL_SHA256, - "Normal Map Model", + spec.filename, + spec.url, + spec.sha256, + spec.name, ) .await?; + // Free the other checkpoint before loading this one. + if let Some(state) = ai_state_mutex.lock().unwrap().as_mut() { + state.normal_model = None; + } + let _ = ort::init().with_name("AI-Normal").commit(); - let model_path = models_dir.join(NORMAL_FILENAME); + let model_path = models_dir.join(spec.filename); let session = Session::builder()?.commit_from_file(model_path)?; let normal_model = Arc::new(Mutex::new(session)); @@ -806,14 +855,14 @@ pub async fn get_or_init_normal_model( let mut ai_state_lock = ai_state_mutex.lock().unwrap(); if let Some(state) = ai_state_lock.as_mut() { - state.normal_model = Some(normal_model.clone()); + state.normal_model = Some((quality, normal_model.clone())); } else { *ai_state_lock = Some(AiState { models: None, denoise_model: None, clip_models: None, lama_model: None, - normal_model: Some(normal_model.clone()), + normal_model: Some((quality, normal_model.clone())), embeddings: None, depth_map: None, }); diff --git a/src-tauri/src/cache_utils.rs b/src-tauri/src/cache_utils.rs index 6f86bafe44..b5f81e4b29 100644 --- a/src-tauri/src/cache_utils.rs +++ b/src-tauri/src/cache_utils.rs @@ -103,6 +103,7 @@ fn hash_relight( "relightAmbient", "relightSoftness", "relightSpecular", + "relightDetail", "relightNormalMap", "relightDepthMap", "relightDepthScale", diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 1d94907125..1f54bb14ec 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2181,6 +2181,7 @@ pub fn run() { ai_commands::test_ai_connector_connection, ai_commands::generate_full_image_depth_map, ai_commands::generate_relight_maps, + ai_commands::generate_relight_normal_preview, ai_commands::cancel_ai_task, apple_raw::is_raw9_available, inpainting::invoke_generative_replace_with_mask_def, diff --git a/src-tauri/src/relight.rs b/src-tauri/src/relight.rs index 1532c1ea17..e4092d95e0 100644 --- a/src-tauri/src/relight.rs +++ b/src-tauri/src/relight.rs @@ -13,6 +13,8 @@ const MAX_FRONT_DISTANCE: f32 = 1.5; // Below this size there is nothing to gain from snapping the maps to the image edges. const MIN_REFINE_SIZE: u32 = 64; const MAX_LIGHTS: usize = 8; +// Brightness slope to surface tilt, for a detail amount of 100. +const DETAIL_GAIN: f32 = 2.0; type DepthMap = ImageBuffer, Vec>; @@ -21,6 +23,30 @@ struct SurfaceMaps { normals: Vec<[i8; 3]>, coverage: Vec, depth: Vec, + // Fine relief (hair, pores, fabric) read from the photo's own texture, as an X / Y tilt + // to add to the normal. The AI maps are too coarse to carry it. + detail: Vec<[i8; 2]>, +} + +impl SurfaceMaps { + #[inline(always)] + fn normal(&self, i: usize, detail_strength: f32) -> (f32, f32, f32) { + let n = self.normals[i]; + let (nx, ny, nz) = ( + n[0] as f32 / 127.0, + n[1] as f32 / 127.0, + n[2] as f32 / 127.0, + ); + if detail_strength <= 0.0 { + return (nx, ny, nz); + } + + let d = self.detail[i]; + let nx = nx + d[0] as f32 / 127.0 * detail_strength; + let ny = ny + d[1] as f32 / 127.0 * detail_strength; + let inv_len = 1.0 / (nx * nx + ny * ny + nz * nz).sqrt().max(1e-6); + (nx * inv_len, ny * inv_len, nz * inv_len) + } } // Building the maps only depends on the photo and the AI maps, never on the light, so the @@ -47,6 +73,7 @@ struct RelightParams { ambient: f32, wrap: f32, specular: f32, + detail: f32, } pub fn apply_relight<'a>( @@ -63,11 +90,6 @@ pub fn apply_relight<'a>( return image; } - let normal_b64 = adjustments["relightNormalMap"].as_str().unwrap_or(""); - if normal_b64.is_empty() { - return image; - } - let (w, h) = image.dimensions(); if w < 2 || h < 2 { return image; @@ -76,9 +98,65 @@ pub fn apply_relight<'a>( let start = std::time::Instant::now(); let mut out = image.as_ref().to_rgb32f(); + let Some(maps) = cached_surface_maps(&out, adjustments) else { + return image; + }; + + let has_depth_map = !adjustments["relightDepthMap"] + .as_str() + .unwrap_or("") + .is_empty(); + let params = read_params(adjustments, has_depth_map); + relight_in_place(&mut out, &maps, ¶ms); + + log::info!("relight ({}x{}) took {:.2?}", w, h, start.elapsed()); + + Cow::Owned(DynamicImage::ImageRgb32F(out)) +} + +/// Renders the normal map that the relight pass actually uses for `image` (scaled up, +/// snapped to the photo's edges and with the detail amount applied), for inspection. +pub fn render_normal_preview( + image: &DynamicImage, + adjustments: &serde_json::Value, +) -> Option { + let (w, h) = image.dimensions(); + if w < 2 || h < 2 { + return None; + } + + let source = image.to_rgb32f(); + let maps = cached_surface_maps(&source, adjustments)?; + drop(source); + + let detail = read_params(adjustments, false).detail; + // Areas without geometry (sky) are shown dark instead of as a made-up normal. + const EMPTY: f32 = 32.0; + + let mut raw = vec![0u8; w as usize * h as usize * 3]; + raw.par_chunks_exact_mut(3).enumerate().for_each(|(i, px)| { + let (nx, ny, nz) = maps.normal(i, detail); + let coverage = maps.coverage[i] as f32 / 255.0; + for (c, n) in [nx, ny, nz].into_iter().enumerate() { + let packed = (n * 0.5 + 0.5).clamp(0.0, 1.0) * 255.0; + px[c] = (EMPTY + (packed - EMPTY) * coverage).round() as u8; + } + }); + + image::RgbImage::from_raw(w, h, raw) +} + +fn cached_surface_maps( + image: &Rgb32FImage, + adjustments: &serde_json::Value, +) -> Option> { + let normal_b64 = adjustments["relightNormalMap"].as_str().unwrap_or(""); + if normal_b64.is_empty() { + return None; + } let depth_b64 = adjustments["relightDepthMap"].as_str().unwrap_or(""); let depth_scale = adjustments["relightDepthScale"].as_f64().unwrap_or(0.0) as f32; - let cache_key = surface_maps_key(&out, normal_b64, depth_b64, depth_scale); + let cache_key = surface_maps_key(image, normal_b64, depth_b64, depth_scale); let cached = SURFACE_MAPS_CACHE .lock() @@ -86,47 +164,38 @@ pub fn apply_relight<'a>( .as_ref() .filter(|(key, _)| *key == cache_key) .map(|(_, maps)| Arc::clone(maps)); + if cached.is_some() { + return cached; + } - let maps = match cached { - Some(maps) => maps, - None => { - let Some(normal_map) = decode_data_url(normal_b64).map(|img| img.into_rgba8()) else { - return image; - }; - if normal_map.width() < 2 || normal_map.height() < 2 { - return image; - } - let depth_map = decode_data_url(depth_b64) - .map(|img| img.into_luma16()) - .filter(|depth| depth.dimensions() == normal_map.dimensions()); - - let refine = w >= MIN_REFINE_SIZE && h >= MIN_REFINE_SIZE; - let maps = Arc::new(build_surface_maps( - &out, - &normal_map, - depth_map.as_ref(), - depth_scale, - refine, - )); - log::info!( - "relight maps ({}x{}) built in {:.2?}", - w, - h, - start.elapsed() - ); - - *SURFACE_MAPS_CACHE.lock().unwrap_or_else(|e| e.into_inner()) = - Some((cache_key, Arc::clone(&maps))); - maps - } - }; - - let params = read_params(adjustments, !depth_b64.is_empty()); - relight_in_place(&mut out, &maps, ¶ms); - - log::info!("relight ({}x{}) took {:.2?}", w, h, start.elapsed()); + let start = std::time::Instant::now(); + let normal_map = decode_data_url(normal_b64)?.into_rgba8(); + if normal_map.width() < 2 || normal_map.height() < 2 { + return None; + } + let depth_map = decode_data_url(depth_b64) + .map(|img| img.into_luma16()) + .filter(|depth| depth.dimensions() == normal_map.dimensions()); - Cow::Owned(DynamicImage::ImageRgb32F(out)) + let (w, h) = image.dimensions(); + let refine = w >= MIN_REFINE_SIZE && h >= MIN_REFINE_SIZE; + let maps = Arc::new(build_surface_maps( + image, + &normal_map, + depth_map.as_ref(), + depth_scale, + refine, + )); + log::info!( + "relight maps ({}x{}) built in {:.2?}", + w, + h, + start.elapsed() + ); + + *SURFACE_MAPS_CACHE.lock().unwrap_or_else(|e| e.into_inner()) = + Some((cache_key, Arc::clone(&maps))); + Some(maps) } fn decode_data_url(data_url: &str) -> Option { @@ -311,79 +380,115 @@ fn build_surface_maps( let mut normals = vec![[0i8; 3]; w * h]; let mut coverage = vec![0u8; w * h]; let mut depth = vec![0u16; w * h]; + let mut detail = vec![[0i8; 2]; w * h]; + + // A feature that is one pixel wide in a small render spans several pixels in a large + // one, so slopes are measured against a fixed reference size. + let detail_gain = DETAIL_GAIN * w.max(h) as f32 / 1000.0; normals .par_chunks_exact_mut(w) .zip(coverage.par_chunks_exact_mut(w)) .zip(depth.par_chunks_exact_mut(w)) + .zip(detail.par_chunks_exact_mut(w)) .enumerate() - .for_each(|(y, ((normal_row, coverage_row), depth_row))| { - let ty = &y_taps[y]; - let (r0, r1) = (ty.i0 * gw, ty.i1 * gw); - - for x in 0..w { - let tx = &x_taps[x]; - let (i00, i10, i01, i11) = (r0 + tx.i0, r0 + tx.i1, r1 + tx.i0, r1 + tx.i1); - let bilerp = |v00: f32, v10: f32, v01: f32, v11: f32| { - let top = v00 + (v10 - v00) * tx.t; - let bot = v01 + (v11 - v01) * tx.t; - top + (bot - top) * ty.t - }; + .for_each( + |(y, (((normal_row, coverage_row), depth_row), detail_row))| { + let ty = &y_taps[y]; + let (r0, r1) = (ty.i0 * gw, ty.i1 * gw); + + for x in 0..w { + let tx = &x_taps[x]; + let (i00, i10, i01, i11) = (r0 + tx.i0, r0 + tx.i1, r1 + tx.i0, r1 + tx.i1); + let bilerp = |v00: f32, v10: f32, v01: f32, v11: f32| { + let top = v00 + (v10 - v00) * tx.t; + let bot = v01 + (v11 - v01) * tx.t; + top + (bot - top) * ty.t + }; - let mut values = [0.0f32; 5]; - for (c, value) in values.iter_mut().enumerate() { - let plane = &planes[c]; - *value = bilerp(plane[i00], plane[i10], plane[i01], plane[i11]); - } + let mut values = [0.0f32; 5]; + for (c, value) in values.iter_mut().enumerate() { + let plane = &planes[c]; + *value = bilerp(plane[i00], plane[i10], plane[i01], plane[i11]); + } - if refine { - let weight = bilerp( - edge_weight[i00], - edge_weight[i10], - edge_weight[i01], - edge_weight[i11], - ); - if weight > 1e-3 { - let luma = luma_full[y * w + x]; - for (c, value) in values.iter_mut().enumerate() { - let model = &models[c]; - let m = |k: usize| { - bilerp( - model[i00 * 4 + k], - model[i10 * 4 + k], - model[i01 * 4 + k], - model[i11 * 4 + k], - ) - }; - let mut guided = (m(0) * luma + m(1)).clamp(m(2), m(3)).clamp(0.0, 1.0); - if c == 3 { - // Coverage is a matte: push it back towards solid / empty so the - // photo's texture does not leak into it as a halo. - guided = smoothstep(0.25, 0.75, guided); + if refine { + let weight = bilerp( + edge_weight[i00], + edge_weight[i10], + edge_weight[i01], + edge_weight[i11], + ); + + // Treat brightness as height and take its slope (Sobel). A silhouette is + // a change of object rather than relief, so it is left out. + let (xl, xr) = (x.saturating_sub(1), (x + 1).min(w - 1)); + let (yu, yd) = (y.saturating_sub(1) * w, (y + 1).min(h - 1) * w); + let yc = y * w; + let l = |row: usize, col: usize| luma_full[row + col]; + let slope_x = (l(yu, xr) + 2.0 * l(yc, xr) + l(yd, xr) + - l(yu, xl) + - 2.0 * l(yc, xl) + - l(yd, xl)) + * 0.125; + let slope_y = (l(yd, xl) + 2.0 * l(yd, x) + l(yd, xr) + - l(yu, xl) + - 2.0 * l(yu, x) + - l(yu, xr)) + * 0.125; + // The surface leans away from the uphill direction; image Y points down. + let tilt_x = -slope_x * detail_gain; + let tilt_y = slope_y * detail_gain; + let soft_clip = (1.0 - weight).max(0.0) + / (1.0 + (tilt_x * tilt_x + tilt_y * tilt_y).sqrt()); + detail_row[x] = [ + (tilt_x * soft_clip * 127.0).round() as i8, + (tilt_y * soft_clip * 127.0).round() as i8, + ]; + if weight > 1e-3 { + let luma = luma_full[y * w + x]; + for (c, value) in values.iter_mut().enumerate() { + let model = &models[c]; + let m = |k: usize| { + bilerp( + model[i00 * 4 + k], + model[i10 * 4 + k], + model[i01 * 4 + k], + model[i11 * 4 + k], + ) + }; + let mut guided = + (m(0) * luma + m(1)).clamp(m(2), m(3)).clamp(0.0, 1.0); + if c == 3 { + // Coverage is a matte: push it back towards solid / empty so the + // photo's texture does not leak into it as a halo. + guided = smoothstep(0.25, 0.75, guided); + } + *value += (guided - *value) * weight; } - *value += (guided - *value) * weight; } } - } - let nx = values[0] * 2.0 - 1.0; - let ny = values[1] * 2.0 - 1.0; - let nz = values[2] * 2.0 - 1.0; - let scale = 127.0 / (nx * nx + ny * ny + nz * nz).sqrt().max(1e-6); - normal_row[x] = [ - (nx * scale).round() as i8, - (ny * scale).round() as i8, - (nz * scale).round() as i8, - ]; - coverage_row[x] = (values[3].clamp(0.0, 1.0) * 255.0).round() as u8; - depth_row[x] = (values[4].clamp(0.0, 1.0) * 65535.0).round() as u16; - } - }); + let nx = values[0] * 2.0 - 1.0; + let ny = values[1] * 2.0 - 1.0; + let nz = values[2] * 2.0 - 1.0; + let scale = 127.0 / (nx * nx + ny * ny + nz * nz).sqrt().max(1e-6); + normal_row[x] = [ + (nx * scale).round() as i8, + (ny * scale).round() as i8, + (nz * scale).round() as i8, + ]; + coverage_row[x] = (values[3].clamp(0.0, 1.0) * 255.0).round() as u8; + depth_row[x] = (values[4].clamp(0.0, 1.0) * 65535.0).round() as u16; + } + }, + ); SurfaceMaps { normals, coverage, depth, + detail, } } @@ -474,6 +579,7 @@ fn read_params(adjustments: &serde_json::Value, has_depth_map: bool) -> RelightP ambient: (num("relightAmbient", 80.0) / 100.0).clamp(0.0, 1.0), wrap: (num("relightSoftness", 30.0) / 100.0).clamp(0.0, 1.0), specular: (num("relightSpecular", 0.0) / 100.0).clamp(0.0, 1.0), + detail: (num("relightDetail", 30.0) / 100.0).clamp(0.0, 1.0) * 2.0, } } @@ -513,12 +619,7 @@ fn relight_in_place(image: &mut Rgb32FImage, maps: &SurfaceMaps, p: &RelightPara continue; } - let n = maps.normals[i]; - let (nx, ny, nz) = ( - n[0] as f32 / 127.0, - n[1] as f32 / 127.0, - n[2] as f32 / 127.0, - ); + let (nx, ny, nz) = maps.normal(i, p.detail); let surface_z = maps.depth[i] as f32 * depth_to_z; let mut lit = [0.0f32; 3]; @@ -694,6 +795,48 @@ mod tests { assert!(refined > soft * 1.5, "refined {refined} vs bilinear {soft}"); } + #[test] + fn detail_tilts_the_surface_along_the_photo_texture() { + // A bright ridge down the middle of a flat, camera-facing surface. + let image = Rgb32FImage::from_fn(1000, 64, |x, _| { + let v = 0.2 + 0.6 * (-((x as f32 - 500.0) / 6.0).powi(2)).exp(); + Rgb([v, v, v]) + }); + let normal_map = RgbaImage::from_pixel(16, 16, Rgba([128, 128, 255, 255])); + let maps = build_surface_maps(&image, &normal_map, None, 0.0, true); + let row = 32 * 1000; + + // Left of the ridge the surface climbs to the right, so it faces left; and vice versa. + let (left_nx, _, _) = maps.normal(row + 494, 1.0); + let (right_nx, _, _) = maps.normal(row + 506, 1.0); + assert!(left_nx < -0.05, "left side nx {left_nx}"); + assert!(right_nx > 0.05, "right side nx {right_nx}"); + + // Without detail the AI normal is used as is. + let (plain_nx, _, plain_nz) = maps.normal(row + 494, 0.0); + assert!(plain_nx.abs() < 0.02 && plain_nz > 0.98); + } + + #[test] + fn normal_preview_shows_normals_and_darkens_empty_areas() { + let image = DynamicImage::ImageRgb32F(Rgb32FImage::from_pixel(8, 8, Rgb([0.5, 0.5, 0.5]))); + let encode = |alpha: u8| { + let mut buf = std::io::Cursor::new(Vec::new()); + DynamicImage::ImageRgba8(flat_normal_map([128, 128, 255], alpha)) + .write_to(&mut buf, image::ImageFormat::Png) + .unwrap(); + format!("data:image/png;base64,{}", BASE64.encode(buf.get_ref())) + }; + + let solid = render_normal_preview(&image, &json!({ "relightNormalMap": encode(255) })); + let empty = render_normal_preview(&image, &json!({ "relightNormalMap": encode(0) })); + + let solid = solid.unwrap().get_pixel(3, 3).0; + assert!(solid[2] > 250 && (solid[0] as i32 - 128).abs() <= 2); + assert_eq!(empty.unwrap().get_pixel(3, 3).0, [32, 32, 32]); + assert!(render_normal_preview(&image, &json!({})).is_none()); + } + #[test] fn refinement_keeps_normals_away_from_edges() { let image = Rgb32FImage::from_fn(128, 128, |x, y| { diff --git a/src/components/adjustments/Effects.tsx b/src/components/adjustments/Effects.tsx index 3ce2592294..6478e6d4a1 100644 --- a/src/components/adjustments/Effects.tsx +++ b/src/components/adjustments/Effects.tsx @@ -15,6 +15,7 @@ import { getAdjustmentToolOrder, getHiddenAdjustmentTools, RelightLight, + RelightQuality, DEFAULT_RELIGHT_LIGHT, MAX_RELIGHT_LIGHTS, } from '../../utils/adjustments'; @@ -182,15 +183,19 @@ export default function EffectsPanel({ } }; - const handleGenerateRelightNormalMap = async () => { + const relightQuality: RelightQuality = adjustments.relightQuality ?? 'standard'; + + const handleGenerateRelightNormalMap = async (quality: RelightQuality = relightQuality) => { setIsGeneratingNormal(true); try { const maps: { normalMap: string; depthMap: string; depthScale: number } = await invoke('generate_relight_maps', { jsAdjustments: adjustments, + quality, }); setAdjustments((prev: Partial) => ({ ...prev, relightNormalMap: maps.normalMap, + relightQuality: quality, relightDepthMap: maps.depthMap, relightDepthScale: maps.depthScale, })); @@ -233,6 +238,7 @@ export default function EffectsPanel({ const relightLights = adjustments.relightLights ?? []; const activeRelightLightId = useEditorStore((state) => state.activeRelightLightId); const setEditor = useEditorStore((state) => state.setEditor); + const showRelightNormalMap = useEditorStore((state) => state.showRelightNormalMap); const activeRelightLight = relightLights.find((light) => light.id === activeRelightLightId) ?? relightLights[0]; const relightColor = (activeRelightLight?.color || '#ffffff').toLowerCase(); @@ -623,6 +629,26 @@ export default function EffectsPanel({ fillOrigin="min" /> + handleAdjustmentChange(Effect.RelightDetail, e.target.value)} + step={1} + value={adjustments.relightDetail ?? 30} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + {adjustments.relightNormalMap && ( + setEditor({ showRelightNormalMap: checked })} + /> + )} + {adjustments.relightNormalMap && ( )} - + ))} +
+
+ + - )} - {relightLights.length > 1 && ( - - )} -
-
- - updateActiveRelightLight({ intensity: parseInt(e.target.value, 10) })} - step={1} - value={activeRelightLight?.intensity ?? 50} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - updateActiveRelightLight({ depth: parseInt(e.target.value, 10) })} - step={1} - value={activeRelightLight?.depth ?? 30} - onDragStateChange={onDragStateChange} - /> - - updateActiveRelightLight({ range: parseInt(e.target.value, 10) })} - step={1} - value={activeRelightLight?.range ?? 60} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - -
- - {t('adjustments.effects.relightColor')} - -
- {RELIGHT_TEMPERATURE_PRESETS.map((kelvin) => { - const presetColor = kelvinToHex(kelvin); - return ( - - ); - })} -
- -
- - handleAdjustmentChange(Effect.RelightAmbient, e.target.value)} - step={1} - value={adjustments.relightAmbient ?? 80} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - handleAdjustmentChange(Effect.RelightSoftness, e.target.value)} - step={1} - value={adjustments.relightSoftness ?? 30} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - handleAdjustmentChange(Effect.RelightSpecular, e.target.value)} - step={1} - value={adjustments.relightSpecular ?? 0} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - handleAdjustmentChange(Effect.RelightDetail, e.target.value)} - step={1} - value={adjustments.relightDetail ?? 30} - onDragStateChange={onDragStateChange} - fillOrigin="min" - /> - - {adjustments.relightNormalMap && ( - setEditor({ showRelightNormalMap: checked })} - /> - )} - - {adjustments.relightNormalMap && ( - {t('adjustments.effects.relightNormalMap')} - )} - -
- - {t('adjustments.effects.relightQuality')} - -
- {(['standard', 'high'] as const).map((quality) => ( - - ))} -
-
- - - - )} -
- - - - )} - {!hiddenTools.includes('lut') && ( = { [Panel.Masks]: Layers, [Panel.Ai]: Paintbrush, [Panel.Presets]: SwatchBook, + [Panel.Relight]: Lightbulb, [Panel.Export]: FileInput, [Panel.FolderTree]: FolderIcon, [Panel.Tethering]: Camera, @@ -37,6 +39,7 @@ const PANEL_TITLES: Record = { [Panel.Masks]: 'editor.switcher.tooltips.masks', [Panel.Ai]: 'editor.switcher.tooltips.inpaint', [Panel.Presets]: 'editor.switcher.tooltips.presets', + [Panel.Relight]: 'editor.switcher.tooltips.relight', [Panel.Export]: 'editor.switcher.tooltips.export', [Panel.FolderTree]: 'library.folders.sourcesTitle', [Panel.Tethering]: 'editor.switcher.tooltips.tethering', @@ -256,6 +259,7 @@ export function MobilePanelSwitcher({ Panel.Masks, Panel.Ai, Panel.Presets, + Panel.Relight, Panel.Export, ]; diff --git a/src/components/panel/editor/ImageCanvas.tsx b/src/components/panel/editor/ImageCanvas.tsx index 7c6b448619..d6a7be23f5 100644 --- a/src/components/panel/editor/ImageCanvas.tsx +++ b/src/components/panel/editor/ImageCanvas.tsx @@ -15,7 +15,7 @@ import { RelightLight, } from '../../../utils/adjustments'; import { Mask, SubMask, SubMaskMode, ToolType } from '../right/Masks'; -import { AppSettings, BrushSettings, SelectedImage } from '../../ui/AppProperties'; +import { AppSettings, BrushSettings, Panel, SelectedImage } from '../../ui/AppProperties'; import { RenderSize } from '../../../hooks/useImageRenderSize'; import { useOsPlatform } from '../../../hooks/useOsPlatform'; import { useTranslation } from 'react-i18next'; @@ -1675,7 +1675,7 @@ const ImageCanvas = memo( : activeCrop.y : 0; - const isEffectsSectionOpen = useUIStore((state) => state.collapsibleSectionsState.effects); + const isRelightPanelActive = useUIStore((state) => Object.values(state.activePanels).includes(Panel.Relight)); const activeRotation = liveRotation !== null && liveRotation !== undefined ? liveRotation : adjustments.rotation || 0; const relightGeometry = useMemo( @@ -1718,8 +1718,7 @@ const ImageCanvas = memo( const isRelightHandleVisible = !!adjustments.relightEnabled && !!adjustments.relightNormalMap && - adjustments.sectionVisibility?.effects !== false && - isEffectsSectionOpen && + isRelightPanelActive && !isCropping && !isMasking && !isAiEditing && @@ -1744,27 +1743,39 @@ const ImageCanvas = memo( adjustments.orientationSteps, ]); + const showRelightNormalPreview = useCallback((url: string | null) => { + setRelightNormalPreviewUrl((previous) => { + if (previous) URL.revokeObjectURL(previous); + return url; + }); + }, []); + + useEffect(() => () => showRelightNormalPreview(null), [showRelightNormalPreview]); + useEffect(() => { const request = ++relightNormalPreviewRequest.current; if (!isRelightNormalMapVisible) { - setRelightNormalPreviewUrl(null); + showRelightNormalPreview(null); return; } - // Wait for the detail slider to settle: each render covers the full-size image. + // A short pause so that dragging the detail slider does not queue a render per step. const timeout = window.setTimeout(async () => { try { - const url: string = await invoke('generate_relight_normal_preview', { - jsAdjustments: relightPreviewAdjustmentsRef.current, - }); - if (request === relightNormalPreviewRequest.current) setRelightNormalPreviewUrl(url); + // Masks, patches and LUT data play no part here and only make the request heavier. + const payload: Record = { ...relightPreviewAdjustmentsRef.current }; + for (const key of ['masks', 'aiPatches', 'lutData', 'lensBlurDepthMap']) delete payload[key]; + + const bytes: ArrayBuffer = await invoke('generate_relight_normal_preview', { jsAdjustments: payload }); + if (request !== relightNormalPreviewRequest.current) return; + showRelightNormalPreview(URL.createObjectURL(new Blob([new Uint8Array(bytes)], { type: 'image/jpeg' }))); } catch (e) { if (request === relightNormalPreviewRequest.current) { - setRelightNormalPreviewUrl(null); + showRelightNormalPreview(null); toast.error(`Failed to render normal map: ${e}`); } } - }, 150); + }, 40); return () => window.clearTimeout(timeout); }, [ isRelightNormalMapVisible, @@ -1773,6 +1784,7 @@ const ImageCanvas = memo( adjustments.relightDepthScale, adjustments.relightDetail, relightPreviewLayoutKey, + showRelightNormalPreview, ]); const effectiveZoomScale = transformState.scale > 0 ? transformState.scale : 1; diff --git a/src/components/panel/right/RelightPanel.tsx b/src/components/panel/right/RelightPanel.tsx new file mode 100644 index 0000000000..97536e94bf --- /dev/null +++ b/src/components/panel/right/RelightPanel.tsx @@ -0,0 +1,428 @@ +import { useCallback, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { invoke } from '@tauri-apps/api/core'; +import { toast } from 'react-toastify'; +import { Loader2, Plus, RotateCcw, Trash2 } from 'lucide-react'; +import clsx from 'clsx'; +import { v4 as uuidv4 } from 'uuid'; +import Slider from '../../ui/Slider'; +import Switch from '../../ui/Switch'; +import Button from '../../ui/Button'; +import Text from '../../ui/Text'; +import { TextColors, TextVariants, TextWeights } from '../../../types/typography'; +import { + Adjustments, + Effect, + INITIAL_ADJUSTMENTS, + RelightLight, + RelightQuality, + DEFAULT_RELIGHT_LIGHT, + MAX_RELIGHT_LIGHTS, +} from '../../../utils/adjustments'; +import { RELIGHT_TEMPERATURE_PRESETS, kelvinToHex } from '../../../utils/relightUtils'; +import { useEditorStore } from '../../../store/useEditorStore'; +import { useProcessStore } from '../../../store/useProcessStore'; +import { useEditorActions } from '../../../hooks/useEditorActions'; + +const NEW_RELIGHT_LIGHT_TEMPERATURES = [5600, 3200, 8000, 2700, 10000, 4000]; + +const RELIGHT_ADJUSTMENT_KEYS: Array = [ + Effect.RelightEnabled, + Effect.RelightLights, + Effect.RelightAmbient, + Effect.RelightSoftness, + Effect.RelightSpecular, + Effect.RelightDetail, + Effect.RelightNormalMap, + Effect.RelightQuality, + Effect.RelightDepthMap, + Effect.RelightDepthScale, +]; + +export default function RelightPanel() { + const { t } = useTranslation(); + const { setAdjustments } = useEditorActions(); + const adjustments = useEditorStore((state) => state.adjustments); + const selectedImage = useEditorStore((state) => state.selectedImage); + const activeRelightLightId = useEditorStore((state) => state.activeRelightLightId); + const showRelightNormalMap = useEditorStore((state) => state.showRelightNormalMap); + const setEditor = useEditorStore((state) => state.setEditor); + const aiModelDownloadStatus = useProcessStore((state) => state.aiModelDownloadStatus); + const [isGeneratingNormal, setIsGeneratingNormal] = useState(false); + + const onDragStateChange = useCallback( + (isDragging: boolean) => setEditor({ isSliderDragging: isDragging }), + [setEditor], + ); + + const relightQuality: RelightQuality = adjustments.relightQuality ?? 'standard'; + + const handleGenerateRelightNormalMap = async (quality: RelightQuality = relightQuality) => { + setIsGeneratingNormal(true); + try { + const maps: { normalMap: string; depthMap: string; depthScale: number } = await invoke('generate_relight_maps', { + jsAdjustments: adjustments, + quality, + }); + setAdjustments((prev: Adjustments) => ({ + ...prev, + relightNormalMap: maps.normalMap, + relightQuality: quality, + relightDepthMap: maps.depthMap, + relightDepthScale: maps.depthScale, + })); + } catch (e) { + toast.error(`Failed to generate normal map: ${e}`); + setAdjustments((prev: Adjustments) => ({ ...prev, relightEnabled: false })); + } finally { + setIsGeneratingNormal(false); + } + }; + + const handleSliderChange = (key: keyof Adjustments, value: string | number) => { + setAdjustments((prev: Adjustments) => ({ ...prev, [key]: Number(value) })); + }; + + const relightLights = adjustments.relightLights ?? []; + const activeRelightLight = relightLights.find((light) => light.id === activeRelightLightId) ?? relightLights[0]; + const relightColor = (activeRelightLight?.color || '#ffffff').toLowerCase(); + + const updateActiveRelightLight = (changes: Partial) => { + if (!activeRelightLight) return; + const id = activeRelightLight.id; + setAdjustments((prev: Adjustments) => ({ + ...prev, + relightLights: (prev.relightLights ?? []).map((light) => (light.id === id ? { ...light, ...changes } : light)), + })); + }; + + const handleRelightColorChange = (color: string) => { + updateActiveRelightLight({ color: color.toLowerCase() }); + }; + + const handleAddRelightLight = () => { + if (relightLights.length >= MAX_RELIGHT_LIGHTS) return; + // Each new light starts somewhere else and with another tint, so it is easy to tell apart. + const index = relightLights.length; + const light: RelightLight = { + ...DEFAULT_RELIGHT_LIGHT, + id: uuidv4(), + x: [0.7, 0.5, 0.3, 0.7, 0.5][index % 5], + y: [0.3, 0.7, 0.7, 0.7, 0.2][index % 5], + color: kelvinToHex(NEW_RELIGHT_LIGHT_TEMPERATURES[index % NEW_RELIGHT_LIGHT_TEMPERATURES.length]), + }; + setAdjustments((prev: Adjustments) => ({ + ...prev, + relightLights: [...(prev.relightLights ?? []), light], + })); + setEditor({ activeRelightLightId: light.id }); + }; + + const handleRemoveRelightLight = () => { + if (!activeRelightLight || relightLights.length <= 1) return; + const id = activeRelightLight.id; + setAdjustments((prev: Adjustments) => ({ + ...prev, + relightLights: (prev.relightLights ?? []).filter((light) => light.id !== id), + })); + setEditor({ activeRelightLightId: null }); + }; + + const handleRelightToggle = (enabled: boolean) => { + setAdjustments((prev: Adjustments) => ({ ...prev, relightEnabled: enabled })); + if (enabled && (!adjustments.relightNormalMap || !adjustments.relightDepthMap)) { + handleGenerateRelightNormalMap(); + } + }; + + const handleReset = () => { + setAdjustments((prev: Adjustments) => ({ + ...prev, + ...Object.fromEntries(RELIGHT_ADJUSTMENT_KEYS.map((key) => [key, INITIAL_ADJUSTMENTS[key]])), + })); + setEditor({ activeRelightLightId: null, showRelightNormalMap: false }); + }; + + return ( +
+
+ {t('adjustments.effects.relight')} + +
+ +
+ {selectedImage ? ( + <> + + +
+
+
+ {isGeneratingNormal ? ( +
+
+ + + {aiModelDownloadStatus + ? t('editor.masks.settings.aiModelDownloading') + : t('editor.ai.generatingNormalMap')} + +
+ {aiModelDownloadStatus && ( + + {aiModelDownloadStatus} + + )} +
+ ) : ( + <> +
+ + {t('adjustments.effects.relightLights')} + +
+ {relightLights.map((light) => ( + + )} + {relightLights.length > 1 && ( + + )} +
+
+ + updateActiveRelightLight({ intensity: Number(e.target.value) })} + step={1} + value={activeRelightLight?.intensity ?? 50} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + updateActiveRelightLight({ depth: Number(e.target.value) })} + step={1} + value={activeRelightLight?.depth ?? 30} + onDragStateChange={onDragStateChange} + /> + + updateActiveRelightLight({ range: Number(e.target.value) })} + step={1} + value={activeRelightLight?.range ?? 60} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + +
+ + {t('adjustments.effects.relightColor')} + +
+ {RELIGHT_TEMPERATURE_PRESETS.map((kelvin) => { + const presetColor = kelvinToHex(kelvin); + return ( + + ); + })} +
+ +
+ + handleSliderChange(Effect.RelightAmbient, e.target.value)} + step={1} + value={adjustments.relightAmbient ?? 80} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleSliderChange(Effect.RelightSoftness, e.target.value)} + step={1} + value={adjustments.relightSoftness ?? 30} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleSliderChange(Effect.RelightSpecular, e.target.value)} + step={1} + value={adjustments.relightSpecular ?? 0} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + handleSliderChange(Effect.RelightDetail, e.target.value)} + step={1} + value={adjustments.relightDetail ?? 30} + onDragStateChange={onDragStateChange} + fillOrigin="min" + /> + + {adjustments.relightNormalMap && ( + setEditor({ showRelightNormalMap: checked })} + /> + )} + + {adjustments.relightNormalMap && ( + {t('adjustments.effects.relightNormalMap')} + )} + +
+ + {t('adjustments.effects.relightQuality')} + +
+ {(['standard', 'high'] as const).map((quality) => ( + + ))} +
+
+ + + + )} +
+
+
+ + ) : ( +
+ + {t('editor.ai.noImageSelected')} + +
+ )} +
+
+ ); +} diff --git a/src/components/ui/AppProperties.tsx b/src/components/ui/AppProperties.tsx index 9ae2229d4d..93fe380728 100644 --- a/src/components/ui/AppProperties.tsx +++ b/src/components/ui/AppProperties.tsx @@ -133,6 +133,7 @@ export enum Panel { Masks = 'masks', Metadata = 'metadata', Presets = 'presets', + Relight = 'relight', FolderTree = 'folderTree', Tethering = 'tethering', } diff --git a/src/i18n/locales/ca.json b/src/i18n/locales/ca.json index 007dc9613c..b88172df07 100644 --- a/src/i18n/locales/ca.json +++ b/src/i18n/locales/ca.json @@ -787,6 +787,7 @@ "inpaint": "Inpainting", "masks": "Màscares", "presets": "Ajustos preestablerts", + "relight": "Reil·luminar", "tethering": "Captura connectada" } }, diff --git a/src/i18n/locales/de.json b/src/i18n/locales/de.json index d1c8c56e50..4f20b7f731 100644 --- a/src/i18n/locales/de.json +++ b/src/i18n/locales/de.json @@ -787,6 +787,7 @@ "inpaint": "Reparieren", "masks": "Masken", "presets": "Vorgaben", + "relight": "Neu beleuchten", "tethering": "Tethering" } }, diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index 20a5a8b1dd..64c31b6d95 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -787,6 +787,7 @@ "inpaint": "Inpaint", "masks": "Masks", "presets": "Presets", + "relight": "Relight", "tethering": "Tethering" } }, diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index ebe58faf2a..a7fac296ca 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -808,6 +808,7 @@ "inpaint": "Inpainting", "masks": "Máscaras", "presets": "Ajustes preestablecidos", + "relight": "Reiluminar", "tethering": "Captura conectada" } }, diff --git a/src/i18n/locales/fr.json b/src/i18n/locales/fr.json index fb46791bce..df89e66604 100644 --- a/src/i18n/locales/fr.json +++ b/src/i18n/locales/fr.json @@ -808,6 +808,7 @@ "inpaint": "Inpainting", "masks": "Masques", "presets": "Paramètres prédéfinis", + "relight": "Rééclairage", "tethering": "Prise de vue connectée" } }, diff --git a/src/i18n/locales/it.json b/src/i18n/locales/it.json index 6bd7b2519f..bf8a467ab1 100644 --- a/src/i18n/locales/it.json +++ b/src/i18n/locales/it.json @@ -808,6 +808,7 @@ "inpaint": "Inpainting", "masks": "Maschere", "presets": "Predefiniti", + "relight": "Reilluminazione", "tethering": "Acquisizione diretta" } }, diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index 387340b22a..6fd045c9b4 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -785,6 +785,7 @@ "inpaint": "インペイント", "masks": "マスク", "presets": "プリセット", + "relight": "リライト", "tethering": "テザー撮影" } }, diff --git a/src/i18n/locales/ko.json b/src/i18n/locales/ko.json index 83173bb79a..fac98aea8f 100644 --- a/src/i18n/locales/ko.json +++ b/src/i18n/locales/ko.json @@ -785,6 +785,7 @@ "inpaint": "인페인팅", "masks": "마스크", "presets": "프리셋", + "relight": "재조명", "tethering": "테더링" } }, diff --git a/src/i18n/locales/pl.json b/src/i18n/locales/pl.json index 7a8fc18e8b..decef92474 100644 --- a/src/i18n/locales/pl.json +++ b/src/i18n/locales/pl.json @@ -829,6 +829,7 @@ "inpaint": "Inpaint", "masks": "Maski", "presets": "Presety", + "relight": "Ponowne oświetlenie", "tethering": "Tethering (Połączenie przewodowe)" } }, diff --git a/src/i18n/locales/pt.json b/src/i18n/locales/pt.json index cd896b3833..964217b3c3 100644 --- a/src/i18n/locales/pt.json +++ b/src/i18n/locales/pt.json @@ -808,6 +808,7 @@ "inpaint": "Preencher", "masks": "Máscaras", "presets": "Predefinições", + "relight": "Reiluminar", "tethering": "Captura Conectada (Tethering)" } }, diff --git a/src/i18n/locales/ru.json b/src/i18n/locales/ru.json index 25e65028c0..8331d9d979 100644 --- a/src/i18n/locales/ru.json +++ b/src/i18n/locales/ru.json @@ -829,6 +829,7 @@ "inpaint": "Восстановление", "masks": "Маски", "presets": "Пресеты", + "relight": "Переосвещение", "tethering": "Съемка на ПК (Тетеринг)" } }, diff --git a/src/i18n/locales/zh-CN.json b/src/i18n/locales/zh-CN.json index 590ff072de..853dc70140 100644 --- a/src/i18n/locales/zh-CN.json +++ b/src/i18n/locales/zh-CN.json @@ -785,6 +785,7 @@ "inpaint": "修复", "masks": "蒙版", "presets": "预设", + "relight": "重新打光", "tethering": "联机拍摄" } }, diff --git a/src/i18n/locales/zh-TW.json b/src/i18n/locales/zh-TW.json index f793c1a708..08026a0ca6 100644 --- a/src/i18n/locales/zh-TW.json +++ b/src/i18n/locales/zh-TW.json @@ -785,6 +785,7 @@ "inpaint": "修復", "masks": "遮色片", "presets": "預設", + "relight": "重新打光", "tethering": "連線拍攝" } }, diff --git a/src/store/useUIStore.ts b/src/store/useUIStore.ts index 76aecad690..541a8a40ac 100644 --- a/src/store/useUIStore.ts +++ b/src/store/useUIStore.ts @@ -93,6 +93,7 @@ const ALL_PANELS: Panel[] = [ Panel.Masks, Panel.Ai, Panel.Presets, + Panel.Relight, ]; const DEFAULT_PANEL_DEFAULT_REGIONS: Record = { @@ -105,6 +106,7 @@ const DEFAULT_PANEL_DEFAULT_REGIONS: Record = { [Panel.Masks]: 'rightTop', [Panel.Ai]: 'rightTop', [Panel.Presets]: 'rightTop', + [Panel.Relight]: 'rightTop', }; export const DEFAULT_PANEL_WIDTH = 350; @@ -125,7 +127,7 @@ export function reconcileWorkspace( panelLayout: { leftTop: [Panel.Metadata, Panel.FolderTree, Panel.Export, ...(isTetheringSupported ? [Panel.Tethering] : [])], leftBottom: [], - rightTop: [Panel.Adjustments, Panel.Crop, Panel.Masks, Panel.Ai, Panel.Presets], + rightTop: [Panel.Adjustments, Panel.Crop, Panel.Masks, Panel.Ai, Panel.Presets, Panel.Relight], rightBottom: [], }, activePanels: { @@ -290,7 +292,7 @@ export const useUIStore = create((set, get) => ({ panelLayout: { leftTop: [Panel.Metadata, Panel.FolderTree, Panel.Export], leftBottom: [], - rightTop: [Panel.Adjustments, Panel.Crop, Panel.Masks, Panel.Ai, Panel.Presets], + rightTop: [Panel.Adjustments, Panel.Crop, Panel.Masks, Panel.Ai, Panel.Presets, Panel.Relight], rightBottom: [], }, activePanels: { diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index e6421031e2..b7047eb283 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -1041,16 +1041,6 @@ export const ADJUSTMENT_SECTIONS: Sections = { Effect.LensBlurMaxDepth, Effect.LensBlurMinFade, Effect.LensBlurMaxFade, - Effect.RelightEnabled, - Effect.RelightLights, - Effect.RelightAmbient, - Effect.RelightSoftness, - Effect.RelightSpecular, - Effect.RelightNormalMap, - Effect.RelightQuality, - Effect.RelightDepthMap, - Effect.RelightDepthScale, - Effect.RelightDetail, ], }; @@ -1098,7 +1088,6 @@ export const ADJUSTMENT_SECTION_TOOLS: Record Date: Mon, 5 Oct 2026 17:38:12 -0500 Subject: [PATCH 09/16] feat: add RelightPanel component and Tauri backend commands for AI relighting --- src-tauri/src/ai_commands.rs | 193 +++++++++++++++++++- src-tauri/src/relight.rs | 89 +++++++-- src/components/panel/right/RelightPanel.tsx | 7 +- 3 files changed, 267 insertions(+), 22 deletions(-) diff --git a/src-tauri/src/ai_commands.rs b/src-tauri/src/ai_commands.rs index bb6347abff..d4f2e3923d 100644 --- a/src-tauri/src/ai_commands.rs +++ b/src-tauri/src/ai_commands.rs @@ -301,7 +301,7 @@ pub async fn generate_full_image_depth_map( Ok(format!("data:image/png;base64,{}", base64_str)) } -#[derive(serde::Serialize)] +#[derive(serde::Serialize, serde::Deserialize, Debug, Clone, PartialEq)] #[serde(rename_all = "camelCase")] pub struct RelightMapsPayload { normal_map: String, @@ -318,18 +318,113 @@ fn encode_png_data_url(image: &image::DynamicImage) -> Result { Ok(format!("data:image/png;base64,{}", base64_str)) } +// Generated maps are kept on disk, one file per image, geometry and model quality, so that +// switching quality back and forth or reopening a photo does not run the model again. +const RELIGHT_MAPS_CACHE_VERSION: &str = "v1"; +const RELIGHT_MAPS_CACHE_MAX_FILES: usize = 60; + +fn relight_maps_cache_file( + cache_dir: &std::path::Path, + image_path: &str, + js_adjustments: &serde_json::Value, + quality: crate::ai_processing::NormalModelQuality, +) -> std::path::PathBuf { + let mut hasher = blake3::Hasher::new(); + hasher.update(RELIGHT_MAPS_CACHE_VERSION.as_bytes()); + hasher.update(image_path.as_bytes()); + hasher.update(format!("{:?}", quality).as_bytes()); + hasher.update(&crate::cache_utils::calculate_geometry_hash(js_adjustments).to_le_bytes()); + + // A file that was replaced on disk must not reuse the maps of the old one. + let (source_path, _) = crate::file_management::parse_virtual_path(image_path); + if let Ok(metadata) = std::fs::metadata(&source_path) { + hasher.update(&metadata.len().to_le_bytes()); + if let Ok(elapsed) = metadata.modified().and_then(|m| { + m.duration_since(std::time::UNIX_EPOCH) + .map_err(std::io::Error::other) + }) { + hasher.update(&elapsed.as_nanos().to_le_bytes()); + } + } + + cache_dir.join(format!("{}.json", hasher.finalize().to_hex())) +} + +fn read_relight_maps_cache(file: &std::path::Path) -> Option { + let bytes = std::fs::read(file).ok()?; + serde_json::from_slice(&bytes).ok() +} + +fn write_relight_maps_cache( + file: &std::path::Path, + maps: &RelightMapsPayload, +) -> Result<(), String> { + let dir = file.parent().ok_or("Invalid relight cache path")?; + std::fs::create_dir_all(dir).map_err(|e| e.to_string())?; + + // Written next to the target and renamed, so a crash never leaves half a file behind. + let tmp = file.with_extension("tmp"); + std::fs::write(&tmp, serde_json::to_vec(maps).map_err(|e| e.to_string())?) + .map_err(|e| e.to_string())?; + std::fs::rename(&tmp, file).map_err(|e| e.to_string())?; + + // Drop the least recently written maps once there are too many. + let mut entries: Vec<(std::time::SystemTime, std::path::PathBuf)> = std::fs::read_dir(dir) + .map_err(|e| e.to_string())? + .filter_map(|entry| entry.ok()) + .filter(|entry| entry.path().extension().is_some_and(|ext| ext == "json")) + .filter_map(|entry| Some((entry.metadata().ok()?.modified().ok()?, entry.path()))) + .collect(); + if entries.len() > RELIGHT_MAPS_CACHE_MAX_FILES { + entries.sort(); + for (_, path) in &entries[..entries.len() - RELIGHT_MAPS_CACHE_MAX_FILES] { + let _ = std::fs::remove_file(path); + } + } + Ok(()) +} + +/// Returns the relight maps for the open image, generating them only when there is no saved +/// copy for this image and quality (or when `force` asks for a fresh run). #[tauri::command] pub async fn generate_relight_maps( js_adjustments: serde_json::Value, quality: Option, + force: Option, state: tauri::State<'_, AppState>, app_handle: tauri::AppHandle, ) -> Result { + use tauri::Manager; + + let quality = quality.unwrap_or_default(); + let image_path = state + .original_image + .lock() + .unwrap() + .as_ref() + .map(|loaded| loaded.path.clone()) + .ok_or("No original image loaded")?; + + let cache_file = app_handle.path().app_cache_dir().ok().map(|dir| { + relight_maps_cache_file( + &dir.join("relight_maps"), + &image_path, + &js_adjustments, + quality, + ) + }); + + if !force.unwrap_or(false) + && let Some(maps) = cache_file.as_deref().and_then(read_relight_maps_cache) + { + return Ok(maps); + } + let normal_model = crate::ai_processing::get_or_init_normal_model( &app_handle, &state.ai_state, &state.ai_init_lock, - quality.unwrap_or_default(), + quality, ) .await .map_err(|e| e.to_string())?; @@ -339,11 +434,23 @@ pub async fn generate_relight_maps( let maps = crate::ai_processing::run_normal_model(warped_image.as_ref(), &normal_model) .map_err(|e| e.to_string())?; - Ok(RelightMapsPayload { + let payload = RelightMapsPayload { normal_map: encode_png_data_url(&image::DynamicImage::ImageRgba8(maps.normal))?, depth_map: encode_png_data_url(&image::DynamicImage::ImageLuma16(maps.depth))?, depth_scale: maps.depth_scale, - }) + }; + + if let Some(file) = &cache_file + && let Err(error) = write_relight_maps_cache(file, &payload) + { + log::warn!( + "Could not save relight maps to {}: {}", + file.display(), + error + ); + } + + Ok(payload) } /// Normal map as the relight pass uses it, laid out like the edited image (orientation, @@ -651,3 +758,81 @@ pub async fn test_ai_connector_connection(address: String) -> Result<(), String> Err(e) => Err(e.to_string()), } } + +#[cfg(test)] +mod relight_maps_cache_tests { + use super::*; + use crate::ai_processing::NormalModelQuality; + use serde_json::json; + + fn payload(tag: &str) -> RelightMapsPayload { + RelightMapsPayload { + normal_map: format!("normal-{tag}"), + depth_map: format!("depth-{tag}"), + depth_scale: 1.5, + } + } + + #[test] + fn saved_maps_are_read_back() { + let dir = tempfile::tempdir().unwrap(); + let file = relight_maps_cache_file( + dir.path(), + "photo.raw", + &json!({}), + NormalModelQuality::High, + ); + + assert_eq!(read_relight_maps_cache(&file), None); + write_relight_maps_cache(&file, &payload("a")).unwrap(); + assert_eq!(read_relight_maps_cache(&file), Some(payload("a"))); + } + + #[test] + fn each_image_quality_and_geometry_gets_its_own_file() { + let dir = std::path::Path::new("cache"); + let base = + relight_maps_cache_file(dir, "photo.raw", &json!({}), NormalModelQuality::Standard); + + let same = relight_maps_cache_file( + dir, + "photo.raw", + &json!({ "exposure": 2 }), + NormalModelQuality::Standard, + ); + let other_quality = + relight_maps_cache_file(dir, "photo.raw", &json!({}), NormalModelQuality::High); + let other_image = + relight_maps_cache_file(dir, "other.raw", &json!({}), NormalModelQuality::Standard); + let other_geometry = relight_maps_cache_file( + dir, + "photo.raw", + &json!({ "transformRotate": 5 }), + NormalModelQuality::Standard, + ); + + assert_eq!(base, same); + assert_ne!(base, other_quality); + assert_ne!(base, other_image); + assert_ne!(base, other_geometry); + } + + #[test] + fn oldest_maps_are_dropped_when_the_cache_is_full() { + let dir = tempfile::tempdir().unwrap(); + for i in 0..RELIGHT_MAPS_CACHE_MAX_FILES + 3 { + let file = dir.path().join(format!("{i:04}.json")); + write_relight_maps_cache(&file, &payload("x")).unwrap(); + let stamp = filetime::FileTime::from_unix_time(1_700_000_000 + i as i64, 0); + filetime::set_file_mtime(&file, stamp).unwrap(); + } + // One more write triggers the clean-up with every timestamp in place. + let newest = dir.path().join("newest.json"); + write_relight_maps_cache(&newest, &payload("x")).unwrap(); + + let remaining = std::fs::read_dir(dir.path()).unwrap().count(); + assert_eq!(remaining, RELIGHT_MAPS_CACHE_MAX_FILES); + assert!(!dir.path().join("0000.json").exists()); + assert!(newest.exists()); + } +} diff --git a/src-tauri/src/relight.rs b/src-tauri/src/relight.rs index 17cdf372cd..fc5b78ca0f 100644 --- a/src-tauri/src/relight.rs +++ b/src-tauri/src/relight.rs @@ -357,6 +357,30 @@ fn local_deviation(plane: &[f32], w: usize, h: usize, radius: usize) -> Vec .collect() } +/// How well the photo explains a plane around each pixel: the squared correlation between +/// the two, from 0 (unrelated) to 1 (the photo's edges are the plane's edges). +fn guide_agreement(guide: &[f32], plane: &[f32], w: usize, h: usize, radius: usize) -> Vec { + let blurred = |values: Vec| { + let mut values = values; + dof_box_filter(&mut values, w, h, 1, radius); + values + }; + let mean_i = blurred(guide.to_vec()); + let mean_p = blurred(plane.to_vec()); + let mean_ii = blurred(guide.iter().map(|v| v * v).collect()); + let mean_pp = blurred(plane.iter().map(|v| v * v).collect()); + let mean_ip = blurred(guide.iter().zip(plane).map(|(i, p)| i * p).collect()); + + (0..w * h) + .map(|k| { + let cov = mean_ip[k] - mean_i[k] * mean_p[k]; + let var_i = (mean_ii[k] - mean_i[k] * mean_i[k]).max(0.0); + let var_p = (mean_pp[k] - mean_p[k] * mean_p[k]).max(0.0); + (cov * cov / (var_i * var_p + 1e-9)).min(1.0) + }) + .collect() +} + /// Scales the AI maps up to the image. With `refine`, silhouettes (hair, object outlines) /// are snapped to the edges of the photo instead of being left as a soft blur, while the /// inside of each surface keeps the detail the model predicted. @@ -404,25 +428,26 @@ fn build_surface_maps( let radius = ((gw.max(gh) as f32) * 0.015).round().max(2.0) as usize; - // Only trust the photo's edges where the geometry really breaks: a depth jump or the - // border of the valid area. Elsewhere the filter would just blur the normals. - let coverage_dev = local_deviation(&planes[3], gw, gh, radius); + // Only trust the photo's edges where the geometry really breaks (a depth jump or the + // border of the valid area) and the photo actually shows that break. Anywhere else + // the filter has nothing to follow and would just smear the maps. let has_depth = depth_map.is_some() && depth_scale > 0.0; - let other_dev = if has_depth { - local_deviation(&planes[4], gw, gh, radius) - } else { - local_deviation(&planes[2], gw, gh, radius) - }; - let mut edge_weight: Vec = coverage_dev - .iter() - .zip(&other_dev) - .map(|(cov, other)| { - let geometry = if has_depth { - smoothstep(0.01, 0.06, other * depth_scale) + let geometry_plane = if has_depth { &planes[4] } else { &planes[2] }; + let coverage_dev = local_deviation(&planes[3], gw, gh, radius); + let geometry_dev = local_deviation(geometry_plane, gw, gh, radius); + let coverage_agreement = guide_agreement(&guide, &planes[3], gw, gh, radius); + let geometry_agreement = guide_agreement(&guide, geometry_plane, gw, gh, radius); + + let mut edge_weight: Vec = (0..gw * gh) + .map(|k| { + let geometry_edge = if has_depth { + smoothstep(0.01, 0.06, geometry_dev[k] * depth_scale) } else { - smoothstep(0.05, 0.2, *other) + smoothstep(0.05, 0.2, geometry_dev[k]) }; - smoothstep(0.05, 0.25, *cov).max(geometry) + let coverage_edge = smoothstep(0.05, 0.25, coverage_dev[k]); + (coverage_edge * smoothstep(0.25, 0.6, coverage_agreement[k])) + .max(geometry_edge * smoothstep(0.25, 0.6, geometry_agreement[k])) }) .collect(); dof_box_filter(&mut edge_weight, gw, gh, 1, (radius / 2).max(1)); @@ -924,6 +949,38 @@ mod tests { assert_eq!(preview, again); } + #[test] + fn refinement_leaves_edges_alone_when_the_photo_does_not_show_them() { + // The geometry has a hard edge (depth and normal change at the middle), but the photo + // is evenly lit noise there, so it gives no hint about where that edge really is. + let image = Rgb32FImage::from_fn(512, 128, |x, y| { + let v = 0.4 + 0.1 * (((x * 7 + y * 13) % 11) as f32 / 11.0); + Rgb([v, v, v]) + }); + let normal_map = RgbaImage::from_fn(256, 64, |x, _| { + if x < 128 { + Rgba([60, 128, 230, 255]) + } else { + Rgba([200, 128, 230, 255]) + } + }); + let depth_map = + DepthMap::from_fn(256, 64, |x, _| Luma([if x < 128 { 0 } else { u16::MAX }])); + + let plain = build_surface_maps(&image, &normal_map, Some(&depth_map), 2.0, false); + let refined = build_surface_maps(&image, &normal_map, Some(&depth_map), 2.0, true); + + // Guessing would smear the edge over a wide band; it has to stay as sharp as it was. + let row = 64 * 512; + for x in 236..276 { + let (a, b) = (plain.normals[row + x], refined.normals[row + x]); + assert!( + (a[0] as i32 - b[0] as i32).abs() <= 6, + "normal at x={x} changed from {a:?} to {b:?}" + ); + } + } + #[test] fn refinement_keeps_normals_away_from_edges() { let image = Rgb32FImage::from_fn(128, 128, |x, y| { diff --git a/src/components/panel/right/RelightPanel.tsx b/src/components/panel/right/RelightPanel.tsx index 97536e94bf..a65b909a6f 100644 --- a/src/components/panel/right/RelightPanel.tsx +++ b/src/components/panel/right/RelightPanel.tsx @@ -57,12 +57,15 @@ export default function RelightPanel() { const relightQuality: RelightQuality = adjustments.relightQuality ?? 'standard'; - const handleGenerateRelightNormalMap = async (quality: RelightQuality = relightQuality) => { + // Maps are saved per image and quality, so this only runs the model when there is no + // saved copy yet, or when `force` asks for a fresh one. + const handleGenerateRelightNormalMap = async (quality: RelightQuality = relightQuality, force = false) => { setIsGeneratingNormal(true); try { const maps: { normalMap: string; depthMap: string; depthScale: number } = await invoke('generate_relight_maps', { jsAdjustments: adjustments, quality, + force, }); setAdjustments((prev: Adjustments) => ({ ...prev, @@ -399,7 +402,7 @@ export default function RelightPanel() { - + + ) : ( + + )} + {denoiseError && ( + + {denoiseError} + + )} + + )} )} diff --git a/src/components/ui/AppProperties.tsx b/src/components/ui/AppProperties.tsx index 2a3005c8b0..e85057c515 100644 --- a/src/components/ui/AppProperties.tsx +++ b/src/components/ui/AppProperties.tsx @@ -36,6 +36,9 @@ export enum Invokes { ApplyAdjustmentsToPaths = 'apply_adjustments_to_paths', ApplyAutoAdjustmentsToPaths = 'apply_auto_adjustments_to_paths', ApplyDenoising = 'apply_denoising', + DeleteAiDenoiseLayer = 'delete_ai_denoise_layer', + GenerateAiDenoiseLayer = 'generate_ai_denoise_layer', + GetAiDenoiseLayerStatus = 'get_ai_denoise_layer_status', CalculateAutoAdjustments = 'calculate_auto_adjustments', SampleWhiteBalance = 'sample_white_balance', CancelAiTask = 'cancel_ai_task', diff --git a/src/i18n/locales/ca.json b/src/i18n/locales/ca.json index 6844cb2b62..fe2abd3d85 100644 --- a/src/i18n/locales/ca.json +++ b/src/i18n/locales/ca.json @@ -96,6 +96,11 @@ "resetPoint": "Restablir corba de punts de {{channel}}" }, "details": { + "aiDenoise": "Reducció de soroll amb IA", + "aiDenoiseGenerate": "Generar reducció de soroll amb IA", + "aiDenoiseGenerating": "Generant...", + "aiDenoiseRawOnly": "Només disponible per a fitxers RAW", + "aiDenoiseRemove": "Eliminar la reducció de soroll amb IA", "blueYellow": "Blau/Groc", "centre": "Centre", "chromaticAberration": "Aberració cromàtica", diff --git a/src/i18n/locales/cs.json b/src/i18n/locales/cs.json index 15e17da54c..91324fd05b 100644 --- a/src/i18n/locales/cs.json +++ b/src/i18n/locales/cs.json @@ -95,6 +95,11 @@ "resetPoint": "Resetovat bodovou křivku ({{channel}})" }, "details": { + "aiDenoise": "AI odšumění", + "aiDenoiseGenerate": "Vygenerovat AI odšumění", + "aiDenoiseGenerating": "Generování...", + "aiDenoiseRawOnly": "Dostupné pouze pro soubory RAW", + "aiDenoiseRemove": "Odstranit AI odšumění", "blueYellow": "Modrá/Žlutá", "centre": "Střed", "chromaticAberration": "Chromatická aberace", diff --git a/src/i18n/locales/de.json b/src/i18n/locales/de.json index 9050f1fca3..7722ba690c 100644 --- a/src/i18n/locales/de.json +++ b/src/i18n/locales/de.json @@ -96,6 +96,11 @@ "resetPoint": "Punktkurve ({{channel}}) zurücksetzen" }, "details": { + "aiDenoise": "KI-Entrauschen", + "aiDenoiseGenerate": "KI-Entrauschen erzeugen", + "aiDenoiseGenerating": "Wird erzeugt...", + "aiDenoiseRawOnly": "Nur für RAW-Dateien verfügbar", + "aiDenoiseRemove": "KI-Entrauschen entfernen", "blueYellow": "Blau/Gelb", "centre": "Zentriert", "chromaticAberration": "Chromatische Aberration", diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index 8300f01221..4b94ae5a80 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -96,6 +96,11 @@ "resetPoint": "Reset {{channel}} Point Curve" }, "details": { + "aiDenoise": "AI Denoise", + "aiDenoiseGenerate": "Generate AI Denoise", + "aiDenoiseGenerating": "Generating...", + "aiDenoiseRawOnly": "Only available for RAW files", + "aiDenoiseRemove": "Remove AI denoise", "blueYellow": "Blue/Yellow", "centre": "Centré", "chromaticAberration": "Chromatic Aberration", diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index 3276c471da..986d3b9224 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -96,6 +96,11 @@ "resetPoint": "Restablecer curva de puntos de {{channel}}" }, "details": { + "aiDenoise": "Reducción de ruido con IA", + "aiDenoiseGenerate": "Generar reducción de ruido con IA", + "aiDenoiseGenerating": "Generando...", + "aiDenoiseRawOnly": "Solo disponible para archivos RAW", + "aiDenoiseRemove": "Eliminar la reducción de ruido con IA", "blueYellow": "Azul/Amarillo", "centre": "Centro", "chromaticAberration": "Aberración cromática", diff --git a/src/i18n/locales/fr.json b/src/i18n/locales/fr.json index 4bb2f8d0d6..5a7280b162 100644 --- a/src/i18n/locales/fr.json +++ b/src/i18n/locales/fr.json @@ -96,6 +96,11 @@ "resetPoint": "Réinitialiser la courbe à points {{channel}}" }, "details": { + "aiDenoise": "Débruitage IA", + "aiDenoiseGenerate": "Générer le débruitage IA", + "aiDenoiseGenerating": "Génération...", + "aiDenoiseRawOnly": "Disponible uniquement pour les fichiers RAW", + "aiDenoiseRemove": "Supprimer le débruitage IA", "blueYellow": "Bleu/Jaune", "centre": "Centre", "chromaticAberration": "Aberration chromatique", diff --git a/src/i18n/locales/it.json b/src/i18n/locales/it.json index 927201c899..7b9507eda1 100644 --- a/src/i18n/locales/it.json +++ b/src/i18n/locales/it.json @@ -96,6 +96,11 @@ "resetPoint": "Ripristina Curva per Punti {{channel}}" }, "details": { + "aiDenoise": "Riduzione rumore IA", + "aiDenoiseGenerate": "Genera riduzione rumore IA", + "aiDenoiseGenerating": "Generazione...", + "aiDenoiseRawOnly": "Disponibile solo per i file RAW", + "aiDenoiseRemove": "Rimuovi riduzione rumore IA", "blueYellow": "Blu/Giallo", "centre": "Centro", "chromaticAberration": "Aberrazione Cromatica", diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index d1eea3ab44..3f4a473ed1 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -96,6 +96,11 @@ "resetPoint": "{{channel}} ポイントカーブをリセット" }, "details": { + "aiDenoise": "AIノイズ低減", + "aiDenoiseGenerate": "AIノイズ低減を生成", + "aiDenoiseGenerating": "生成中...", + "aiDenoiseRawOnly": "RAWファイルでのみ利用できます", + "aiDenoiseRemove": "AIノイズ低減を削除", "blueYellow": "ブルー/イエロー", "centre": "中央", "chromaticAberration": "色収差", diff --git a/src/i18n/locales/ko.json b/src/i18n/locales/ko.json index 6e022e75f9..431b440076 100644 --- a/src/i18n/locales/ko.json +++ b/src/i18n/locales/ko.json @@ -96,6 +96,11 @@ "resetPoint": "{{channel}} 포인트 커브 초기화" }, "details": { + "aiDenoise": "AI 노이즈 제거", + "aiDenoiseGenerate": "AI 노이즈 제거 생성", + "aiDenoiseGenerating": "생성 중...", + "aiDenoiseRawOnly": "RAW 파일에서만 사용할 수 있습니다", + "aiDenoiseRemove": "AI 노이즈 제거 삭제", "blueYellow": "파랑/노랑", "centre": "중앙", "chromaticAberration": "색수차", diff --git a/src/i18n/locales/nl.json b/src/i18n/locales/nl.json index 071121d0f4..025d21a0b8 100644 --- a/src/i18n/locales/nl.json +++ b/src/i18n/locales/nl.json @@ -95,6 +95,11 @@ "resetPoint": "{{channel}} puntcurve herstellen" }, "details": { + "aiDenoise": "AI-ruisonderdrukking", + "aiDenoiseGenerate": "AI-ruisonderdrukking genereren", + "aiDenoiseGenerating": "Genereren...", + "aiDenoiseRawOnly": "Alleen beschikbaar voor RAW-bestanden", + "aiDenoiseRemove": "AI-ruisonderdrukking verwijderen", "blueYellow": "Blauw/Geel", "centre": "Centrum", "chromaticAberration": "Chromatische aberratie", diff --git a/src/i18n/locales/pl.json b/src/i18n/locales/pl.json index fe93124699..a1420519f1 100644 --- a/src/i18n/locales/pl.json +++ b/src/i18n/locales/pl.json @@ -96,6 +96,11 @@ "resetPoint": "Resetuj krzywą punktową ({{channel}})" }, "details": { + "aiDenoise": "Odszumianie AI", + "aiDenoiseGenerate": "Generuj odszumianie AI", + "aiDenoiseGenerating": "Generowanie...", + "aiDenoiseRawOnly": "Dostępne tylko dla plików RAW", + "aiDenoiseRemove": "Usuń odszumianie AI", "blueYellow": "Niebieski/Żółty", "centre": "Środek", "chromaticAberration": "Aberracja chromatyczna", diff --git a/src/i18n/locales/pt.json b/src/i18n/locales/pt.json index beebc888c6..c1aae5ea4c 100644 --- a/src/i18n/locales/pt.json +++ b/src/i18n/locales/pt.json @@ -96,6 +96,11 @@ "resetPoint": "Redefinir Curva de Pontos {{channel}}" }, "details": { + "aiDenoise": "Redução de ruído com IA", + "aiDenoiseGenerate": "Gerar redução de ruído com IA", + "aiDenoiseGenerating": "Gerando...", + "aiDenoiseRawOnly": "Disponível apenas para arquivos RAW", + "aiDenoiseRemove": "Remover a redução de ruído com IA", "blueYellow": "Azul/Amarelo", "centre": "Centro", "chromaticAberration": "Aberração Cromática", diff --git a/src/i18n/locales/ru.json b/src/i18n/locales/ru.json index b6f854748a..ea929f04ec 100644 --- a/src/i18n/locales/ru.json +++ b/src/i18n/locales/ru.json @@ -96,6 +96,11 @@ "resetPoint": "Сбросить точечную кривую {{channel}}" }, "details": { + "aiDenoise": "ИИ-шумоподавление", + "aiDenoiseGenerate": "Создать ИИ-шумоподавление", + "aiDenoiseGenerating": "Создание...", + "aiDenoiseRawOnly": "Доступно только для файлов RAW", + "aiDenoiseRemove": "Удалить ИИ-шумоподавление", "blueYellow": "Синий/Желтый", "centre": "Центр", "chromaticAberration": "Хроматическая аберрация", diff --git a/src/i18n/locales/zh-CN.json b/src/i18n/locales/zh-CN.json index b80b5a67ef..4f7738b832 100644 --- a/src/i18n/locales/zh-CN.json +++ b/src/i18n/locales/zh-CN.json @@ -96,6 +96,11 @@ "resetPoint": "重置 {{channel}} 点曲线" }, "details": { + "aiDenoise": "AI 降噪", + "aiDenoiseGenerate": "生成 AI 降噪", + "aiDenoiseGenerating": "正在生成...", + "aiDenoiseRawOnly": "仅适用于 RAW 文件", + "aiDenoiseRemove": "移除 AI 降噪", "blueYellow": "蓝/黄", "centre": "居中", "chromaticAberration": "色差", diff --git a/src/i18n/locales/zh-TW.json b/src/i18n/locales/zh-TW.json index 4ec007556e..8de9fbf383 100644 --- a/src/i18n/locales/zh-TW.json +++ b/src/i18n/locales/zh-TW.json @@ -96,6 +96,11 @@ "resetPoint": "重置 {{channel}} 點曲線" }, "details": { + "aiDenoise": "AI 降噪", + "aiDenoiseGenerate": "產生 AI 降噪", + "aiDenoiseGenerating": "正在產生...", + "aiDenoiseRawOnly": "僅適用於 RAW 檔案", + "aiDenoiseRemove": "移除 AI 降噪", "blueYellow": "藍/黃", "centre": "置中", "chromaticAberration": "色差", diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index 75dccc4ae9..9150a0c6fc 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -241,6 +241,7 @@ export interface ParametricCurve { export interface Adjustments { [index: string]: any; + aiDenoiseEnabled: boolean; aiPatches: Array; aspectRatio: number | null; blacks: number; @@ -599,6 +600,7 @@ export const INITIAL_MASK_CONTAINER: MaskContainer = { }; export const INITIAL_ADJUSTMENTS: Adjustments = { + aiDenoiseEnabled: false, aiPatches: [], aspectRatio: null, blacks: 0, @@ -828,6 +830,7 @@ export const normalizeLoadedAdjustments = (loadedAdjustments: Adjustments): any lensBlurDiffusion: loadedAdjustments.lensBlurDiffusion ?? INITIAL_ADJUSTMENTS.lensBlurDiffusion, lensBlurShape: loadedAdjustments.lensBlurShape ?? INITIAL_ADJUSTMENTS.lensBlurShape, lensBlurDepthMap: loadedAdjustments.lensBlurDepthMap ?? INITIAL_ADJUSTMENTS.lensBlurDepthMap, + aiDenoiseEnabled: loadedAdjustments.aiDenoiseEnabled ?? INITIAL_ADJUSTMENTS.aiDenoiseEnabled, lensBlurEnabled: loadedAdjustments.lensBlurEnabled ?? INITIAL_ADJUSTMENTS.lensBlurEnabled, lensBlurMaxDepth: loadedAdjustments.lensBlurMaxDepth ?? INITIAL_ADJUSTMENTS.lensBlurMaxDepth, lensBlurMaxFade: loadedAdjustments.lensBlurMaxFade ?? INITIAL_ADJUSTMENTS.lensBlurMaxFade, From ccbefb481d8ec77032806f09ef3fdda345d5331e Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Fri, 9 Oct 2026 15:01:54 -0500 Subject: [PATCH 14/16] add amount, color and detail controls to the ai denoise layer Blend the generated denoise layer with the original image so the effect can be dialed back without regenerating it: separate luminance and color amounts plus an edge-weighted detail recovery. The blend is applied to the loaded base image and to exports. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01JDN3t3qVMHMPyx7kWpwfQU --- src-tauri/src/denoising.rs | 181 ++++++++++++++++++++++--- src-tauri/src/image_loader.rs | 18 ++- src/components/adjustments/Details.tsx | 37 +++++ src/i18n/locales/ca.json | 3 + src/i18n/locales/cs.json | 3 + src/i18n/locales/de.json | 3 + src/i18n/locales/en.json | 3 + src/i18n/locales/es.json | 3 + src/i18n/locales/fr.json | 3 + src/i18n/locales/it.json | 3 + src/i18n/locales/ja.json | 3 + src/i18n/locales/ko.json | 3 + src/i18n/locales/nl.json | 3 + src/i18n/locales/pl.json | 3 + src/i18n/locales/pt.json | 3 + src/i18n/locales/ru.json | 3 + src/i18n/locales/zh-CN.json | 3 + src/i18n/locales/zh-TW.json | 3 + src/utils/adjustments.ts | 9 ++ 19 files changed, 265 insertions(+), 25 deletions(-) diff --git a/src-tauri/src/denoising.rs b/src-tauri/src/denoising.rs index 4adfcd1390..4ad9c0a82c 100644 --- a/src-tauri/src/denoising.rs +++ b/src-tauri/src/denoising.rs @@ -1125,10 +1125,41 @@ const AI_DENOISE_TILE_QUALITY: f32 = 0.5; type DenoiseLayerEntry = (String, Arc); +#[derive(Clone, Copy, PartialEq)] +pub struct DenoiseBlend { + luma: f32, + color: f32, + detail: f32, +} + +impl DenoiseBlend { + pub fn from_adjustments(adjustments: &serde_json::Value) -> Self { + let read = |key: &str, default: f64| { + (adjustments + .get(key) + .and_then(|v| v.as_f64()) + .unwrap_or(default) + / 100.0) + .clamp(0.0, 1.0) as f32 + }; + Self { + luma: read("aiDenoiseAmount", 100.0), + color: read("aiDenoiseColor", 100.0), + detail: read("aiDenoiseDetail", 0.0), + } + } + + pub fn is_full(&self) -> bool { + self.luma >= 1.0 && self.color >= 1.0 && self.detail <= 0.0 + } +} + struct ActiveDenoiseLayer { path: String, pristine: Arc, denoised: Arc, + applied: Arc, + blend: DenoiseBlend, } static DENOISE_LAYER_APP_HANDLE: OnceLock = OnceLock::new(); @@ -1207,7 +1238,7 @@ fn restore_pristine_base(state: &AppState) { .unwrap_or_else(|e| e.into_inner()); if let Some(loaded) = guard.as_mut() && loaded.path == active.path - && Arc::ptr_eq(&loaded.image, &active.denoised) + && Arc::ptr_eq(&loaded.image, &active.applied) { loaded.image = active.pristine; restored = true; @@ -1227,6 +1258,94 @@ fn clear_denoise_layer_caches(state: &AppState) { clear_base_image_caches(state); } +const DETAIL_EDGE_LOW: f32 = 0.004; +const DETAIL_EDGE_HIGH: f32 = 0.03; + +fn float_pixels(image: &DynamicImage) -> Option<(&[f32], usize)> { + match image { + DynamicImage::ImageRgb32F(buffer) => Some((buffer.as_raw().as_slice(), 3)), + DynamicImage::ImageRgba32F(buffer) => Some((buffer.as_raw().as_slice(), 4)), + _ => None, + } +} + +pub fn blend_denoised( + pristine: &Arc, + denoised: &Arc, + blend: DenoiseBlend, +) -> Arc { + if blend.is_full() || pristine.dimensions() != denoised.dimensions() { + return Arc::clone(denoised); + } + let (Some((source_raw, source_channels)), Some((clean_raw, channels))) = + (float_pixels(pristine), float_pixels(denoised)) + else { + return Arc::clone(denoised); + }; + if source_channels != channels { + return Arc::clone(denoised); + } + + let (image_width, image_height) = denoised.dimensions(); + let (width, height) = (image_width as usize, image_height as usize); + let perceptual_luma = |x: usize, y: usize| { + let i = (y * width + x) * channels; + (0.2126 * clean_raw[i] + 0.7152 * clean_raw[i + 1] + 0.0722 * clean_raw[i + 2]) + .max(0.0) + .sqrt() + }; + + let mut output = clean_raw.to_vec(); + output + .par_chunks_mut(width * channels) + .zip(source_raw.par_chunks(width * channels)) + .enumerate() + .for_each(|(y, (out_row, source_row))| { + for x in 0..width { + let edge = if blend.detail > 0.0 { + let (x0, x1) = (x.saturating_sub(1), (x + 1).min(width - 1)); + let (y0, y1) = (y.saturating_sub(1), (y + 1).min(height - 1)); + let dx = perceptual_luma(x1, y) - perceptual_luma(x0, y); + let dy = perceptual_luma(x, y1) - perceptual_luma(x, y0); + let gradient = (dx * dx + dy * dy).sqrt(); + let t = ((gradient - DETAIL_EDGE_LOW) / (DETAIL_EDGE_HIGH - DETAIL_EDGE_LOW)) + .clamp(0.0, 1.0); + t * t * (3.0 - 2.0 * t) + } else { + 0.0 + }; + + let i = x * channels; + let residual = [ + source_row[i] - out_row[i], + source_row[i + 1] - out_row[i + 1], + source_row[i + 2] - out_row[i + 2], + ]; + let residual_luma = + 0.2126 * residual[0] + 0.7152 * residual[1] + 0.0722 * residual[2]; + let keep_luma = + ((1.0 - blend.luma) + blend.luma * blend.detail * edge).clamp(0.0, 1.0); + let keep_color = 1.0 - blend.color; + + for c in 0..3 { + out_row[i + c] += + keep_luma * residual_luma + keep_color * (residual[c] - residual_luma); + } + } + }); + + let blended = if channels == 3 { + Rgb32FImage::from_raw(image_width, image_height, output).map(DynamicImage::ImageRgb32F) + } else { + image::Rgba32FImage::from_raw(image_width, image_height, output) + .map(DynamicImage::ImageRgba32F) + }; + match blended { + Some(image) => Arc::new(image), + None => Arc::clone(denoised), + } +} + pub fn sync_denoised_base(state: &AppState, adjustments: &serde_json::Value) { let Some((path, current)) = state .original_image @@ -1238,31 +1357,47 @@ pub fn sync_denoised_base(state: &AppState, adjustments: &serde_json::Value) { return; }; - let is_active = ACTIVE_DENOISE_LAYER - .lock() - .unwrap_or_else(|e| e.into_inner()) - .as_ref() - .is_some_and(|active| active.path == path && Arc::ptr_eq(&active.denoised, ¤t)); + let blend = DenoiseBlend::from_adjustments(adjustments); let wants_denoised = ai_denoise_enabled(adjustments); - if !is_active { - ACTIVE_DENOISE_LAYER + let active_sources = { + let mut active_lock = ACTIVE_DENOISE_LAYER .lock() - .unwrap_or_else(|e| e.into_inner()) - .take(); - } - if wants_denoised == is_active { - return; - } - if !wants_denoised { - restore_pristine_base(state); - return; - } + .unwrap_or_else(|e| e.into_inner()); + let is_active = active_lock + .as_ref() + .is_some_and(|active| active.path == path && Arc::ptr_eq(&active.applied, ¤t)); + if !is_active { + active_lock.take(); + } + active_lock.as_ref().map(|active| { + ( + Arc::clone(&active.pristine), + Arc::clone(&active.denoised), + active.blend, + ) + }) + }; - let Some(denoised) = denoised_base(&path, adjustments, Some(current.dimensions())) else { - return; + let (pristine, denoised) = match (wants_denoised, active_sources) { + (false, None) => return, + (false, Some(_)) => { + restore_pristine_base(state); + return; + } + (true, Some((_, _, active_blend))) if active_blend == blend => return, + (true, Some((pristine, denoised, _))) => (pristine, denoised), + (true, None) => { + let Some(denoised) = denoised_base(&path, adjustments, Some(current.dimensions())) + else { + return; + }; + (Arc::clone(¤t), denoised) + } }; + let applied = blend_denoised(&pristine, &denoised, blend); + let mut swapped = false; { let mut guard = state @@ -1273,7 +1408,7 @@ pub fn sync_denoised_base(state: &AppState, adjustments: &serde_json::Value) { && loaded.path == path && Arc::ptr_eq(&loaded.image, ¤t) { - loaded.image = Arc::clone(&denoised); + loaded.image = Arc::clone(&applied); swapped = true; } } @@ -1283,8 +1418,10 @@ pub fn sync_denoised_base(state: &AppState, adjustments: &serde_json::Value) { .lock() .unwrap_or_else(|e| e.into_inner()) = Some(ActiveDenoiseLayer { path, - pristine: current, + pristine, denoised, + applied, + blend, }); clear_base_image_caches(state); } diff --git a/src-tauri/src/image_loader.rs b/src-tauri/src/image_loader.rs index a1538096b1..78f675bc47 100644 --- a/src-tauri/src/image_loader.rs +++ b/src-tauri/src/image_loader.rs @@ -74,14 +74,26 @@ pub fn load_and_composite( settings: &AppSettings, cancel_token: Option<(Arc, usize)>, ) -> Result { - if !use_fast_raw_dev - && let Some(denoised) = crate::denoising::denoised_base(path, adjustments, None) + let denoised = if use_fast_raw_dev { + None + } else { + crate::denoising::denoised_base(path, adjustments, None) + }; + let blend = crate::denoising::DenoiseBlend::from_adjustments(adjustments); + if let Some(denoised) = &denoised + && blend.is_full() { - return composite_patches_on_image(&denoised, adjustments); + return composite_patches_on_image(denoised, adjustments); } let base_image = load_base_image_from_bytes(base_image, path, use_fast_raw_dev, settings, cancel_token)?; + if let Some(denoised) = denoised + && denoised.dimensions() == base_image.dimensions() + { + let blended = crate::denoising::blend_denoised(&Arc::new(base_image), &denoised, blend); + return composite_patches_on_image(&blended, adjustments); + } composite_patches_on_image(&base_image, adjustments) } diff --git a/src/components/adjustments/Details.tsx b/src/components/adjustments/Details.tsx index 6d9cce2ceb..610cb4434e 100644 --- a/src/components/adjustments/Details.tsx +++ b/src/components/adjustments/Details.tsx @@ -258,6 +258,43 @@ export default function DetailsPanel({ )} + {hasDenoiseLayer && adjustments.aiDenoiseEnabled && ( + <> + handleAdjustmentChange('aiDenoiseAmount', e.target.value)} + step={1} + value={adjustments.aiDenoiseAmount ?? 100} + defaultValue={100} + fillOrigin="min" + onDragStateChange={onDragStateChange} + /> + handleAdjustmentChange('aiDenoiseColor', e.target.value)} + step={1} + value={adjustments.aiDenoiseColor ?? 100} + defaultValue={100} + fillOrigin="min" + onDragStateChange={onDragStateChange} + /> + handleAdjustmentChange('aiDenoiseDetail', e.target.value)} + step={1} + value={adjustments.aiDenoiseDetail ?? 0} + defaultValue={0} + fillOrigin="min" + onDragStateChange={onDragStateChange} + /> + + )} {denoiseError && ( {denoiseError} diff --git a/src/i18n/locales/ca.json b/src/i18n/locales/ca.json index fe2abd3d85..373788e116 100644 --- a/src/i18n/locales/ca.json +++ b/src/i18n/locales/ca.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "Reducció de soroll amb IA", + "aiDenoiseAmount": "Quantitat (IA)", + "aiDenoiseColor": "Soroll de color (IA)", + "aiDenoiseDetail": "Recuperació de detall", "aiDenoiseGenerate": "Generar reducció de soroll amb IA", "aiDenoiseGenerating": "Generant...", "aiDenoiseRawOnly": "Només disponible per a fitxers RAW", diff --git a/src/i18n/locales/cs.json b/src/i18n/locales/cs.json index 91324fd05b..2823b6f586 100644 --- a/src/i18n/locales/cs.json +++ b/src/i18n/locales/cs.json @@ -96,6 +96,9 @@ }, "details": { "aiDenoise": "AI odšumění", + "aiDenoiseAmount": "Míra (AI)", + "aiDenoiseColor": "Barevný šum (AI)", + "aiDenoiseDetail": "Obnovení detailů", "aiDenoiseGenerate": "Vygenerovat AI odšumění", "aiDenoiseGenerating": "Generování...", "aiDenoiseRawOnly": "Dostupné pouze pro soubory RAW", diff --git a/src/i18n/locales/de.json b/src/i18n/locales/de.json index 7722ba690c..1f5f57cee9 100644 --- a/src/i18n/locales/de.json +++ b/src/i18n/locales/de.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "KI-Entrauschen", + "aiDenoiseAmount": "Stärke (KI)", + "aiDenoiseColor": "Farbrauschen (KI)", + "aiDenoiseDetail": "Detailwiederherstellung", "aiDenoiseGenerate": "KI-Entrauschen erzeugen", "aiDenoiseGenerating": "Wird erzeugt...", "aiDenoiseRawOnly": "Nur für RAW-Dateien verfügbar", diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index 4b94ae5a80..655fa3ec09 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "AI Denoise", + "aiDenoiseAmount": "Amount (AI)", + "aiDenoiseColor": "Color Noise (AI)", + "aiDenoiseDetail": "Detail Recovery", "aiDenoiseGenerate": "Generate AI Denoise", "aiDenoiseGenerating": "Generating...", "aiDenoiseRawOnly": "Only available for RAW files", diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index 986d3b9224..d0d7fb6d08 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "Reducción de ruido con IA", + "aiDenoiseAmount": "Cantidad (IA)", + "aiDenoiseColor": "Ruido de color (IA)", + "aiDenoiseDetail": "Recuperación de detalle", "aiDenoiseGenerate": "Generar reducción de ruido con IA", "aiDenoiseGenerating": "Generando...", "aiDenoiseRawOnly": "Solo disponible para archivos RAW", diff --git a/src/i18n/locales/fr.json b/src/i18n/locales/fr.json index 5a7280b162..df99cdd298 100644 --- a/src/i18n/locales/fr.json +++ b/src/i18n/locales/fr.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "Débruitage IA", + "aiDenoiseAmount": "Quantité (IA)", + "aiDenoiseColor": "Bruit de couleur (IA)", + "aiDenoiseDetail": "Récupération des détails", "aiDenoiseGenerate": "Générer le débruitage IA", "aiDenoiseGenerating": "Génération...", "aiDenoiseRawOnly": "Disponible uniquement pour les fichiers RAW", diff --git a/src/i18n/locales/it.json b/src/i18n/locales/it.json index 7b9507eda1..425996e7d0 100644 --- a/src/i18n/locales/it.json +++ b/src/i18n/locales/it.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "Riduzione rumore IA", + "aiDenoiseAmount": "Quantità (IA)", + "aiDenoiseColor": "Rumore colore (IA)", + "aiDenoiseDetail": "Recupero dettagli", "aiDenoiseGenerate": "Genera riduzione rumore IA", "aiDenoiseGenerating": "Generazione...", "aiDenoiseRawOnly": "Disponibile solo per i file RAW", diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index 3f4a473ed1..9d6e4e46ad 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "AIノイズ低減", + "aiDenoiseAmount": "適用量 (AI)", + "aiDenoiseColor": "カラーノイズ (AI)", + "aiDenoiseDetail": "ディテールの復元", "aiDenoiseGenerate": "AIノイズ低減を生成", "aiDenoiseGenerating": "生成中...", "aiDenoiseRawOnly": "RAWファイルでのみ利用できます", diff --git a/src/i18n/locales/ko.json b/src/i18n/locales/ko.json index 431b440076..743ad635e0 100644 --- a/src/i18n/locales/ko.json +++ b/src/i18n/locales/ko.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "AI 노이즈 제거", + "aiDenoiseAmount": "적용량 (AI)", + "aiDenoiseColor": "색상 노이즈 (AI)", + "aiDenoiseDetail": "디테일 복원", "aiDenoiseGenerate": "AI 노이즈 제거 생성", "aiDenoiseGenerating": "생성 중...", "aiDenoiseRawOnly": "RAW 파일에서만 사용할 수 있습니다", diff --git a/src/i18n/locales/nl.json b/src/i18n/locales/nl.json index 025d21a0b8..82c48885c1 100644 --- a/src/i18n/locales/nl.json +++ b/src/i18n/locales/nl.json @@ -96,6 +96,9 @@ }, "details": { "aiDenoise": "AI-ruisonderdrukking", + "aiDenoiseAmount": "Hoeveelheid (AI)", + "aiDenoiseColor": "Kleurruis (AI)", + "aiDenoiseDetail": "Detailherstel", "aiDenoiseGenerate": "AI-ruisonderdrukking genereren", "aiDenoiseGenerating": "Genereren...", "aiDenoiseRawOnly": "Alleen beschikbaar voor RAW-bestanden", diff --git a/src/i18n/locales/pl.json b/src/i18n/locales/pl.json index a1420519f1..bd61ce02f7 100644 --- a/src/i18n/locales/pl.json +++ b/src/i18n/locales/pl.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "Odszumianie AI", + "aiDenoiseAmount": "Ilość (AI)", + "aiDenoiseColor": "Szum kolorów (AI)", + "aiDenoiseDetail": "Odzyskiwanie szczegółów", "aiDenoiseGenerate": "Generuj odszumianie AI", "aiDenoiseGenerating": "Generowanie...", "aiDenoiseRawOnly": "Dostępne tylko dla plików RAW", diff --git a/src/i18n/locales/pt.json b/src/i18n/locales/pt.json index c1aae5ea4c..10c8e1c345 100644 --- a/src/i18n/locales/pt.json +++ b/src/i18n/locales/pt.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "Redução de ruído com IA", + "aiDenoiseAmount": "Quantidade (IA)", + "aiDenoiseColor": "Ruído de cor (IA)", + "aiDenoiseDetail": "Recuperação de detalhe", "aiDenoiseGenerate": "Gerar redução de ruído com IA", "aiDenoiseGenerating": "Gerando...", "aiDenoiseRawOnly": "Disponível apenas para arquivos RAW", diff --git a/src/i18n/locales/ru.json b/src/i18n/locales/ru.json index ea929f04ec..3ba09e194a 100644 --- a/src/i18n/locales/ru.json +++ b/src/i18n/locales/ru.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "ИИ-шумоподавление", + "aiDenoiseAmount": "Степень (ИИ)", + "aiDenoiseColor": "Цветовой шум (ИИ)", + "aiDenoiseDetail": "Восстановление деталей", "aiDenoiseGenerate": "Создать ИИ-шумоподавление", "aiDenoiseGenerating": "Создание...", "aiDenoiseRawOnly": "Доступно только для файлов RAW", diff --git a/src/i18n/locales/zh-CN.json b/src/i18n/locales/zh-CN.json index 4f7738b832..550c38f59b 100644 --- a/src/i18n/locales/zh-CN.json +++ b/src/i18n/locales/zh-CN.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "AI 降噪", + "aiDenoiseAmount": "数量 (AI)", + "aiDenoiseColor": "颜色噪点 (AI)", + "aiDenoiseDetail": "细节恢复", "aiDenoiseGenerate": "生成 AI 降噪", "aiDenoiseGenerating": "正在生成...", "aiDenoiseRawOnly": "仅适用于 RAW 文件", diff --git a/src/i18n/locales/zh-TW.json b/src/i18n/locales/zh-TW.json index 8de9fbf383..001a002da7 100644 --- a/src/i18n/locales/zh-TW.json +++ b/src/i18n/locales/zh-TW.json @@ -97,6 +97,9 @@ }, "details": { "aiDenoise": "AI 降噪", + "aiDenoiseAmount": "總量 (AI)", + "aiDenoiseColor": "顏色雜訊 (AI)", + "aiDenoiseDetail": "細節復原", "aiDenoiseGenerate": "產生 AI 降噪", "aiDenoiseGenerating": "正在產生...", "aiDenoiseRawOnly": "僅適用於 RAW 檔案", diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index 9150a0c6fc..861c40f986 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -241,6 +241,9 @@ export interface ParametricCurve { export interface Adjustments { [index: string]: any; + aiDenoiseAmount: number; + aiDenoiseColor: number; + aiDenoiseDetail: number; aiDenoiseEnabled: boolean; aiPatches: Array; aspectRatio: number | null; @@ -600,6 +603,9 @@ export const INITIAL_MASK_CONTAINER: MaskContainer = { }; export const INITIAL_ADJUSTMENTS: Adjustments = { + aiDenoiseAmount: 100, + aiDenoiseColor: 100, + aiDenoiseDetail: 0, aiDenoiseEnabled: false, aiPatches: [], aspectRatio: null, @@ -830,6 +836,9 @@ export const normalizeLoadedAdjustments = (loadedAdjustments: Adjustments): any lensBlurDiffusion: loadedAdjustments.lensBlurDiffusion ?? INITIAL_ADJUSTMENTS.lensBlurDiffusion, lensBlurShape: loadedAdjustments.lensBlurShape ?? INITIAL_ADJUSTMENTS.lensBlurShape, lensBlurDepthMap: loadedAdjustments.lensBlurDepthMap ?? INITIAL_ADJUSTMENTS.lensBlurDepthMap, + aiDenoiseAmount: loadedAdjustments.aiDenoiseAmount ?? INITIAL_ADJUSTMENTS.aiDenoiseAmount, + aiDenoiseColor: loadedAdjustments.aiDenoiseColor ?? INITIAL_ADJUSTMENTS.aiDenoiseColor, + aiDenoiseDetail: loadedAdjustments.aiDenoiseDetail ?? INITIAL_ADJUSTMENTS.aiDenoiseDetail, aiDenoiseEnabled: loadedAdjustments.aiDenoiseEnabled ?? INITIAL_ADJUSTMENTS.aiDenoiseEnabled, lensBlurEnabled: loadedAdjustments.lensBlurEnabled ?? INITIAL_ADJUSTMENTS.lensBlurEnabled, lensBlurMaxDepth: loadedAdjustments.lensBlurMaxDepth ?? INITIAL_ADJUSTMENTS.lensBlurMaxDepth, From 53cffa2f13e387ece071f0f5c6ffc9fa0c9119ca Mon Sep 17 00:00:00 2001 From: Fabian Grajales Date: Fri, 9 Oct 2026 20:28:56 -0500 Subject: [PATCH 15/16] feat: add AI-powered facial refinement, skin retouching, and inpainting features with multilingual localization --- src-tauri/src/ai_processing.rs | 97 ++ src-tauri/src/face_parsing.rs | 767 ++++++++++++ src-tauri/src/face_refine.rs | 458 ++++++++ src-tauri/src/inpainting.rs | 151 +-- src-tauri/src/lib.rs | 5 + src-tauri/src/skin_retouch.rs | 1169 +++++++++++++++++++ src/components/adjustments/Effects.tsx | 15 + src/components/adjustments/FacialRefine.tsx | 378 ++++++ src/components/panel/right/AIPanel.tsx | 40 +- src/i18n/locales/en.json | 26 +- src/i18n/locales/es.json | 26 +- src/utils/adjustments.ts | 11 + 12 files changed, 3015 insertions(+), 128 deletions(-) create mode 100644 src-tauri/src/face_parsing.rs create mode 100644 src-tauri/src/face_refine.rs create mode 100644 src-tauri/src/skin_retouch.rs create mode 100644 src/components/adjustments/FacialRefine.tsx diff --git a/src-tauri/src/ai_processing.rs b/src-tauri/src/ai_processing.rs index 7f822c9dde..103d1c991b 100644 --- a/src-tauri/src/ai_processing.rs +++ b/src-tauri/src/ai_processing.rs @@ -19,6 +19,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; + const ENCODER_URL: &str = "https://huggingface.co/CyberTimon/RapidRAW-Models/resolve/main/sam_vit_b_01ec64_encoder.onnx?download=true"; const DECODER_URL: &str = "https://huggingface.co/CyberTimon/RapidRAW-Models/resolve/main/sam_vit_b_01ec64_decoder.onnx?download=true"; const ENCODER_FILENAME: &str = "sam_vit_b_01ec64_encoder.onnx"; @@ -74,6 +77,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"; + pub struct AiModels { pub sam_encoder: Mutex, pub sam_decoder: Mutex, @@ -109,6 +122,8 @@ pub struct AiState { pub normal_model: Option>>, pub embeddings: Option, pub depth_map: Option, + pub face_models: Option>, + pub face_crops: Option>, } fn edt_1d(f: &mut [f32], v: &mut [usize], z: &mut [f32], d: &mut [f32]) { @@ -568,6 +583,8 @@ pub async fn get_or_init_ai_models( normal_model: None, embeddings: None, depth_map: None, + face_models: None, + face_crops: None, }); } @@ -727,6 +744,8 @@ pub async fn get_or_init_denoise_model( normal_model: None, embeddings: None, depth_map: None, + face_models: None, + face_crops: None, }); } @@ -800,6 +819,8 @@ pub async fn get_or_init_clip_models( normal_model: None, embeddings: None, depth_map: None, + face_models: None, + face_crops: None, }); } @@ -861,12 +882,86 @@ pub async fn get_or_init_lama_model( normal_model: None, embeddings: None, depth_map: None, + face_models: None, + face_crops: None, }); } 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_FILENAME, + FACE_DETECTOR_URL, + FACE_DETECTOR_SHA256, + "Face Detection Model", + ) + .await?; + download_and_verify_model( + app_handle, + &models_dir, + FACE_PARSER_FILENAME, + FACE_PARSER_URL, + FACE_PARSER_SHA256, + "Face Regions Model", + ) + .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(); + + let mut ai_state_lock = ai_state_mutex.lock().unwrap(); + if let Some(state) = ai_state_lock.as_mut() { + state.face_models = Some(face_models.clone()); + } else { + *ai_state_lock = Some(AiState { + models: None, + denoise_model: None, + clip_models: None, + lama_model: None, + normal_model: None, + embeddings: None, + depth_map: None, + face_models: Some(face_models.clone()), + face_crops: None, + }); + } + + Ok(face_models) +} + pub async fn get_or_init_normal_model( app_handle: &tauri::AppHandle, ai_state_mutex: &Mutex>, @@ -922,6 +1017,8 @@ pub async fn get_or_init_normal_model( normal_model: Some(normal_model.clone()), embeddings: None, depth_map: None, + face_models: None, + face_crops: None, }); } 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/inpainting.rs b/src-tauri/src/inpainting.rs index 2751e7cdab..3cd188bb24 100644 --- a/src-tauri/src/inpainting.rs +++ b/src-tauri/src/inpainting.rs @@ -13,8 +13,9 @@ use crate::image_loader::composite_patches_on_image; 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>, @@ -103,7 +104,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, @@ -1137,7 +1138,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]; @@ -1147,10 +1155,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 { @@ -1232,7 +1250,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; @@ -1252,112 +1272,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 dfbc481edc..d9fff45816 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -37,6 +37,8 @@ mod denoising; mod effect_maps; mod exif_processing; mod export_processing; +mod face_parsing; +mod face_refine; mod file_management; mod focus_stacking; mod fog; @@ -59,6 +61,7 @@ mod panorama_utils; mod preset_converter; mod raw_processing; mod relight; +mod skin_retouch; mod tagging; mod tagging_utils; mod white_balance; @@ -2389,6 +2392,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, denoising::apply_denoising, denoising::batch_denoise_images, denoising::save_denoised_image, 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/components/adjustments/Effects.tsx b/src/components/adjustments/Effects.tsx index 734f35dc33..77da0f90bc 100644 --- a/src/components/adjustments/Effects.tsx +++ b/src/components/adjustments/Effects.tsx @@ -24,6 +24,7 @@ import { ColorSwatch } from './Color'; import { AppSettings } from '../ui/AppProperties'; import Text from '../ui/Text'; import AdjustmentSubSection from './AdjustmentSubSection'; +import FacialRefine from './FacialRefine'; import { TextVariants } from '../../types/typography'; import { DepthRangePicker } from '../ui/DepthRangePicker'; import { useProcessStore } from '../../store/useProcessStore'; @@ -885,6 +886,20 @@ export default function EffectsPanel({ )} + {!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((state) => state.aiModelDownloadStatus); + + 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/right/AIPanel.tsx b/src/components/panel/right/AIPanel.tsx index 98bda5ea8e..a90a9bee9b 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); @@ -776,11 +796,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); @@ -974,7 +990,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/i18n/locales/en.json b/src/i18n/locales/en.json index 655fa3ec09..af349e4f85 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -127,6 +127,25 @@ "bokehRing": "Bubble (Ring)", "bokehShape": "Bokeh Shape", "creative": "Creative", + "facialRefine": "Facial Refine", + "facialRefineApplyToAll": "Apply These Settings to All Faces", + "facialRefineDetect": "Detect Faces", + "facialRefineDetecting": "Detecting faces...", + "facialRefineFace": "Face {{index}}", + "facialRefineFailed": "Facial Refine failed: {{error}}", + "facialRefineHint": "Detect the faces in this photo to retouch each one with its own settings.", + "facialRefineNoFaces": "No faces were found in this photo.", + "facialRefineParams": { + "blemish": "Blemish Removal (Forehead & Cheeks)", + "evenTone": "Even Skin Tone", + "eyes": "Eye Enhancement", + "intensity": "Skin Smoothing", + "shine": "Shine & Sweat Reduction", + "teeth": "Teeth Whitening", + "texture": "Skin Texture" + }, + "facialRefineRedetect": "Detect Faces Again", + "facialRefineRemove": "Remove this face", "feather": "Feather", "fog": "Fog", "fogDensity": "Density", @@ -429,10 +448,15 @@ "manualCleanupTitle": "Manual Cleanup", "noImageSelected": "No image selected.", "params": { + "blemish": "Blemish Removal", + "evenTone": "Even Skin Tone", "feather": "Feather", "grow": "Grow", "intensity": "Intensity", - "pressure": "Pressure" + "pressure": "Pressure", + "shine": "Shine & Sweat Reduction", + "skinProtection": "Protect Eyes, Lips & Hair", + "texture": "Skin Texture" }, "patches": { "aiEdit": "AI Edit {{count}}", diff --git a/src/i18n/locales/es.json b/src/i18n/locales/es.json index d0d7fb6d08..3da19314b3 100644 --- a/src/i18n/locales/es.json +++ b/src/i18n/locales/es.json @@ -127,6 +127,25 @@ "bokehRing": "Burbuja (Anillo)", "bokehShape": "Forma del bokeh", "creative": "Creativo", + "facialRefine": "Facial Refine", + "facialRefineApplyToAll": "Aplicar estos ajustes a todas las caras", + "facialRefineDetect": "Detectar caras", + "facialRefineDetecting": "Detectando caras...", + "facialRefineFace": "Cara {{index}}", + "facialRefineFailed": "Facial Refine falló: {{error}}", + "facialRefineHint": "Detecta las caras de esta foto para retocar cada una con sus propios ajustes.", + "facialRefineNoFaces": "No se encontraron caras en esta foto.", + "facialRefineParams": { + "blemish": "Eliminar imperfecciones (frente y pómulos)", + "evenTone": "Unificar tono de piel", + "eyes": "Realzar ojos", + "intensity": "Suavizar piel", + "shine": "Reducir brillos y sudor", + "teeth": "Blanquear dientes", + "texture": "Textura de la piel" + }, + "facialRefineRedetect": "Volver a detectar caras", + "facialRefineRemove": "Quitar esta cara", "feather": "Desvanecimiento", "fog": "Niebla", "fogDensity": "Densidad", @@ -445,10 +464,15 @@ "manualCleanupTitle": "Limpieza manual", "noImageSelected": "No hay imagen seleccionada.", "params": { + "blemish": "Eliminar imperfecciones", + "evenTone": "Unificar tono de piel", "feather": "Desvanecimiento", "grow": "Expandir", "intensity": "Intensidad", - "pressure": "Presión" + "pressure": "Presión", + "shine": "Reducir brillos y sudor", + "skinProtection": "Proteger ojos, labios y pelo", + "texture": "Textura de la piel" }, "patches": { "aiEdit": "Edición IA {{count}}", diff --git a/src/utils/adjustments.ts b/src/utils/adjustments.ts index 861c40f986..ee93920645 100644 --- a/src/utils/adjustments.ts +++ b/src/utils/adjustments.ts @@ -361,6 +361,8 @@ export interface Adjustments { } export interface AiPatch { + /** Set on the per-face patches that the Facial Refine effect creates and manages. */ + faceRefine?: boolean; id: string; isLoading: boolean; invert: boolean; @@ -371,6 +373,14 @@ export interface AiPatch { visible: boolean; } +/** Inserts a patch at `index` (or last) among the regular patches, keeping Facial Refine patches at the end. */ +export const insertAiPatch = (patches: Array, patch: AiPatch, index?: number): Array => { + const regular = patches.filter((p) => !p.faceRefine); + const target = Math.max(0, Math.min(index ?? regular.length, regular.length)); + regular.splice(target, 0, patch); + return [...regular, ...patches.filter((p) => p.faceRefine)]; +}; + export interface Color { color: string; name: string; @@ -1158,6 +1168,7 @@ export const ADJUSTMENT_SECTION_TOOLS: Record Date: Sat, 31 Oct 2026 13:19:01 -0500 Subject: [PATCH 16/16] feat(library): add picked and rejected views with counts to the sources panel Add a Flags section to the folder tree that lists the flagged images of the current folder and its subfolders, or of the current album, with a live count for each flag. The counts come from a new cancellable background scan so that selecting a folder does not block the window. Co-Authored-By: Claude Opus 5.5 --- src-tauri/src/file_management.rs | 209 ++++++++++++++++++++-- src-tauri/src/lib.rs | 1 + src/App.tsx | 3 + src/components/panel/MainLibrary.tsx | 21 ++- src/components/panel/right/FolderTree.tsx | 107 ++++++++++- src/components/ui/AppProperties.tsx | 1 + src/hooks/useAppNavigation.ts | 111 ++++++++++-- src/hooks/useSortedLibrary.ts | 11 +- src/i18n/locales/ca.json | 6 + src/i18n/locales/cs.json | 6 + src/i18n/locales/de.json | 6 + src/i18n/locales/en.json | 6 + src/i18n/locales/es.json | 6 + src/i18n/locales/fr.json | 6 + src/i18n/locales/it.json | 6 + src/i18n/locales/ja.json | 6 + src/i18n/locales/ko.json | 6 + src/i18n/locales/nl.json | 6 + src/i18n/locales/pl.json | 6 + src/i18n/locales/pt.json | 6 + src/i18n/locales/ru.json | 6 + src/i18n/locales/zh-CN.json | 6 + src/i18n/locales/zh-TW.json | 6 + src/store/useLibraryStore.ts | 20 ++- src/utils/imageFlags.ts | 31 ++++ 25 files changed, 569 insertions(+), 36 deletions(-) diff --git a/src-tauri/src/file_management.rs b/src-tauri/src/file_management.rs index 39325e564e..2c65cba2db 100644 --- a/src-tauri/src/file_management.rs +++ b/src-tauri/src/file_management.rs @@ -9,7 +9,7 @@ use std::io::Cursor; use std::path::{Path, PathBuf}; 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; @@ -692,22 +692,21 @@ pub fn list_images_in_dir(path: String, app_handle: AppHandle) -> Result Result, String> { - let settings = load_settings(app_handle.clone()).unwrap_or_default(); - let enable_xmp_sync = settings.enable_xmp_sync.unwrap_or(false); +type RecursiveImageTask = (String, String, PathBuf, Vec>); - update_rotational_disk_flag(&path, &app_handle); - - let root_path = Path::new(&path); +fn collect_recursive_image_tasks( + root_path: &Path, + is_cancelled: &(dyn Fn() -> bool + Sync), +) -> Vec { 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 Vec::new(); + } + let entry_path = entry.path(); if !entry_path.is_file() { continue; @@ -738,7 +737,7 @@ pub fn list_images_recursive( } } - let tasks: Vec<_> = images + images .into_iter() .map(|path_buf| { let sidecars = sidecars_by_path @@ -752,7 +751,89 @@ pub fn list_images_recursive( .into_owned(); (path_str, file_name, path_buf, sidecars) }) - .collect(); + .collect() +} + +// Cloud placeholders are skipped so a background count never forces downloads. +fn collect_flagged_paths( + root_path: &Path, + enable_xmp_sync: bool, + settings: &AppSettings, + is_cancelled: &(dyn Fn() -> bool + Sync), +) -> HashMap { + collect_recursive_image_tasks(root_path, is_cancelled) + .into_par_iter() + .flat_map(|(path_str, file_name, path_buf, sidecars)| { + let mut flagged = Vec::new(); + if is_cancelled() { + return flagged; + } + + let xmp_is_placeholder = enable_xmp_sync + && resolve_xmp_path(&path_buf).is_some_and(|p| is_cloud_placeholder(&p)); + if xmp_is_placeholder { + return flagged; + } + + for copy_id_opt in sidecars { + let (virtual_path, sidecar_filename) = match copy_id_opt { + Some(id) => ( + format!("{}?vc={}", path_str, id), + format!("{}.{}.rrdata", file_name, id), + ), + None => (path_str.clone(), format!("{}.rrdata", file_name)), + }; + + let sidecar_path = path_buf.with_file_name(sidecar_filename); + if is_cloud_placeholder(&sidecar_path) { + continue; + } + + let metadata = + resolve_image_metadata(&path_buf, &sidecar_path, 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(); + let enable_xmp_sync = settings.enable_xmp_sync.unwrap_or(false); + collect_flagged_paths(Path::new(&path), enable_xmp_sync, &settings, &|| { + FLAG_SCAN_GENERATION.load(Ordering::SeqCst) != generation + }) + }) + .await + .map_err(|e| e.to_string()) +} + +#[tauri::command] +pub fn list_images_recursive( + path: String, + app_handle: AppHandle, +) -> Result, String> { + let settings = load_settings(app_handle.clone()).unwrap_or_default(); + let enable_xmp_sync = settings.enable_xmp_sync.unwrap_or(false); + + update_rotational_disk_flag(&path, &app_handle); + + let root_path = Path::new(&path); + let tasks = collect_recursive_image_tasks(root_path, &|| false); let mut result_list: Vec = tasks .into_par_iter() @@ -4356,3 +4437,105 @@ 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(), false, &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"), + false, + &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(), false, &AppSettings::default(), &|| true); + + assert!(flagged.is_empty()); + } + + #[test] + fn listing_and_scan_agree_on_paths() { + let dir = tempfile::tempdir().unwrap(); + create_library(dir.path()); + + let listed: HashSet = collect_recursive_image_tasks(dir.path(), &|| false) + .into_iter() + .map(|(path_str, ..)| path_str) + .collect(); + let flagged = collect_flagged_paths(dir.path(), false, &AppSettings::default(), &|| false); + + for path in flagged.keys() { + let source = path.split("?vc=").next().unwrap(); + assert!(listed.contains(source), "{} was not listed", source); + } + } +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index d9fff45816..e8a01eb933 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2417,6 +2417,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/App.tsx b/src/App.tsx index 804a708027..18f9d76fcb 100644 --- a/src/App.tsx +++ b/src/App.tsx @@ -347,6 +347,7 @@ function App() { handleImageSelect, handleSelectSubfolder, handleSelectAlbum, + handleSelectFlagView, handleOpenFolder, handleNavBack, handleNavForward, @@ -721,6 +722,7 @@ function App() { onFolderSelect={(path, skipHistory) => handleSelectSubfolder(path, false, undefined, true, false, skipHistory) } + onSelectFlagView={handleSelectFlagView} onToggleFolder={handleToggleFolder} onOpenFolder={handleOpenFolder} onNavBack={handleNavBack} @@ -767,6 +769,7 @@ function App() { handleAlbumTreeContextMenu, handleSelectAlbum, handleSelectSubfolder, + handleSelectFlagView, handleToggleFolder, handleOpenFolder, handleNavBack, diff --git a/src/components/panel/MainLibrary.tsx b/src/components/panel/MainLibrary.tsx index 3e95481368..9298dd8015 100644 --- a/src/components/panel/MainLibrary.tsx +++ b/src/components/panel/MainLibrary.tsx @@ -32,6 +32,7 @@ import { RawStatus, EditedStatus, FlagStatus, + ImageFlag, LibraryDisplayMode, } from '../ui/AppProperties'; import { GroupBadgeInfo, GroupId } from '../../utils/imageGrouping'; @@ -184,6 +185,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 translatedRatingFilterOptions = useMemo( () => [ @@ -523,7 +527,15 @@ export default function MainLibrary(props: MainLibraryProps) { {!props.isAndroid && (
{props.currentFolderPath ? ( - {props.currentFolderPath} + + {props.currentFolderPath} + {flagView && + ` · ${t( + flagView === ImageFlag.Pick + ? 'library.folders.flagViews.picked' + : 'library.folders.flagViews.rejected', + )}`} + ) : (

)} @@ -623,6 +635,7 @@ export default function MainLibrary(props: MainLibraryProps) { ) : ( @@ -649,6 +662,10 @@ 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/FolderTree.tsx b/src/components/panel/right/FolderTree.tsx index 6936f3ebaa..af81b7c0cf 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'; export interface FolderTree { children: FolderTree[]; @@ -59,6 +68,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; @@ -102,6 +112,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; @@ -337,6 +352,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} + + )} +