remove watermarks but sucks
This commit is contained in:
Julian Freeman
@@ -455,31 +455,101 @@ async fn detect_watermark(path: String) -> Result<DetectionResult, String> {
|
||||
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<DetectionResult, String> {
|
||||
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<DetectionResult, String> {
|
||||
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<StrokePoint>,
|
||||
width: f64,
|
||||
},
|
||||
#[serde(rename = "rect")]
|
||||
Rect {
|
||||
rect: StrokeRect,
|
||||
},
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
async fn run_inpainting(path: String, strokes: Vec<MaskStroke>) -> Result<String, String> {
|
||||
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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user