Replace Ollama/Windows OCR with ocrs, add clipboard image deduplication
This commit is contained in:
Riz Ashraf committed
1 parent
548f0a06ea
commit
5c5aead1a5
3 files changed
+262
-82
No files matched your search
Generated
+144
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
+116
-82
@@ -289,32 +289,70 @@ pub fn get_windows_html_clipboard() -> Option<String> {
|
||||
}
|
||||
None
|
||||
}
|
||||
use ocrs::{OcrEngine, OcrEngineParams};
|
||||
use std::sync::OnceLock;
|
||||
|
||||
#[cfg(target_os = "windows")]
|
||||
pub fn perform_windows_ocr(img: &image::DynamicImage) -> Option<String> {
|
||||
use windows::Graphics::Imaging::BitmapDecoder;
|
||||
use windows::Media::Ocr::OcrEngine;
|
||||
use windows::Storage::Streams::{DataWriter, InMemoryRandomAccessStream};
|
||||
static OCR_ENGINE: OnceLock<OcrEngine> = 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 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 cache_dir = dirs::home_dir()?.join(".gemini/mcp_memory/models");
|
||||
let _ = std::fs::create_dir_all(&cache_dir);
|
||||
|
||||
let decoder = BitmapDecoder::CreateAsync(&stream).ok()?.join().ok()?;
|
||||
let software_bitmap = decoder.GetSoftwareBitmapAsync().ok()?.join().ok()?;
|
||||
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 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();
|
||||
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 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);
|
||||
|
||||
if let Ok(engine) = OcrEngine::new(params) {
|
||||
let _ = OCR_ENGINE.set(engine);
|
||||
OCR_ENGINE.get()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
pub fn perform_ocrs_ocr(engine: &OcrEngine, img: &image::DynamicImage) -> Option<String> {
|
||||
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<MemoryState>) {
|
||||
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<String> = tokio::task::spawn_blocking(move || {
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
perform_windows_ocr(&img_clone)
|
||||
}
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
None::<String>
|
||||
}
|
||||
})
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
let mut ocr_text_opt: Option<String> = 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;
|
||||
|
||||
// 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
|
||||
};
|
||||
|
||||
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();
|
||||
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<String, Value>, Option<String>)> {
|
||||
let engine = ensure_ocr_engine().await;
|
||||
|
||||
let out =
|
||||
tokio::task::spawn_blocking(move || -> crate::error::Result<serde_json::Map<String, Value>> {
|
||||
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::<String, crate::error::AppError>(serde_json::to_string_pretty(&Value::Object(
|
||||
out,
|
||||
|
||||
Reference in new issue
Block a user