Compare commits

..

7 Commits

Author SHA1 Message Date
Julian Freeman
2ce9e5471e update version 2026-01-19 15:38:58 -04:00
Julian Freeman
2c89ef42b3 support gpu 2026-01-19 15:35:53 -04:00
Julian Freeman
330a62027c import ocr ares 2026-01-19 15:16:38 -04:00
Julian Freeman
302f106ae6 css 2026-01-19 15:12:59 -04:00
Julian Freeman
3ed264f349 left and right key 2026-01-19 15:00:47 -04:00
Julian Freeman
99c442663a import scrollbar 2026-01-19 14:52:54 -04:00
Julian Freeman
a01277a11c adjust ui 2026-01-19 14:47:44 -04:00
13 changed files with 164 additions and 57 deletions

View File

@@ -1,7 +1,7 @@
{
"name": "watermark-wizard",
"private": true,
"version": "0.1.0",
"version": "1.0.0",
"type": "module",
"scripts": {
"dev": "vite",

2
src-tauri/Cargo.lock generated
View File

@@ -5117,7 +5117,7 @@ dependencies = [
[[package]]
name = "watermark-wizard"
version = "0.1.0"
version = "1.0.0"
dependencies = [
"ab_glyph",
"image",

View File

@@ -1,6 +1,6 @@
[package]
name = "watermark-wizard"
version = "0.1.0"
version = "1.0.0"
description = "A Tauri App handles watermarks"
authors = ["you"]
edition = "2021"
@@ -26,6 +26,6 @@ rayon = "1.10"
tauri-plugin-dialog = "2"
imageproc = "0.25"
ab_glyph = "0.2.23"
ort = { version = "=2.0.0-rc.11", features = ["load-dynamic"] }
ort = { version = "=2.0.0-rc.11", features = ["load-dynamic", "directml"] }
ndarray = "0.16"

View File

@@ -1,5 +1,5 @@
use crate::ort_ops;
use image::{Rgba, RgbaImage};
use ort::session::{Session, builder::GraphOptimizationLevel};
use ort::value::Value;
pub fn run_lama_inpainting(
@@ -8,14 +8,8 @@ pub fn run_lama_inpainting(
mask_image: &image::GrayImage,
) -> Result<RgbaImage, String> {
// 1. Initialize Session
let mut session = Session::builder()
.map_err(|e| format!("Failed to create session builder: {}", e))?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|e| format!("Failed to set opt level: {}", e))?
.with_intra_threads(4)
.map_err(|e| format!("Failed to set threads: {}", e))?
.commit_from_file(model_path)
.map_err(|e| format!("Failed to load model from {:?}: {}", model_path, e))?;
let mut session = ort_ops::create_session(model_path)
.map_err(|e| format!("Failed to create ORT session for LAMA: {}", e))?;
// 2. Preprocess
let target_size = (512, 512);

View File

@@ -94,8 +94,9 @@ use imageproc::drawing::draw_text_mut;
use ab_glyph::{FontRef, PxScale};
use tauri::{AppHandle, Manager};
mod lama;
mod ocr;
pub mod lama;
pub mod ocr;
pub mod ort_ops;
// Embed the font to ensure it's always available without path issues
const FONT_DATA: &[u8] = include_bytes!("../assets/fonts/Roboto-Regular.ttf");
@@ -258,7 +259,7 @@ async fn export_batch(images: Vec<ExportImageTask>, watermark: WatermarkSettings
None => input_path.file_name().unwrap_or_default()
};
let output_path = Path::new(&output_dir).join(file_name);
let output_path = Path::new(&output_dir).join(file_name);
let _output_path = Path::new(&output_dir).join(file_name);
// Handle format specific saving
// JPEG does not support Alpha channel. If we save Rgba8 to Jpeg, it might fail or look wrong.

View File

@@ -1,5 +1,5 @@
use image::RgbaImage;
use ort::session::{Session, builder::GraphOptimizationLevel};
use crate::ort_ops;
use image::{GenericImageView, Rgba, RgbaImage};
use ort::value::Value;
use std::path::Path;
@@ -15,15 +15,9 @@ pub fn run_ocr_detection(
model_path: &Path,
input_image: &RgbaImage,
) -> Result<Vec<DetectedBox>, String> {
// 1. Load Model
let mut session = Session::builder()
.map_err(|e| format!("Failed to create session: {}", e))?
.with_optimization_level(GraphOptimizationLevel::Level3)
.map_err(|e| format!("Failed to set opt: {}", e))?
.with_intra_threads(4)
.map_err(|e| format!("Failed to set threads: {}", e))?
.commit_from_file(model_path)
.map_err(|e| format!("Failed to load OCR model: {}", e))?;
// 1. Load Model using the shared function
let mut session = ort_ops::create_session(model_path)
.map_err(|e| format!("Failed to create ORT session: {}", e))?;
// 2. Preprocess
// DBNet expects standard normalization: (img - mean) / std
@@ -31,7 +25,26 @@ pub fn run_ocr_detection(
// And usually resized to multiple of 32. Limit max size for speed.
let max_side = 1600; // Increase resolution limit
let (orig_w, orig_h) = input_image.dimensions();
// --- CROP to top 5% and bottom 5% ---
let crop_height = (orig_h as f64 * 0.05).ceil() as u32;
let mut masked_image = RgbaImage::from_pixel(orig_w, orig_h, Rgba([0, 0, 0, 255]));
if crop_height > 0 {
// Copy top part
let top_view = input_image.view(0, 0, orig_w, crop_height).to_image();
image::imageops::replace(&mut masked_image, &top_view, 0, 0);
// Copy bottom part
let bottom_y = orig_h.saturating_sub(crop_height);
let bottom_view = input_image.view(0, bottom_y, orig_w, crop_height).to_image();
image::imageops::replace(&mut masked_image, &bottom_view, 0, bottom_y as i64);
} else {
// If image is very short, just use the original
masked_image = input_image.clone();
}
// --- End CROP ---
let mut resize_w = orig_w;
let mut resize_h = orig_h;
@@ -50,7 +63,7 @@ pub fn run_ocr_detection(
resize_w = resize_w.max(32);
resize_h = resize_h.max(32);
let resized = image::imageops::resize(input_image, resize_w, resize_h, image::imageops::FilterType::Triangle);
let resized = image::imageops::resize(&masked_image, resize_w, resize_h, image::imageops::FilterType::Triangle);
let channel_stride = (resize_w * resize_h) as usize;
let mut input_data = Vec::with_capacity(1 * 3 * channel_stride);
@@ -114,7 +127,7 @@ pub fn run_ocr_detection(
let idx = (y * map_w + x) as usize;
if binary_map[idx] && !visited[idx] {
// Flood fill
let mut stack = vec![(x, y)];
let mut stack = vec![(x as u32, y as u32)];
visited[idx] = true;
let mut min_x = x;

41
src-tauri/src/ort_ops.rs Normal file
View File

@@ -0,0 +1,41 @@
use ort::session::{Session, builder::GraphOptimizationLevel};
use ort::execution_providers::DirectMLExecutionProvider;
use std::path::Path;
/// Attempts to create an ORT session with GPU (DirectML) acceleration.
/// If GPU initialization fails, it falls back to a CPU-only session.
pub fn create_session(model_path: &Path) -> Result<Session, ort::Error> {
// Try to build with DirectML
let dm_provider = DirectMLExecutionProvider::default().build();
let session_builder = Session::builder()?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.with_intra_threads(4)?;
match session_builder.with_execution_providers([dm_provider]) {
Ok(builder_with_dm) => {
println!("Attempting to commit session with DirectML provider...");
match builder_with_dm.commit_from_file(model_path) {
Ok(session) => {
println!("Successfully created ORT session with DirectML GPU acceleration.");
return Ok(session);
},
Err(e) => {
println!("Failed to create session with DirectML: {:?}. Falling back to CPU.", e);
// Fall through to CPU execution
}
}
},
Err(e) => {
println!("Failed to build session with DirectML provider: {:?}. Falling back to CPU.", e);
// Fall through to CPU execution
}
};
// Fallback to CPU
println!("Creating ORT session with CPU provider.");
Session::builder()?
.with_optimization_level(GraphOptimizationLevel::Level3)?
.with_intra_threads(4)?
.commit_from_file(model_path)
}

View File

@@ -1,7 +1,7 @@
{
"$schema": "https://schema.tauri.app/config/2",
"productName": "watermark-wizard",
"version": "0.1.0",
"version": "1.0.0",
"identifier": "top.volan.watermark-wizard",
"build": {
"beforeDevCommand": "pnpm dev",
@@ -12,7 +12,7 @@
"app": {
"windows": [
{
"title": "水印精灵 v0.1.0",
"title": "水印精灵 v1.0.0",
"width": 1650,
"height": 1000
}

View File

@@ -6,11 +6,27 @@ import { useGalleryStore } from "./stores/gallery";
import { open } from '@tauri-apps/plugin-dialog';
import { invoke } from '@tauri-apps/api/core';
import { FolderOpen, Download } from 'lucide-vue-next';
import { ref } from "vue";
import { ref, onMounted, onUnmounted } from "vue";
const store = useGalleryStore();
const isExporting = ref(false);
function handleKeydown(event: KeyboardEvent) {
if (event.key === 'ArrowRight') {
store.nextImage();
} else if (event.key === 'ArrowLeft') {
store.prevImage();
}
}
onMounted(() => {
window.addEventListener('keydown', handleKeydown);
});
onUnmounted(() => {
window.removeEventListener('keydown', handleKeydown);
});
async function openFolder() {
try {
const selected = await open({
@@ -94,7 +110,7 @@ async function exportBatch() {
<div class="h-screen w-screen bg-gray-900 text-white overflow-hidden flex flex-col">
<header class="h-14 bg-gray-800 flex items-center justify-between px-6 border-b border-gray-700 shrink-0 shadow-md z-10">
<div class="flex items-center gap-4">
<h1 class="text-lg font-bold tracking-wider bg-gradient-to-r from-blue-400 to-purple-500 bg-clip-text text-transparent">水印精灵</h1>
<h1 class="text-lg font-bold tracking-wider bg-linear-to-r from-blue-400 to-purple-500 bg-clip-text text-transparent">水印精灵</h1>
</div>
<div class="flex items-center gap-3">

View File

@@ -242,24 +242,28 @@ const applyAll = () => {
<!-- Execution Controls -->
<div class="flex flex-col gap-2">
<button
@click="store.processInpainting(store.selectedIndex)"
class="w-full bg-red-600 hover:bg-red-500 text-white py-3 rounded font-bold shadow-lg transition-colors flex items-center justify-center gap-2"
:disabled="store.isProcessing || store.selectedIndex < 0"
>
<div v-if="store.isProcessing" class="w-4 h-4 border-2 border-white/30 border-t-white rounded-full animate-spin"></div>
<Eraser v-else class="w-5 h-5" />
执行移除 (当前)
</button>
<button
@click="store.processAllInpainting()"
class="w-full bg-red-800 hover:bg-red-700 text-white py-2 rounded text-sm font-medium shadow transition-colors flex items-center justify-center gap-2"
:disabled="store.isProcessing || store.images.length === 0"
>
<Eraser class="w-4 h-4" />
执行移除 (全部有遮罩)
</button>
<label class="text-sm text-gray-400 uppercase tracking-wider font-semibold">执行移除</label>
<div class="flex gap-2">
<button
@click="store.processInpainting(store.selectedIndex)"
class="flex-1 bg-red-600 hover:bg-red-500 text-white py-2 rounded flex items-center justify-center gap-2 text-sm transition-colors shadow-lg"
:disabled="store.isProcessing || store.selectedIndex < 0"
>
<div v-if="store.isProcessing" class="w-3 h-3 border-2 border-white/30 border-t-white rounded-full animate-spin"></div>
<Eraser v-else class="w-4 h-4" />
当前
</button>
<button
@click="store.processAllInpainting()"
class="flex-1 bg-red-800 hover:bg-red-700 text-white py-2 rounded flex items-center justify-center gap-2 text-sm transition-colors shadow"
:disabled="store.isProcessing || store.images.length === 0"
>
<div v-if="store.isProcessing" class="w-3 h-3 border-2 border-white/30 border-t-white rounded-full animate-spin"></div>
<Eraser v-else class="w-4 h-4" />
全部
</button>
</div>
</div>
<button

View File

@@ -23,7 +23,7 @@ const onSelect = (index: number) => {
>
<template #default="{ item, index }">
<div
class="h-full w-[100px] p-2 cursor-pointer transition-colors"
class="h-full w-25 p-2 cursor-pointer transition-colors"
:class="{'bg-blue-600': store.selectedIndex === index, 'hover:bg-gray-700': store.selectedIndex !== index}"
@click="onSelect(index)"
>

View File

@@ -236,10 +236,14 @@ export const useGalleryStore = defineStore("gallery", () => {
isProcessing.value = true;
progress.value = { current: 0, total: candidates.length };
try {
// Sequential processing to avoid freezing UI or overloading backend
for (const img of candidates) {
await runInpaintingForImage(img);
progress.value.current++;
// Parallel processing in batches to avoid overwhelming backend
const batchSize = 4;
for (let i = 0; i < candidates.length; i += batchSize) {
const batch = candidates.slice(i, i + batchSize).map(async (img) => {
await runInpaintingForImage(img);
progress.value.current++;
});
await Promise.all(batch);
}
alert("批量处理完成!");
} catch (e) {
@@ -287,6 +291,18 @@ export const useGalleryStore = defineStore("gallery", () => {
}
}
function nextImage() {
if (images.value.length === 0) return;
const nextIndex = (selectedIndex.value + 1) % images.value.length;
selectImage(nextIndex);
}
function prevImage() {
if (images.value.length === 0) return;
const prevIndex = (selectedIndex.value - 1 + images.value.length) % images.value.length;
selectImage(prevIndex);
}
return {
images,
selectedIndex,
@@ -310,6 +326,8 @@ export const useGalleryStore = defineStore("gallery", () => {
recalcAllWatermarks,
processInpainting,
processAllInpainting,
restoreImage
restoreImage,
nextImage,
prevImage
};
});

View File

@@ -5,4 +5,24 @@
src: url('/fonts/Roboto-Regular.ttf') format('truetype');
font-weight: normal;
font-style: normal;
}
/* Custom scrollbar styles */
::-webkit-scrollbar {
width: 8px;
height: 8px;
}
::-webkit-scrollbar-track {
background-color: #1f2937; /* gray-800 */
border-radius: 10px;
}
::-webkit-scrollbar-thumb {
background-color: #4b5563; /* gray-600 */
border-radius: 10px;
}
::-webkit-scrollbar-thumb:hover {
background-color: #6b7280; /* gray-500 */
}