From 5c5aead1a5ad20d5fa732623a0df72a548251ac6 Mon Sep 17 00:00:00 2001 From: Riz Ashraf Date: Thu, 8 Oct 2026 21:48:12 +0100 Subject: [PATCH] Replace Ollama/Windows OCR with ocrs, add clipboard image deduplication --- Cargo.lock | 144 ++++++++++++++++++++++++ server/Cargo.toml | 2 + server/src/handlers/vision.rs | 200 ++++++++++++++++++++-------------- 3 files changed, 263 insertions(+), 83 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 9aa590b..3bf68ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1370,6 +1370,16 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef25905e51abafe4dcea6c15fec58c57b601cdbd0ee53d22ea1d3016c587d39b" +[[package]] +name = "flatbuffers" +version = "24.12.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" +dependencies = [ + "bitflags 1.3.2", + "rustc_version", +] + [[package]] name = "flate2" version = "1.1.10" @@ -2636,9 +2646,11 @@ dependencies = [ "ignore", "image", "notify", + "ocrs", "redb 4.3.0", "reqwest", "rmcp", + "rten", "schemars", "serde", "serde_json", @@ -3107,6 +3119,21 @@ dependencies = [ "objc2-core-foundation", ] +[[package]] +name = "ocrs" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e4192a7dbc7570398ef97ab2606799703b6ea11da53f637d705d44546614435" +dependencies = [ + "anyhow", + "rayon", + "rten", + "rten-imageproc", + "rten-tensor", + "thiserror 2.0.20", + "wasm-bindgen", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -3872,6 +3899,123 @@ dependencies = [ "rmp", ] +[[package]] +name = "rten" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf7741649977427aa13009537ac9a8fe1a503b0aa58784069ce8f117e47eb6ab" +dependencies = [ + "flatbuffers", + "num_cpus", + "rayon", + "rten-base", + "rten-gemm", + "rten-model-file", + "rten-onnx", + "rten-parallel", + "rten-shape-inference", + "rten-simd", + "rten-tensor", + "rten-vecmath", + "rustc-hash", + "smallvec", + "typeid", + "wasm-bindgen", +] + +[[package]] +name = "rten-base" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413dc975c7a3d0e6da0b146cd0387e3d389cceab0fc48328bfa716c7f2ae1b1e" + +[[package]] +name = "rten-gemm" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8efc07c654d4b3f8d81e2efd8b13aea7475568395289e5bb35393411aa6a451e" +dependencies = [ + "rayon", + "rten-base", + "rten-parallel", + "rten-simd", + "rten-tensor", +] + +[[package]] +name = "rten-imageproc" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18726eae9614e23a400560eb3a5e3fdb06526914c31aadaf767f14afde7680dd" +dependencies = [ + "rten-tensor", +] + +[[package]] +name = "rten-model-file" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04ea9e630b377f9249ff3c290466fef282319f86d13d560964209104cd3af542" +dependencies = [ + "flatbuffers", + "rten-base", +] + +[[package]] +name = "rten-onnx" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96944374568990b8d1bbc5bae5957585866f9dd8256130460405730948bfd42e" + +[[package]] +name = "rten-parallel" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba9075b130f81df7be82496ea3a9bbf20fe9e75f106a22b0ece8700850f8f664" +dependencies = [ + "rayon", + "rten-base", +] + +[[package]] +name = "rten-shape-inference" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89289c26b6a5248fccf4d67627596848a03ed31888e6ce4ba6f666c689258122" +dependencies = [ + "rten-tensor", + "smallvec", +] + +[[package]] +name = "rten-simd" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f0c01294eb782adca256aa130b48a04006d03763073a68f45a9c33b982ea556" + +[[package]] +name = "rten-tensor" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf72556ba231d7af50238d26dd264856a2149c6369694eb79b22644c8280c81b" +dependencies = [ + "rayon", + "rten-base", + "rten-parallel", + "smallvec", + "typeid", +] + +[[package]] +name = "rten-vecmath" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1007f05160f9ea2ce4e9c67324e5bafb27ae1e5302d0c6b1ea4db0d35fed78ee" +dependencies = [ + "rten-base", + "rten-simd", +] + [[package]] name = "rust-stemmers" version = "1.2.0" diff --git a/server/Cargo.toml b/server/Cargo.toml index 2b4b51e..f68ac6f 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -49,6 +49,8 @@ notify = "7" blake3 = "1.8" tokio-stream = { version = "0.1", features = ["sync"] } chrono = { version = "0.4.45", features = ["serde"] } +ocrs = "0.13.1" +rten = "0.26.0" [build-dependencies] chrono = "0.4.45" diff --git a/server/src/handlers/vision.rs b/server/src/handlers/vision.rs index 07d7f27..7b8bb3d 100644 --- a/server/src/handlers/vision.rs +++ b/server/src/handlers/vision.rs @@ -289,32 +289,70 @@ pub fn get_windows_html_clipboard() -> Option { } None } +use ocrs::{OcrEngine, OcrEngineParams}; +use std::sync::OnceLock; -#[cfg(target_os = "windows")] -pub fn perform_windows_ocr(img: &image::DynamicImage) -> Option { - use windows::Graphics::Imaging::BitmapDecoder; - use windows::Media::Ocr::OcrEngine; - use windows::Storage::Streams::{DataWriter, InMemoryRandomAccessStream}; +static OCR_ENGINE: OnceLock = OnceLock::new(); - let mut png_bytes = std::io::Cursor::new(Vec::new()); - img.write_to(&mut png_bytes, image::ImageFormat::Png).ok()?; - let bytes = png_bytes.into_inner(); +pub async fn ensure_ocr_engine() -> Option<&'static OcrEngine> { + if let Some(engine) = OCR_ENGINE.get() { + return Some(engine); + } + + let cache_dir = dirs::home_dir()?.join(".gemini/mcp_memory/models"); + let _ = std::fs::create_dir_all(&cache_dir); + + let det_path = cache_dir.join("text-detection.rten"); + if !det_path.exists() || std::fs::metadata(&det_path).map(|m| m.len()).unwrap_or(0) < 100_000 { + if let Ok(resp) = reqwest::get("https://huggingface.co/robertknight/ocrs/resolve/main/text-detection.rten").await { + if let Ok(bytes) = resp.bytes().await { + let _ = std::fs::write(&det_path, bytes); + } + } + } + + let rec_path = cache_dir.join("text-recognition.rten"); + if !rec_path.exists() || std::fs::metadata(&rec_path).map(|m| m.len()).unwrap_or(0) < 100_000 { + if let Ok(resp) = reqwest::get("https://huggingface.co/robertknight/ocrs/resolve/main/text-recognition.rten").await { + if let Ok(bytes) = resp.bytes().await { + let _ = std::fs::write(&rec_path, bytes); + } + } + } - let stream = InMemoryRandomAccessStream::new().ok()?; - let writer = DataWriter::CreateDataWriter(&stream).ok()?; - writer.WriteBytes(&bytes).ok()?; - writer.StoreAsync().ok()?.join().ok()?; - writer.FlushAsync().ok()?.join().ok()?; - stream.Seek(0).ok()?; + let mut params = OcrEngineParams::default(); + let det_model = rten::Model::load_file(&det_path).ok()?; + let rec_model = rten::Model::load_file(&rec_path).ok()?; + params.detection_model = Some(det_model); + params.recognition_model = Some(rec_model); - let decoder = BitmapDecoder::CreateAsync(&stream).ok()?.join().ok()?; - let software_bitmap = decoder.GetSoftwareBitmapAsync().ok()?.join().ok()?; + if let Ok(engine) = OcrEngine::new(params) { + let _ = OCR_ENGINE.set(engine); + OCR_ENGINE.get() + } else { + None + } +} - let engine = OcrEngine::TryCreateFromUserProfileLanguages().ok()?; - let result = engine.RecognizeAsync(&software_bitmap).ok()?.join().ok()?; - let text = result.Text().ok()?; - let s = text.to_string(); - let trimmed = s.trim(); +pub fn perform_ocrs_ocr(engine: &OcrEngine, img: &image::DynamicImage) -> Option { + let rgb_img = img.to_rgb8(); + let img_source = ocrs::ImageSource::from_bytes( + rgb_img.as_raw(), + rgb_img.dimensions(), + ).ok()?; + + let ocr_input = engine.prepare_input(img_source).ok()?; + let word_rects = engine.detect_words(&ocr_input).ok()?; + let line_rects = engine.find_text_lines(&ocr_input, &word_rects); + let texts = engine.recognize_text(&ocr_input, &line_rects).ok()?; + + let mut out = String::new(); + for text in texts.into_iter().flatten() { + out.push_str(&text.to_string()); + out.push('\n'); + } + + let trimmed = out.trim(); if trimmed.is_empty() { None } else { @@ -426,50 +464,56 @@ pub async fn process_clipboard_snapshot(state: Arc) { let path_str = file_path.to_string_lossy().to_string(); let wsl_path = to_wsl_path(&path_str); - // Deterministic OCR - let img_clone = img.clone(); - let ocr_text_opt: Option = tokio::task::spawn_blocking(move || { - #[cfg(target_os = "windows")] - { - perform_windows_ocr(&img_clone) - } - #[cfg(not(target_os = "windows"))] - { - None:: - } - }) - .await - .ok() - .flatten(); + let mut ocr_text_opt: Option = None; - // Index OCR text into Tantivy for omni_search - if let Some(ref text) = ocr_text_opt { - let doc_id = format!("clipboard_{}", now_ms); - let _ = state - .search_index - .read() - .await - .index_clipboard_image(&doc_id, &path_str, text); + if let Some(engine) = ensure_ocr_engine().await { + let img_clone = img.clone(); + ocr_text_opt = tokio::task::spawn_blocking(move || { + perform_ocrs_ocr(engine, &img_clone) + }) + .await + .ok() + .flatten(); } - let cached = crate::state::CachedClipboardImage { - file_path: path_str, - file_path_wsl: wsl_path, - captured_at_epoch_ms: now_ms, - age: "0s ago".to_string(), - width: w, - height: h, - size_bytes, - ocr_text: ocr_text_opt, - }; - let mut cache = state.clipboard_cache.write().await; - cache.last_image = Some(cached.clone()); - cache - .history - .push_front(crate::state::ClipboardHistoryItem::Image(cached)); - if cache.history.len() > 20 { - cache.history.pop_back(); + + // Check if this is an identical image captured rapidly + let is_duplicate = if let Some(ref prev) = cache.last_image { + prev.width == w && prev.height == h && prev.size_bytes == size_bytes + } else { + false + }; + + if !is_duplicate { + // Index OCR text into Tantivy for omni_search + if let Some(ref text) = ocr_text_opt { + let doc_id = format!("clipboard_{}", now_ms); + let _ = state + .search_index + .read() + .await + .index_clipboard_image(&doc_id, &path_str, text); + } + + let cached = crate::state::CachedClipboardImage { + file_path: path_str, + file_path_wsl: wsl_path, + captured_at_epoch_ms: now_ms, + age: "0s ago".to_string(), + width: w, + height: h, + size_bytes, + ocr_text: ocr_text_opt, + }; + + cache.last_image = Some(cached.clone()); + cache + .history + .push_front(crate::state::ClipboardHistoryItem::Image(cached)); + if cache.history.len() > 20 { + cache.history.pop_back(); + } } return; } @@ -783,17 +827,18 @@ impl McpTool for ClipboardHandler { .to_string()) } ClipboardAction::Read => { - let (mut out, b64_opt) = - tokio::task::spawn_blocking(move || -> crate::error::Result<(serde_json::Map, Option)> { + let engine = ensure_ocr_engine().await; + + let out = + tokio::task::spawn_blocking(move || -> crate::error::Result> { let mut out = serde_json::Map::new(); - let mut b64_opt = None; if let Some(text) = get_native_clipboard_text() { out.insert("text".into(), json!(text)); } if let Some(dynamic_img) = get_native_clipboard_image() { - let mut img = dynamic_img; + let mut img = dynamic_img.clone(); let max_dim = 1440; if img.width() > max_dim || img.height() > max_dim { img = img.resize(max_dim, max_dim, FilterType::Lanczos3); @@ -804,9 +849,6 @@ impl McpTool for ClipboardHandler { let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut jpeg_bytes, 88); if encoder.encode(&rgb_img, rgb_img.width(), rgb_img.height(), image::ExtendedColorType::Rgb8).is_ok() { let bytes = jpeg_bytes.into_inner(); - use base64::Engine; - let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); - b64_opt = Some(b64); let cache_dir = dirs::home_dir() .unwrap_or_default() @@ -820,26 +862,18 @@ impl McpTool for ClipboardHandler { out.insert("image_path_wsl".into(), json!(wsl_path)); } } + + if let Some(eng) = engine { + if let Some(ocr_text) = perform_ocrs_ocr(eng, &dynamic_img) { + out.insert("image_analysis".to_string(), json!(ocr_text.trim())); + } + } } - Ok((out, b64_opt)) + Ok(out) }) .await .map_err(|e| crate::error::AppError::Internal(format!("Task panic: {}", e)))??; - if let Some(b64_str) = b64_opt { - if state.ollama.is_available().await - && let Ok(analysis) = state - .ollama - .generate_vision( - "Describe the content, text, or UI layout of this image concisely:", - &b64_str, - ) - .await - { - out.insert("image_analysis".to_string(), json!(analysis.trim())); - } - } - state.record_activity("clipboard", "Read contents from OS clipboard", None); Ok::(serde_json::to_string_pretty(&Value::Object( out,