diff --git a/src-tauri/src/image_processing.rs b/src-tauri/src/image_processing.rs index 00f71501a4..5ac707c88f 100644 --- a/src-tauri/src/image_processing.rs +++ b/src-tauri/src/image_processing.rs @@ -1204,15 +1204,16 @@ pub fn apply_cpu_default_raw_processing(image: &mut DynamicImage) { *image = DynamicImage::ImageRgb32F(f32_image); } +fn srgb_channel_to_linear(c: f32) -> f32 { + if c <= 0.04045 { + c / 12.92 + } else { + ((c + 0.055) / 1.055).powf(2.4) + } +} + pub fn apply_srgb_to_linear(mut image: DynamicImage) -> DynamicImage { - let to_linear = |x: f32| -> f32 { - let x = x.max(0.0); - if x <= 0.04045 { - x / 12.92 - } else { - ((x + 0.055) / 1.055).powf(2.4) - } - }; + let to_linear = |x: f32| -> f32 { srgb_channel_to_linear(x.max(0.0)) }; match &mut image { DynamicImage::ImageRgb32F(img) => { @@ -3536,3 +3537,174 @@ pub fn calculate_auto_adjustments( Ok(auto_results_to_json(&results)) } + +#[derive(Deserialize, Clone, Copy)] +pub struct UvPoint { + pub x: f64, + pub y: f64, +} + +#[derive(Serialize, Debug)] +#[serde(rename_all = "camelCase")] +pub struct WhiteBalanceSample { + pub r: f32, + pub g: f32, + pub b: f32, + pub temperature: f32, + pub tint: f32, + pub count: u32, +} + +const MAX_WB_SAMPLES: f64 = 262_144.0; + +fn read_linear_rgb(image: &DynamicImage, x: u32, y: u32, is_raw: bool) -> Option<[f32; 3]> { + let rgb = match image { + DynamicImage::ImageRgb32F(buf) => { + let p = buf.get_pixel(x, y); + [p[0], p[1], p[2]] + } + DynamicImage::ImageRgba32F(buf) => { + let p = buf.get_pixel(x, y); + [p[0], p[1], p[2]] + } + _ => { + let p = image.crop_imm(x, y, 1, 1).to_rgb32f(); + let p = p.get_pixel(0, 0); + [p[0], p[1], p[2]] + } + }; + + if rgb.iter().any(|c| !c.is_finite()) { + return None; + } + + let rgb = rgb.map(|c| c.max(0.0)); + if is_raw { + Some(rgb) + } else { + Some(rgb.map(srgb_channel_to_linear)) + } +} + +fn point_in_convex_quad(px: f64, py: f64, quad: &[(f64, f64)]) -> bool { + let mut has_pos = false; + let mut has_neg = false; + for i in 0..quad.len() { + let (ax, ay) = quad[i]; + let (bx, by) = quad[(i + 1) % quad.len()]; + let cross = (bx - ax) * (py - ay) - (by - ay) * (px - ax); + if cross > 0.0 { + has_pos = true; + } else if cross < 0.0 { + has_neg = true; + } + if has_pos && has_neg { + return false; + } + } + true +} + +fn compute_white_balance_sample( + image: &DynamicImage, + is_raw: bool, + as_shot: WhiteBalance, + corners: &[UvPoint], +) -> Result { + if corners.len() < 3 { + return Err("At least three corners are required".to_string()); + } + + let (width, height) = image.dimensions(); + if width == 0 || height == 0 { + return Err("Image is empty".to_string()); + } + + let quad: Vec<(f64, f64)> = corners + .iter() + .map(|c| (c.x * width as f64, c.y * height as f64)) + .collect(); + + if quad.iter().any(|(x, y)| !x.is_finite() || !y.is_finite()) { + return Err("Invalid sample corners".to_string()); + } + + let min_x = quad.iter().map(|p| p.0).fold(f64::INFINITY, f64::min); + let max_x = quad.iter().map(|p| p.0).fold(f64::NEG_INFINITY, f64::max); + let min_y = quad.iter().map(|p| p.1).fold(f64::INFINITY, f64::min); + let max_y = quad.iter().map(|p| p.1).fold(f64::NEG_INFINITY, f64::max); + + let start_x = min_x.floor().clamp(0.0, width as f64) as u32; + let end_x = max_x.ceil().clamp(0.0, width as f64) as u32; + let start_y = min_y.floor().clamp(0.0, height as f64) as u32; + let end_y = max_y.ceil().clamp(0.0, height as f64) as u32; + + let area = (end_x.saturating_sub(start_x) as f64) * (end_y.saturating_sub(start_y) as f64); + let stride = ((area / MAX_WB_SAMPLES).sqrt().ceil() as usize).max(1); + + let mut sum = [0.0f64; 3]; + let mut count: u32 = 0; + + for y in (start_y..end_y).step_by(stride) { + for x in (start_x..end_x).step_by(stride) { + if !point_in_convex_quad(x as f64 + 0.5, y as f64 + 0.5, &quad) { + continue; + } + if let Some(rgb) = read_linear_rgb(image, x, y, is_raw) { + sum[0] += rgb[0] as f64; + sum[1] += rgb[1] as f64; + sum[2] += rgb[2] as f64; + count += 1; + } + } + } + + if count == 0 { + let cx = quad.iter().map(|p| p.0).sum::() / quad.len() as f64; + let cy = quad.iter().map(|p| p.1).sum::() / quad.len() as f64; + let x = (cx.floor().max(0.0) as u32).min(width - 1); + let y = (cy.floor().max(0.0) as u32).min(height - 1); + let rgb = read_linear_rgb(image, x, y, is_raw).ok_or("Sampled pixel is invalid")?; + sum = [rgb[0] as f64, rgb[1] as f64, rgb[2] as f64]; + count = 1; + } + + let r = (sum[0] / count as f64) as f32; + let g = (sum[1] / count as f64) as f32; + let b = (sum[2] / count as f64) as f32; + let picked = white_balance::pick_white_balance([r as f64, g as f64, b as f64], as_shot) + .unwrap_or(as_shot); + + Ok(WhiteBalanceSample { + r, + g, + b, + temperature: picked.temperature as f32, + tint: picked.tint as f32, + count, + }) +} + +#[tauri::command] +pub async fn sample_white_balance( + corners: Vec, + state: tauri::State<'_, AppState>, +) -> Result { + let loaded_image = state + .original_image + .lock() + .unwrap() + .clone() + .ok_or("No original image loaded")?; + + tokio::task::spawn_blocking(move || { + compute_white_balance_sample( + &loaded_image.image, + loaded_image.is_raw, + loaded_image.as_shot_white_balance, + &corners, + ) + }) + .await + .map_err(|e| format!("Task execution failed: {}", e))? +} diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 71a80f9cc8..60cdf71e92 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2263,7 +2263,6 @@ pub fn run() { focus_stacking::stitch_focus_stack, focus_stacking::save_focus_stack, image_loader::load_image, - white_balance::pick_white_balance, image_loader::is_image_cached, panorama_stitching::stitch_panorama, panorama_stitching::save_panorama, @@ -2271,6 +2270,7 @@ pub fn run() { export_processing::cancel_export, export_processing::estimate_export_sizes, image_processing::calculate_auto_adjustments, + image_processing::sample_white_balance, mask_generation::generate_mask_overlay, file_management::update_exif_fields, file_management::get_supported_file_types, diff --git a/src-tauri/src/white_balance.rs b/src-tauri/src/white_balance.rs index b96b7e2e4d..b519defedc 100644 --- a/src-tauri/src/white_balance.rs +++ b/src-tauri/src/white_balance.rs @@ -247,7 +247,6 @@ pub fn from_adjustments(adjustments: &Value, as_shot: WhiteBalance) -> WhiteBala ) } -#[tauri::command] pub fn pick_white_balance(sample: [f64; 3], current: WhiteBalance) -> Option { let rgb_to_lms = rgb_to_lms().as_dmat3(); let sample_lms = rgb_to_lms * DVec3::from_array(sample); diff --git a/src/components/panel/editor/ImageCanvas.tsx b/src/components/panel/editor/ImageCanvas.tsx index 2612025560..264a66c6b3 100644 --- a/src/components/panel/editor/ImageCanvas.tsx +++ b/src/components/panel/editor/ImageCanvas.tsx @@ -18,9 +18,7 @@ import { calculateStraightenAngle } from '../../../utils/cropUtils'; import { toast } from 'react-toastify'; import { getWhiteBalanceMode, - resolveWhiteBalance, toRelativeWhiteBalance, - WhiteBalance, WhiteBalanceMode, withKelvinWhiteBalance, withRelativeWhiteBalance, @@ -112,6 +110,64 @@ interface MaskOverlayProps { const IDENTITY_3X3 = [1, 0, 0, 0, 1, 0, 0, 0, 1]; +const WB_SAMPLE_SCREEN_SIZE = 16; +const WB_DRAG_THRESHOLD = 4; +const WB_SWATCH_OFFSET = 18; + +interface WbSample { + r: number; + g: number; + b: number; + temperature: number; + tint: number; + count: number; +} + +interface WbDrag { + start: Coord; + end: Coord; + isBox: boolean; +} + +const linearToSrgb8 = (value: number) => { + const c = Math.max(0, Math.min(1, value)); + const encoded = c <= 0.0031308 ? c * 12.92 : 1.055 * Math.pow(c, 1 / 2.4) - 0.055; + return Math.round(encoded * 255); +}; + +interface WbSampleOutlineProps { + x: number; + y: number; + width: number; + height: number; + zoomScale: number; + dashed?: boolean; +} + +const WbSampleOutline = ({ x, y, width, height, zoomScale, dashed = false }: WbSampleOutlineProps) => ( + <> + + + +); + function multiply3x3(a: number[], b: number[]): number[] { if (!a || !b) return IDENTITY_3X3; const out = [0, 0, 0, 0, 0, 0, 0, 0, 0]; @@ -1429,6 +1485,16 @@ const ImageCanvas = memo( const activeStrokeIndex = useRef(null); const [cursorPreview, setCursorPreview] = useState({ x: 0, y: 0, visible: false }); + const [wbHover, setWbHover] = useState({ x: 0, y: 0, visible: false }); + const wbDragRef = useRef(null); + const [wbBox, setWbBox] = useState(null); + const [wbSample, setWbSample] = useState(null); + const wbSampleStateRef = useRef({ + inFlight: false, + pending: null as Coord[] | null, + generation: 0, + session: 0, + }); const [straightenLine, setStraightenLine] = useState(null); const isStraightening = useRef(false); @@ -2051,105 +2117,134 @@ const ImageCanvas = memo( [invH, uncroppedImageRenderSize, selectedImage, adjustments, liveRotation], ); - const handleWbClick = useCallback( - (e: any) => { - const sampleUrl = selectedImage?.thumbnailUrl || finalPreviewUrl; - const asShot = selectedImage?.asShotWhiteBalance; - if (!isWbPickerActive || !sampleUrl || !onWbPicked || !asShot) return; - - const stage = e.target.getStage(); - const pointerPos = getCanvasPointer(stage); - if (!pointerPos) return; + const wbSquareStage = WB_SAMPLE_SCREEN_SIZE / effectiveZoomScale; - const x = pointerPos.x / imageRenderSize.scale; - const y = pointerPos.y / imageRenderSize.scale; - - const imgLogicalWidth = imageRenderSize.width / imageRenderSize.scale; - const imgLogicalHeight = imageRenderSize.height / imageRenderSize.scale; - - if (x < 0 || x > imgLogicalWidth || y < 0 || y > imgLogicalHeight) return; + const isInsideImage = useCallback( + (pos: Coord | null | undefined): pos is Coord => + !!pos && pos.x >= 0 && pos.y >= 0 && pos.x <= imageRenderSize.width && pos.y <= imageRenderSize.height, + [imageRenderSize.width, imageRenderSize.height], + ); - const img = new Image(); - img.crossOrigin = 'Anonymous'; - img.src = sampleUrl; + const mapCanvasPointToUv = useCallback( + (p: Coord): Coord => { + if ( + !uncroppedImageRenderSize?.width || + !uncroppedImageRenderSize?.height || + !effectiveImageDimensions.width || + !effectiveImageDimensions.height + ) { + return { x: 0, y: 0 }; + } + const scale = imageRenderSize.scale || 1; + const stageX = ((p.x / scale + cropX) / effectiveImageDimensions.width) * uncroppedImageRenderSize.width; + const stageY = ((p.y / scale + cropY) / effectiveImageDimensions.height) * uncroppedImageRenderSize.height; + return mapScreenToUv(stageX, stageY); + }, + [uncroppedImageRenderSize, effectiveImageDimensions, imageRenderSize.scale, cropX, cropY, mapScreenToUv], + ); - img.onload = async () => { - const radius = 5; - const side = radius * 2 + 1; + const getWbCorners = useCallback( + (x0: number, y0: number, x1: number, y1: number): Coord[] => { + const minX = Math.max(0, Math.min(x0, x1)); + const maxX = Math.min(imageRenderSize.width, Math.max(x0, x1)); + const minY = Math.max(0, Math.min(y0, y1)); + const maxY = Math.min(imageRenderSize.height, Math.max(y0, y1)); + return [ + { x: minX, y: minY }, + { x: maxX, y: minY }, + { x: maxX, y: maxY }, + { x: minX, y: maxY }, + ].map(mapCanvasPointToUv); + }, + [imageRenderSize.width, imageRenderSize.height, mapCanvasPointToUv], + ); - const canvas = document.createElement('canvas'); - canvas.width = side; - canvas.height = side; - const ctx = canvas.getContext('2d', { willReadFrequently: true }); - if (!ctx) return; + const getWbSquareCorners = useCallback( + (center: Coord) => { + const half = wbSquareStage / 2; + return getWbCorners(center.x - half, center.y - half, center.x + half, center.y + half); + }, + [wbSquareStage, getWbCorners], + ); - const scaleX = img.width / imgLogicalWidth; - const scaleY = img.height / imgLogicalHeight; - const srcX = Math.floor(x * scaleX); - const srcY = Math.floor(y * scaleY); + const requestWbSample = useCallback(async (corners: Coord[]) => { + const state = wbSampleStateRef.current; + if (state.inFlight) { + state.pending = corners; + return; + } - const startX = Math.max(0, srcX - radius); - const startY = Math.max(0, srcY - radius); - const endX = Math.min(img.width, srcX + radius + 1); - const endY = Math.min(img.height, srcY + radius + 1); - const sw = endX - startX; - const sh = endY - startY; + state.inFlight = true; + state.pending = null; + const generation = state.generation; - if (sw <= 0 || sh <= 0) return; + try { + const sample = await invoke(Invokes.SampleWhiteBalance, { corners }); + if (state.generation === generation) { + setWbSample(sample); + } + } catch (err) { + console.error('Failed to sample white balance:', err); + } finally { + state.inFlight = false; + if (state.pending) { + const next = state.pending; + state.pending = null; + requestWbSample(next); + } + } + }, []); - ctx.drawImage(img, startX, startY, sw, sh, 0, 0, sw, sh); + const resetWbSample = useCallback(() => { + const state = wbSampleStateRef.current; + state.generation += 1; + state.pending = null; + setWbSample(null); + setWbHover((p: CursorPreview) => (p.visible ? { ...p, visible: false } : p)); + }, []); - const imageData = ctx.getImageData(0, 0, sw, sh); - const data = imageData.data; + const isKelvinWhiteBalance = getWhiteBalanceMode(appSettings) === WhiteBalanceMode.Kelvin; + const asShotWhiteBalance = selectedImage?.asShotWhiteBalance; - let rTotal = 0, - gTotal = 0, - bTotal = 0; - let count = 0; + const applyWbPick = useCallback( + async (corners: Coord[]) => { + const state = wbSampleStateRef.current; + state.generation += 1; + state.pending = null; + const { generation, session } = state; - for (let i = 0; i < data.length; i += 4) { - rTotal += data[i]; - gTotal += data[i + 1]; - bTotal += data[i + 2]; - count++; + try { + const sample = await invoke(Invokes.SampleWhiteBalance, { corners }); + if (state.session !== session) return; + if (state.generation === generation) { + setWbSample(sample); } - - if (count === 0) return; - - const avgR = rTotal / count; - const avgG = gTotal / count; - const avgB = bTotal / count; - - const sample = [avgR, avgG, avgB].map((c) => Math.pow(c / 255.0, 2.2)); - const picked: WhiteBalance | null = await invoke(Invokes.PickWhiteBalance, { - sample, - current: resolveWhiteBalance(asShot, adjustments), - }); - if (!picked) return; - - setAdjustments((prev: Adjustments) => - getWhiteBalanceMode(appSettings) === WhiteBalanceMode.Kelvin - ? withKelvinWhiteBalance(prev, picked) - : withRelativeWhiteBalance(prev, toRelativeWhiteBalance(asShot, picked)), - ); - - onWbPicked(); - }; + if (asShotWhiteBalance) { + const picked = { temperature: sample.temperature, tint: sample.tint }; + setAdjustments((prev: Adjustments) => + isKelvinWhiteBalance + ? withKelvinWhiteBalance(prev, picked) + : withRelativeWhiteBalance(prev, toRelativeWhiteBalance(asShotWhiteBalance, picked)), + ); + } + onWbPicked?.(); + } catch (err) { + console.error('Failed to pick white balance:', err); + } }, - [ - isWbPickerActive, - selectedImage?.thumbnailUrl, - selectedImage?.asShotWhiteBalance, - adjustments, - appSettings, - finalPreviewUrl, - imageRenderSize, - onWbPicked, - setAdjustments, - getCanvasPointer, - ], + [setAdjustments, onWbPicked, asShotWhiteBalance, isKelvinWhiteBalance], ); + useEffect(() => { + wbSampleStateRef.current.session += 1; + if (wbDragRef.current) { + wbDragRef.current = null; + isDrawing.current = false; + } + setWbBox(null); + resetWbSample(); + }, [isWbPickerActive, selectedImage.path, resetWbSample]); + const handleStart = useCallback( (e: any) => { if (e.evt && typeof e.evt.button === 'number' && e.evt.button !== 0) { @@ -2171,7 +2266,12 @@ const ImageCanvas = memo( } if (isWbPickerActive) { - handleWbClick(e); + const stage = e.target.getStage(); + const pos = getCanvasPointer(stage); + if (!isInsideImage(pos)) return; + wbDragRef.current = { start: pos, end: pos, isBox: false }; + drawingStageRef.current = stage; + isDrawing.current = true; return; } @@ -2370,7 +2470,7 @@ const ImageCanvas = memo( isCropping, mapScreenToUv, isWbPickerActive, - handleWbClick, + isInsideImage, isInitialDrawing, isBrushActive, isCloneOrHealActive, @@ -2412,10 +2512,6 @@ const ImageCanvas = memo( return; } - if (isWbPickerActive) { - return; - } - let pos; if (e && typeof e.target?.getStage === 'function') { const stage = e.target.getStage(); @@ -2428,6 +2524,34 @@ const ImageCanvas = memo( } } + if (isWbPickerActive) { + const drag = wbDragRef.current; + if (drag && isDrawing.current) { + if (!pos) return; + const distance = Math.hypot(pos.x - drag.start.x, pos.y - drag.start.y); + const updated = { + start: drag.start, + end: pos, + isBox: drag.isBox || distance >= WB_DRAG_THRESHOLD / effectiveZoomScale, + }; + wbDragRef.current = updated; + if (updated.isBox) { + setWbBox(updated); + requestWbSample(getWbCorners(updated.start.x, updated.start.y, pos.x, pos.y)); + } + if (e.evt && e.evt.cancelable) e.evt.preventDefault(); + return; + } + + if (isInsideImage(pos)) { + setWbHover({ x: pos.x, y: pos.y, visible: true }); + requestWbSample(getWbSquareCorners(pos)); + } else { + resetWbSample(); + } + return; + } + if (isToolActive) { if (pos) { setCursorPreview({ x: pos.x, y: pos.y, visible: true }); @@ -2640,6 +2764,12 @@ const ImageCanvas = memo( brushImageSpaceSize, baseTool, getCanvasPointer, + effectiveZoomScale, + isInsideImage, + requestWbSample, + resetWbSample, + getWbCorners, + getWbSquareCorners, ], ); @@ -2648,6 +2778,18 @@ const ImageCanvas = memo( return; } + if (isWbPickerActive && wbDragRef.current) { + isDrawing.current = false; + const { start, end, isBox } = wbDragRef.current; + wbDragRef.current = null; + setWbBox(null); + if (isInsideImage(end)) { + setWbHover({ x: end.x, y: end.y, visible: true }); + } + applyWbPick(isBox ? getWbCorners(start.x, start.y, end.x, end.y) : getWbSquareCorners(start)); + return; + } + setIsMaskInteractionActive(false); if (isGuidedPerspectiveActive && isCropping && draftGuideLine) { @@ -2871,6 +3013,11 @@ const ImageCanvas = memo( brushImageSpaceSize, brushStageSize, baseTool, + isWbPickerActive, + isInsideImage, + applyWbPick, + getWbCorners, + getWbSquareCorners, ]); const handleMouseEnter = useCallback(() => { @@ -2881,10 +3028,13 @@ const ImageCanvas = memo( const handleMouseLeave = useCallback(() => { setCursorPreview((p: CursorPreview) => ({ ...p, visible: false })); - }, []); + if (!wbDragRef.current) { + resetWbSample(); + } + }, [resetWbSample]); useEffect(() => { - if (!isToolActive) return; + if (!isToolActive && !isWbPickerActive) return; function onGlobalMove(e: MouseEvent | TouchEvent) { if (!isDrawing.current) return; @@ -2906,7 +3056,7 @@ const ImageCanvas = memo( window.removeEventListener('touchmove', onGlobalMove); window.removeEventListener('touchcancel', onGlobalUp); }; - }, [isToolActive, handleMove, handleUp]); + }, [isToolActive, isWbPickerActive, handleMove, handleUp]); const handleStraightenMouseDown = (e: any) => { if (e.evt.button !== 0 && !e.evt.touches) { @@ -2958,6 +3108,28 @@ const ImageCanvas = memo( const cropPreviewUrl = uncroppedAdjustedPreviewUrl || selectedImage.thumbnailUrl; const isShowingOriginal = showOriginal; + const wbSwatchAnchor = + isWbPickerActive && wbSample && !isShowingOriginal + ? wbBox + ? wbBox.end + : wbHover.visible + ? wbHover + : null + : null; + const wbSwatchNorm = wbSample ? Math.max(1, wbSample.r, wbSample.g, wbSample.b) : 1; + const wbSwatchRgb = wbSample + ? [wbSample.r, wbSample.g, wbSample.b].map((v) => linearToSrgb8(v / wbSwatchNorm)) + : [0, 0, 0]; + const wbSwatchFlipX = !!wbSwatchAnchor && wbSwatchAnchor.x > imageRenderSize.width * 0.75; + const wbSwatchFlipY = !!wbSwatchAnchor && wbSwatchAnchor.y > imageRenderSize.height * 0.75; + const wbSwatchOffset = WB_SWATCH_OFFSET / effectiveZoomScale; + const wbSwatchWhiteBalance = + wbSample && asShotWhiteBalance + ? isKelvinWhiteBalance + ? wbSample + : toRelativeWhiteBalance(asShotWhiteBalance, wbSample) + : null; + const currentTarget = finalPreviewUrl || selectedImage.thumbnailUrl; const baseIsReady = displayState.base === currentTarget && !displayState.fade; @@ -3339,6 +3511,25 @@ const ImageCanvas = memo( listening={false} /> )} + {isWbPickerActive && wbBox && ( + + )} + {isWbPickerActive && wbHover.visible && !wbBox && ( + + )} {isBrushActive && cursorPreview.visible && (!isCloneOrHealActive || @@ -3366,6 +3557,41 @@ const ImageCanvas = memo( )} + + {wbSwatchAnchor && wbSample && ( +
+
+
+ + R {wbSwatchRgb[0]} G {wbSwatchRgb[1]} B {wbSwatchRgb[2]} + + {wbSwatchWhiteBalance && ( + + {t('adjustments.color.temperature')} + {`${Math.round(wbSwatchWhiteBalance.temperature)}${isKelvinWhiteBalance ? 'K' : ''}`} + {t('adjustments.color.tint')} + {Math.round(wbSwatchWhiteBalance.tint)} + + )} +
+
+ )}