diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index b6f4c26..4e8ce44 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -455,31 +455,101 @@ async fn detect_watermark(path: String) -> Result { let (width, height) = img.dimensions(); let gray = img.to_luma8(); - // Heuristic: - // 1. Scan Top 20% and Bottom 20% - // 2. Look for high brightness pixels (> 230) - // 3. Grid based clustering + // "Stroke Detection" Algorithm + // Distinguishes "Text" (Thin White Strokes) from "Solid White Areas" (Walls, Sky) + // Logic: A white text pixel must be "sandwiched" by dark pixels within a short distance. let cell_size = 10; let grid_w = (width + cell_size - 1) / cell_size; let grid_h = (height + cell_size - 1) / cell_size; let mut grid = vec![false; (grid_w * grid_h) as usize]; - let top_limit = (height as f64 * 0.2) as u32; - let bottom_start = (height as f64 * 0.8) as u32; + // Focus Areas: Top 15%, Bottom 25% + let top_limit = (height as f64 * 0.15) as u32; + let bottom_start = (height as f64 * 0.75) as u32; + + let max_stroke_width = 15; // Max pixels for a text stroke thickness + let contrast_threshold = 40; // How much darker the background must be + let brightness_threshold = 200; // Text must be at least this white - for y in 0..height { - // Skip middle section - if y > top_limit && y < bottom_start { - continue; - } + for y in 1..height-1 { + if y > top_limit && y < bottom_start { continue; } - for x in 0..width { - let p = gray.get_pixel(x, y); - if p[0] > 230 { // High brightness threshold - // Check local contrast/edges? - // For now, simple brightness is a good proxy for "white text" - // Mark grid cell + for x in 1..width-1 { + let p = gray.get_pixel(x, y)[0]; + + // 1. Must be Bright + if p < brightness_threshold { continue; } + + // 2. Stroke Check + // We check for "Vertical Stroke" (Dark - Bright - Dark vertically) + // OR "Horizontal Stroke" (Dark - Bright - Dark horizontally) + + let mut is_stroke = false; + + // Check Horizontal Stroke (Vertical boundaries? No, Vertical Stroke has Left/Right boundaries? + // Terminology: "Vertical Stroke" is like 'I'. It has Left/Right boundaries. + // "Horizontal Stroke" is like '-', It has Up/Down boundaries. + + // Let's check Left/Right boundaries (Vertical Stroke) + let mut left_bound = false; + let mut right_bound = false; + + // Search Left + for k in 1..=max_stroke_width { + if x < k { break; } + let neighbor = gray.get_pixel(x - k, y)[0]; + if p > neighbor && (p - neighbor) > contrast_threshold { + left_bound = true; + break; + } + } + // Search Right + if left_bound { + for k in 1..=max_stroke_width { + if x + k >= width { break; } + let neighbor = gray.get_pixel(x + k, y)[0]; + if p > neighbor && (p - neighbor) > contrast_threshold { + right_bound = true; + break; + } + } + } + + if left_bound && right_bound { + is_stroke = true; + } else { + // Check Up/Down boundaries (Horizontal Stroke) + let mut up_bound = false; + let mut down_bound = false; + + // Search Up + for k in 1..=max_stroke_width { + if y < k { break; } + let neighbor = gray.get_pixel(x, y - k)[0]; + if p > neighbor && (p - neighbor) > contrast_threshold { + up_bound = true; + break; + } + } + // Search Down + if up_bound { + for k in 1..=max_stroke_width { + if y + k >= height { break; } + let neighbor = gray.get_pixel(x, y + k)[0]; + if p > neighbor && (p - neighbor) > contrast_threshold { + down_bound = true; + break; + } + } + } + + if up_bound && down_bound { + is_stroke = true; + } + } + + if is_stroke { let gx = x / cell_size; let gy = y / cell_size; grid[(gy * grid_w + gx) as usize] = true; @@ -496,9 +566,6 @@ async fn detect_watermark(path: String) -> Result { let idx = (gy * grid_w + gx) as usize; if grid[idx] && !visited[idx] { // Start a new component - // Simple Flood Fill or just greedy expansion - // Let's do a simple greedy expansion for rectangles - let mut min_gx = gx; let mut max_gx = gx; let mut min_gy = gy; @@ -550,11 +617,178 @@ async fn detect_watermark(path: String) -> Result { Ok(DetectionResult { rects }) } +#[derive(serde::Deserialize)] +struct StrokePoint { + x: f64, + y: f64, +} + +#[derive(serde::Deserialize)] +struct StrokeRect { + x: f64, + y: f64, + w: f64, + h: f64, +} + +#[derive(serde::Deserialize)] +#[serde(tag = "type")] +enum MaskStroke { + #[serde(rename = "path")] + Path { + points: Vec, + width: f64, + }, + #[serde(rename = "rect")] + Rect { + rect: StrokeRect, + }, +} + +#[tauri::command] +async fn run_inpainting(path: String, strokes: Vec) -> Result { + let img = image::open(&path).map_err(|e| e.to_string())?.to_rgba8(); + let (width, height) = img.dimensions(); + + // 1. Create Mask + let mut mask = vec![false; (width * height) as usize]; + + for stroke in strokes { + match stroke { + MaskStroke::Rect { rect } => { + let x1 = (rect.x * width as f64) as i32; + let y1 = (rect.y * height as f64) as i32; + let w = (rect.w * width as f64) as i32; + let h = (rect.h * height as f64) as i32; + + for y in y1..(y1 + h) { + for x in x1..(x1 + w) { + if x >= 0 && x < width as i32 && y >= 0 && y < height as i32 { + mask[(y as u32 * width + x as u32) as usize] = true; + } + } + } + }, + MaskStroke::Path { points, width: stroke_w_pct } => { + if points.is_empty() { continue; } + let r = (stroke_w_pct * width as f64 / 2.0).ceil() as i32; + let r2 = r * r; + + for i in 0..points.len() - 1 { + let p1 = &points[i]; + let p2 = &points[i+1]; + + let x1 = p1.x * width as f64; + let y1 = p1.y * height as f64; + let x2 = p2.x * width as f64; + let y2 = p2.y * height as f64; + + // Simple line interpolation + let dist = ((x2-x1).powi(2) + (y2-y1).powi(2)).sqrt(); + let steps = dist.max(1.0) as i32; + + for s in 0..=steps { + let t = s as f64 / steps as f64; + let cx = (x1 + (x2 - x1) * t) as i32; + let cy = (y1 + (y2 - y1) * t) as i32; + + // Draw circle at cx, cy + for dy in -r..=r { + for dx in -r..=r { + if dx*dx + dy*dy <= r2 { + let nx = cx + dx; + let ny = cy + dy; + if nx >= 0 && nx < width as i32 && ny >= 0 && ny < height as i32 { + mask[(ny as u32 * width + nx as u32) as usize] = true; + } + } + } + } + } + } + } + } + } + + // 2. Diffusion Inpainting (Simple) + // Iteratively replace masked pixels with average of non-masked neighbors + // To make it converge, we update 'mask' as we go (treating filled pixels as valid source) + // But standard diffusion uses double buffering. + // For "Removing Text", simple inward filling works well. + + let iterations = 30; + let mut current_img = img.clone(); + let mut next_img = img.clone(); + + // Convert mask to a distance map-like state? + // Or just simple neighbor average. + + for _ in 0..iterations { + let mut changed = false; + for y in 0..height { + for x in 0..width { + let idx = (y * width + x) as usize; + if mask[idx] { + // It's a hole. Find valid neighbors. + // Valid = Not in ORIGINAL mask (so we pull from original image) + // OR processed in previous iteration? + // Simple logic: Pull from 'current_img'. + + let mut sum_r = 0u32; + let mut sum_g = 0u32; + let mut sum_b = 0u32; + let mut count = 0; + + // Check 4 neighbors + let neighbors = [ + (x.wrapping_sub(1), y), (x + 1, y), + (x, y.wrapping_sub(1)), (x, y + 1) + ]; + + for (nx, ny) in neighbors { + if nx < width && ny < height { + // Weighted check: If neighbor is ALSO masked, it contributes less? + // Or just take everything. + let pixel = current_img.get_pixel(nx, ny); + sum_r += pixel[0] as u32; + sum_g += pixel[1] as u32; + sum_b += pixel[2] as u32; + count += 1; + } + } + + if count > 0 { + let avg = image::Rgba([ + (sum_r / count) as u8, + (sum_g / count) as u8, + (sum_b / count) as u8, + 255 + ]); + next_img.put_pixel(x, y, avg); + changed = true; + } + } + } + } + current_img = next_img.clone(); + if !changed { break; } + } + + // Save to temp + let cache_dir = get_cache_dir(); + let file_name = format!("inpainted_{}.png", std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_millis()); + let out_path = cache_dir.join(file_name); + + current_img.save(&out_path).map_err(|e| e.to_string())?; + + Ok(out_path.to_string_lossy().to_string()) +} + #[cfg_attr(mobile, tauri::mobile_entry_point)] pub fn run() { tauri::Builder::default() .plugin(tauri_plugin_dialog::init()) - .invoke_handler(tauri::generate_handler![scan_dir, get_zca_suggestion, export_batch, detect_watermark, layout_watermark]) + .invoke_handler(tauri::generate_handler![scan_dir, get_zca_suggestion, export_batch, detect_watermark, layout_watermark, run_inpainting]) .run(tauri::generate_context!()) .expect("error while running tauri application"); } diff --git a/src/components/SettingsPanel.vue b/src/components/SettingsPanel.vue index be1f9d4..9735b90 100644 --- a/src/components/SettingsPanel.vue +++ b/src/components/SettingsPanel.vue @@ -183,9 +183,8 @@ const applyAll = () => {
@@ -217,8 +216,17 @@ const applyAll = () => {
- AI Inpainting is not yet connected. Masking preview only. + AI Inpainting (Diffusion) connected.
+ + diff --git a/src/stores/gallery.ts b/src/stores/gallery.ts index ffc0383..1ad2994 100644 --- a/src/stores/gallery.ts +++ b/src/stores/gallery.ts @@ -95,26 +95,31 @@ export const useGalleryStore = defineStore("gallery", () => { } } - async function detectWatermark(index: number) { - const img = images.value[index]; - if (!img) return; - - try { - // @ts-ignore - const result = await invoke<{rects: {x: number, y: number, width: number, height: number}[]}>("detect_watermark", { path: img.path }); - - if (result.rects && result.rects.length > 0) { - if (!img.maskStrokes) img.maskStrokes = []; - - result.rects.forEach(r => { - img.maskStrokes!.push({ - type: 'rect', - rect: { x: r.x, y: r.y, w: r.width, h: r.height } - }); - }); - } - } catch (e) { - console.error("Detection failed", e); + async function detectAllWatermarks() { + if (images.value.length === 0) return; + + const batchSize = 5; + for (let i = 0; i < images.value.length; i += batchSize) { + const batch = images.value.slice(i, i + batchSize).map(async (img) => { + try { + // @ts-ignore + const result = await invoke<{rects: {x: number, y: number, width: number, height: number}[]}>("detect_watermark", { path: img.path }); + + if (result.rects && result.rects.length > 0) { + if (!img.maskStrokes) img.maskStrokes = []; + + result.rects.forEach(r => { + img.maskStrokes!.push({ + type: 'rect', + rect: { x: r.x, y: r.y, w: r.width, h: r.height } + }); + }); + } + } catch (e) { + console.error(`Detection failed for ${img.name}`, e); + } + }); + await Promise.all(batch); } } @@ -127,8 +132,7 @@ export const useGalleryStore = defineStore("gallery", () => { // Process in batches to avoid overwhelming the backend const batchSize = 5; for (let i = 0; i < images.value.length; i += batchSize) { - const batch = images.value.slice(i, i + batchSize).map(async (img, batchIdx) => { - const globalIdx = i + batchIdx; + const batch = images.value.slice(i, i + batchSize).map(async (img) => { try { const result = await invoke<{x: number, y: number, scale: number}>("layout_watermark", { path: img.path, @@ -136,8 +140,14 @@ export const useGalleryStore = defineStore("gallery", () => { baseScale: baseScale }); - setImageManualPosition(globalIdx, result.x, result.y); - setImageSetting(globalIdx, 'scale', result.scale); + // Find index again just in case sort changed? No, we rely on reference or index. + // Img is reference. But setter needs index. + // Let's use image reference finding. + const idx = images.value.indexOf(img); + if (idx >= 0) { + setImageManualPosition(idx, result.x, result.y); + setImageSetting(idx, 'scale', result.scale); + } } catch (e) { console.error(`Layout failed for ${img.name}`, e); } @@ -146,6 +156,33 @@ export const useGalleryStore = defineStore("gallery", () => { } } + async function processInpainting(index: number) { + const img = images.value[index]; + if (!img || !img.maskStrokes || img.maskStrokes.length === 0) return; + + try { + const newPath = await invoke("run_inpainting", { + path: img.path, + strokes: img.maskStrokes + }); + + // Update the image path to point to the processed version + // NOTE: This updates the "Source" of the image in the store. + // In a real app, maybe we keep original? But for "Wizard", modifying is the goal. + img.path = newPath; + // Clear mask after success + img.maskStrokes = []; + + // Force UI refresh (thumbnail) - might be needed + // Generate new thumb? + // Since path changed, converting src should trigger reload. + + } catch (e) { + console.error("Inpainting failed", e); + alert("Inpainting failed: " + e); + } + } + // Applies the settings from the CURRENT image (or global if not overridden) to ALL images // Strategy: Update Global Settings to match current view, and clear individual overrides so everyone follows global. function applySettingsToAll() { @@ -199,7 +236,8 @@ export const useGalleryStore = defineStore("gallery", () => { applySettingsToAll, addMaskStroke, clearMask, - detectWatermark, - recalcAllWatermarks + detectAllWatermarks, + recalcAllWatermarks, + processInpainting }; });