feat(color): area sampling and colour preview for the white balance picker

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<WhiteBalanceSample, String> {
+    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::<f64>() / quad.len() as f64;
+        let cy = quad.iter().map(|p| p.1).sum::<f64>() / 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<UvPoint>,
+    state: tauri::State<'_, AppState>,
+) -> Result<WhiteBalanceSample, String> {
+    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))?
+}

src-tauri/src/lib.rs

@@ -2263,14 +2263,14 @@ 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,
             export_processing::export_images,
             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,

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<WhiteBalance> {
     let rgb_to_lms = rgb_to_lms().as_dmat3();
     let sample_lms = rgb_to_lms * DVec3::from_array(sample);