feat: 优化去水印算法和背景水印图生成逻辑

src-tauri/Cargo.toml

@@ -38,6 +38,7 @@ dirs = { version = "5.0" }
 path-slash = { version = "0.2.1" }
 tokio = { version = "1", features = ["full"] }
 bytes = { version = "1.7.1" }
+parking_lot = { version = "0.12.3", features = ["send_guard"] }
 
 [profile.release]
 strip = true      # Automatically strip symbols from the binary.

src-tauri/src/commands/generate_background.rs

@@ -1,15 +1,14 @@
 use std::path::PathBuf;
-use std::sync::Mutex;
 
 use anyhow::Context;
 use image::RgbImage;
+use parking_lot::Mutex;
 use path_slash::PathBufExt;
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use tauri::AppHandle;
 use walkdir::WalkDir;
 
 use crate::errors::CommandResult;
-use crate::extensions::IgnoreLockPoison;
 use crate::types::{CommandResponse, RectData};
 use crate::utils;
 
@@ -27,81 +26,108 @@ pub fn generate_background(
     height: u32,
 ) -> CommandResult<CommandResponse<()>> {
     let output_dir = utils::get_background_dir_abs_path(&app, manga_dir, width, height)?;
+    // TODO: 给RectData实现Default trait，以替换下面的代码
     let default_rect_data = RectData {
         left: (width as f32 * 0.835) as u32,
         top: (height as f32 * 0.946) as u32,
         right: (width as f32 * 0.994) as u32,
         bottom: (height as f32 * 0.994) as u32,
     };
     let rect_data = rect_data.unwrap_or(default_rect_data);
+    // TODO: 删除下面的代码
     // let res = watermark::generate_background(manga_dir, &rect_data, &output_dir, width, height)?;
     // Ok(res)
 
     // 保证输出目录存在
+    // TODO: 将各种.display()换成 {:?}
     std::fs::create_dir_all(&output_dir)
         .context(format!("创建目录 {} 失败", output_dir.display()))?;
     // 收集尺寸符合width和height的图片的路径
     let image_paths = create_image_paths(manga_dir, width, height);
-    // 用于记录是否找到了黑色背景和白色背景的水印图片
-    let black_status: Mutex<Option<()>> = Mutex::new(None);
-    let white_status: Mutex<Option<()>> = Mutex::new(None);
-    let black_found = || black_status.lock_or_panic().is_some();
-    let white_found = || white_status.lock_or_panic().is_some();
+    // 用于保存各种符合条件的背景水印图
+    let backgrounds = Mutex::new(vec![]);
+    // 用于标记是否找到了黑色和白色背景水印图
+    let background_pair_found = Mutex::new(false);
     // 并发遍历image_paths
     let image_paths = image_paths.par_iter();
     image_paths.try_for_each(|path| -> anyhow::Result<()> {
-        // 如果black_path和white_path都已经找到了，则直接跳过
-        if black_found() && white_found() {
+        // 如果已经找到了黑色和白色背景水印图，则直接返回
+        if *background_pair_found.lock() {
             return Ok(());
         }
+
         let mut img = image::open(path)
             .context(format!("打开图片 {} 失败", path.display()))?
             .to_rgb8();
-        let (left, top) = (rect_data.left, rect_data.top);
-        let (right, bottom) = (rect_data.right, rect_data.bottom);
-        // 检查图片是否满足黑色或白色背景的条件
-        let Some(is_black) = is_black_background(&img, &rect_data) else {
+        // 如果图片不满足背景的条件，则直接跳过
+        if !is_background(&img, &rect_data) {
             return Ok(());
         };
         // 获取左上角的颜色
+        let (left, top) = (rect_data.left, rect_data.top);
+        let (right, bottom) = (rect_data.right, rect_data.bottom);
         let color = *img.get_pixel(left, top);
         // 把截图区域外的像素点设置为左上角的颜色
         for (x, y, pixel) in img.enumerate_pixels_mut() {
             if x < left || x > right || y < top || y > bottom {
                 *pixel = color;
             }
         }
-        let filename = if is_black { "black.png" } else { "white.png" };
-        let output_path = output_dir.join(filename);
-        // 保存黑色背景或白色背景的水印图片
-        let mut background_path = if is_black {
-            black_status.lock_or_panic()
-        } else {
-            white_status.lock_or_panic()
-        };
-        // 如果background_path是None，则把output_path赋值给background_path，并保存图片
-        if background_path.is_none() {
-            *background_path = Some(());
-            // 因为save是耗时操作，所以在这里手动释放锁
-            drop(background_path);
-            img.save(&output_path)
-                .context(format!("保存图片 {} 失败", output_path.display()))?;
+        let mut backgrounds = backgrounds.lock();
+        backgrounds.push(img);
+        // 按照像素值排序，保证黑色背景水印图在前，白色背景水印图在后
+        backgrounds.sort_by(|a, b| {
+            let a_color = a.get_pixel(0, 0);
+            let b_color = b.get_pixel(0, 0);
+            a_color[0].cmp(&b_color[0])
+        });
+        if backgrounds.len() < 2 {
+            return Ok(());
+        }
+
+        let black = &backgrounds[0];
+        let white = &backgrounds[backgrounds.len() - 1];
+        // 如果黑色和白色背景水印图的像素值差异大于50，则认为找到了黑色和白色背景水印图
+        let black_color = black.get_pixel(0, 0);
+        let white_color = white.get_pixel(0, 0);
+        if white_color[0] - black_color[0] > 50 {
+            *background_pair_found.lock() = true;
         }
+
         Ok(())
     })?;
+
+    let backgrounds = std::mem::take(&mut *backgrounds.lock());
+    // 如果有第一张背景水印图，则将其保存为黑色背景
+    if let Some(black) = backgrounds.first() {
+        let black_output_path = output_dir.join("black.png");
+        black
+            .save(&black_output_path)
+            .context(format!("保存图片 {} 失败", black_output_path.display()))?;
+    }
+    // 如果找到了黑色和白色背景水印图
+    if *background_pair_found.lock() {
+        // 把最后一张背景水印图保存为白色背景
+        let white = &backgrounds[backgrounds.len() - 1];
+        let white_output_path = output_dir.join("white.png");
+        white
+            .save(&white_output_path)
+            .context(format!("保存图片 {} 失败", white_output_path.display()))?;
+    }
+
     let mut res = CommandResponse {
         code: 0,
         msg: String::new(),
         data: (),
     };
-    if !black_found() {
+    if backgrounds.is_empty() {
         res.code = -1;
-        res.msg += format!("找不到尺寸为({width}x{height})的黑色背景水印图\n").as_str();
-    };
-    if !white_found() {
+        res.msg += format!("找不到尺寸为({width}x{height})的背景水印图\n").as_str();
+    } else if backgrounds.len() == 1 {
         res.code = -1;
-        res.msg += format!("找不到尺寸为({width}x{height})的白色背景水印图\n").as_str();
+        res.msg += format!("只找到一张尺寸为({width}x{height})的背景水印图\n").as_str();
     };
+
     Ok(res)
 }
 
@@ -132,51 +158,39 @@ fn create_image_paths(manga_dir: &str, width: u32, height: u32) -> Vec<PathBuf>
     image_paths
 }
 
-/// 检查图片`img`是否满足黑色背景的条件，如果返回`None`则表示既不满足黑色背景的条件也不满足白色背景的条件
+/// 检查图片`img`是否满足背景的条件
 #[allow(clippy::cast_precision_loss)]
-fn is_black_background(img: &RgbImage, rect_data: &RectData) -> Option<bool> {
+fn is_background(img: &RgbImage, rect_data: &RectData) -> bool {
     let (left, top) = (rect_data.left, rect_data.top);
     let (right, bottom) = (rect_data.right, rect_data.bottom);
     let inside_rect = |x: u32, y: u32| x >= left && x <= right && y >= top && y <= bottom;
     // 获取左上角的颜色
     let color = *img.get_pixel(left, top);
     let [r, g, b] = color.0;
-    // 如果r,g,b通道之间不相等，则不满足黑色背景或白色背景的条件
+    // 如果r,g,b通道之间不相等，则不满足背景的条件
     if r != g || g != b {
-        return None;
+        return false;
     }
-    // 如果截图区域的左右两边的颜色有一个与左上角的颜色不同，则不满足黑色背景或白色背景的条件
+    // 如果截图区域的左右两边的颜色有一个与左上角的颜色不同，则不满足背景的条件
     for y in top..=bottom {
         if img.get_pixel(left, y) != &color || img.get_pixel(right, y) != &color {
-            return None;
+            return false;
         }
     }
-    // 如果截图区域的上下两边的颜色有一个与左上角的颜色不同，则不满足黑色背景或白色背景的条件
+    // 如果截图区域的上下两边的颜色有一个与左上角的颜色不同，则不满足背景的条件
     for x in left..=right {
         if img.get_pixel(x, top) != &color || img.get_pixel(x, bottom) != &color {
-            return None;
+            return false;
         }
     }
     // 统计rect_data区域内color颜色的像素点数量
     let color_count = img
         .enumerate_pixels()
         .filter(|(x, y, &pixel)| inside_rect(*x, *y) && pixel == color)
         .count();
-    // 如果rect_data区域内的像素点数量大于总数的90%，则返回None
+    // 如果rect_data区域内的像素点数量大于总数的90%，则不满足背景的条件
     if color_count as f32 / ((right - left + 1) * (bottom - top + 1)) as f32 > 0.9 {
-        return None;
-    }
-    // 如果color所有通道的值都小于25，则认为是黑色背景
-    let is_black = r <= 25;
-    // 如果color所有通道的值都大于230，并且截图区域内的通道值都大于100(小于100一般是页码)，则认为是白色背景
-    let is_white = r >= 230
-        && img
-            .enumerate_pixels()
-            .filter(|(x, y, _)| inside_rect(*x, *y)) //矩形区域内的像素
-            .all(|(_, _, pixel)| pixel.0[0] > 100); // 通道值大于100
-    match (is_black, is_white) {
-        (true, false) => Some(true),
-        (false, true) => Some(false),
-        _ => None,
+        return false;
     }
+    true
 }

src-tauri/src/commands/get_config.rs

@@ -1,9 +1,7 @@
-use std::sync::RwLock;
-
+use parking_lot::RwLock;
 use tauri::State;
 
 use crate::config::Config;
-use crate::extensions::IgnoreRwLockPoison;
 use crate::types::CommandResponse;
 
 #[tauri::command(async)]
@@ -13,6 +11,6 @@ pub fn get_config(config: State<'_, RwLock<Config>>) -> CommandResponse<Config>
     CommandResponse {
         code: 0,
         msg: String::new(),
-        data: config.read_or_panic().clone(),
+        data: config.read().clone(),
     }
 }

src-tauri/src/commands/remove_watermark.rs

@@ -1,11 +1,11 @@
 use std::collections::HashMap;
 use std::io::BufWriter;
 use std::path::{Path, PathBuf};
-use std::sync::Mutex;
 
 use anyhow::{anyhow, Context};
-use image::{Rgb, RgbImage};
 use image::codecs::png::PngEncoder;
+use image::{Rgb, RgbImage};
+use parking_lot::Mutex;
 use path_slash::PathBufExt;
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use tauri::AppHandle;
@@ -14,7 +14,6 @@ use walkdir::WalkDir;
 
 use crate::errors::CommandResult;
 use crate::events;
-use crate::extensions::IgnoreLockPoison;
 use crate::types::{CommandResponse, ImageFormat, JpgImageData};
 
 #[tauri::command(async)]
@@ -72,7 +71,7 @@ pub fn remove_watermark(
                 .context(format!("保存图片 {} 失败", out_image_path.display()))?;
             // 更新目录的进度
             let (current, total) = {
-                let mut dir_progress = dir_progress.lock_or_panic();
+                let mut dir_progress = dir_progress.lock();
                 let (current, total) = dir_progress
                     .get_mut(dir)
                     .ok_or(anyhow!("目录 {} 的进度不存在", dir.display()))?;
@@ -201,9 +200,9 @@ fn remove_image_watermark(black: &RgbImage, white: &RgbImage, img: &mut RgbImage
         let [white_r, white_g, white_b] = white.get_pixel(x, y).0;
         // 计算去除水印后的像素点值，将f32转换为u8自带clamp功能
         let watermark_removed_pixel = Rgb([
-            ((img_r as f32 - black_r as f32) / ((white_r - black_r) as f32 / 255.0)) as u8,
-            ((img_g as f32 - black_g as f32) / ((white_g - black_g) as f32 / 255.0)) as u8,
-            ((img_b as f32 - black_b as f32) / ((white_b - black_b) as f32 / 255.0)) as u8,
+            ((img_r as f64 - black_r as f64) / ((white_r - black_r) as f64 / 255.0)).round() as u8,
+            ((img_g as f64 - black_g as f64) / ((white_g - black_g) as f64 / 255.0)).round() as u8,
+            ((img_b as f64 - black_b as f64) / ((white_b - black_b) as f64 / 255.0)).round() as u8,
         ]);
         // 将去除水印后的像素点值写入到图片缓冲区中
         *img_pixel = watermark_removed_pixel;

src-tauri/src/commands/save_config.rs

@@ -1,10 +1,8 @@
-use std::sync::RwLock;
-
+use parking_lot::RwLock;
 use tauri::{AppHandle, State};
 
 use crate::config::Config;
 use crate::errors::CommandResult;
-use crate::extensions::IgnoreRwLockPoison;
 use crate::types::CommandResponse;
 
 #[tauri::command(async)]
@@ -15,7 +13,7 @@ pub fn save_config(
     config_state: State<RwLock<Config>>,
     config: Config,
 ) -> CommandResult<CommandResponse<()>> {
-    let mut config_state = config_state.write_or_panic();
+    let mut config_state = config_state.write();
     *config_state = config;
     config_state.save(&app)?;
     Ok(CommandResponse {

src-tauri/src/extensions.rs

@@ -1,35 +1,3 @@
-use std::sync::{Mutex, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard};
-
-pub trait IgnoreLockPoison<T> {
-    fn lock_or_panic(&self) -> MutexGuard<T>;
-}
-impl<T> IgnoreLockPoison<T> for Mutex<T> {
-    /// 如果发生了lock poison，则直接panic
-    #[allow(clippy::unwrap_used)]
-    fn lock_or_panic(&self) -> MutexGuard<T> {
-        self.lock().unwrap()
-    }
-}
-
-pub trait IgnoreRwLockPoison<T> {
-    fn read_or_panic(&self) -> RwLockReadGuard<T>;
-    fn write_or_panic(&self) -> RwLockWriteGuard<T>;
-}
-
-impl<T> IgnoreRwLockPoison<T> for RwLock<T> {
-    /// 如果发生了lock poison，则直接panic
-    #[allow(clippy::unwrap_used)]
-    fn read_or_panic(&self) -> RwLockReadGuard<T> {
-        self.read().unwrap()
-    }
-
-    /// 如果发生了lock poison，则直接panic
-    #[allow(clippy::unwrap_used)]
-    fn write_or_panic(&self) -> RwLockWriteGuard<T> {
-        self.write().unwrap()
-    }
-}
-
 pub trait AnyhowErrorToStringChain {
     /// 将 `anyhow::Error` 转换为chain格式  
     /// # Example

src-tauri/src/main.rs

@@ -2,6 +2,7 @@
 #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
 #![warn(clippy::unwrap_used)]
 
+use parking_lot::RwLock;
 use tauri::{Context, Manager, Wry};
 
 use crate::commands::prelude::*;
@@ -60,7 +61,7 @@ async fn main() {
         .invoke_handler(builder.invoke_handler())
         .setup(move |app| {
             builder.mount_events(app);
-            let config = std::sync::RwLock::new(Config::new(app.handle())?);
+            let config = RwLock::new(Config::new(app.handle())?);
             app.manage(config);
             Ok(())
         })