diff --git a/src-tauri/src/ocr/engine.rs b/src-tauri/src/ocr/engine.rs index 1d6a7ad..d7d9452 100644 --- a/src-tauri/src/ocr/engine.rs +++ b/src-tauri/src/ocr/engine.rs @@ -1,34 +1,22 @@ -use image::{DynamicImage, GrayImage, ImageBuffer, Luma}; -use imageproc::contrast::{stretch_contrast, threshold}; -use leptess::TessBaseApi; +use image::{DynamicImage}; use std::collections::HashMap; -use tracing::{debug, error, info}; +use tracing::{debug, info}; use crate::ocr::calibration::{OCRRegion, OCRCalibration, PreprocessingConfig}; pub struct OCREngine { - tesseract: TessBaseApi, + leptess: leptess::LepTess, calibration: OCRCalibration, } impl OCREngine { pub fn new() -> Result { - let mut tesseract = TessBaseApi::new(); - - // Initialize with English - if tesseract.init(None, "eng").is_err() { - return Err("Failed to initialize Tesseract".to_string()); - } - - // Set OCR engine mode to LSTM only for better accuracy - tesseract.set_variable("tessedit_ocr_engine_mode", "1").ok(); - - // Whitelist characters for numeric/text recognition - tesseract.set_variable("tessedit_char_whitelist", - "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz.,-:()%/ ").ok(); + // Initialize LepTess with English + let leptess = leptess::LepTess::new(None, "eng") + .map_err(|e| format!("Failed to initialize Tesseract: {:?}", e))?; Ok(Self { - tesseract, + leptess, calibration: OCRCalibration::default(), }) } @@ -42,22 +30,21 @@ impl OCREngine { .ok_or_else(|| format!("Region '{}' not found", region_name))?; // Load image - let mut img = image::load_from_memory(screenshot) + let img = image::load_from_memory(screenshot) .map_err(|e| format!("Failed to load image: {}", e))?; - // Crop to region - let cropped = img.crop( - region.x as u32, - region.y as u32, - region.width as u32, - region.height as u32, - ); - - // Preprocess - let processed = self.preprocess_image(&cropped, ®ion.preprocessing); + // Save to temp file for leptess + let temp_path = std::env::temp_dir().join(format!("ocr_region_{}.png", region_name)); + img.save(&temp_path) + .map_err(|e| format!("Failed to save temp image: {}", e))?; // Perform OCR - let text = self.perform_ocr(&processed)?; + self.leptess.set_image(&temp_path.to_string_lossy()); + let text = self.leptess.get_utf8_text() + .map_err(|e| format!("OCR failed: {:?}", e))?; + + // Cleanup + let _ = std::fs::remove_file(&temp_path); debug!("OCR result for '{}': '{}'", region_name, text); @@ -78,7 +65,6 @@ impl OCREngine { pub fn recognize_mob_name(&mut self, screenshot: &[u8]) -> Result { // Mob names appear as floating text above mobs - // This requires detecting text that appears temporarily let text = self.recognize_region("mob_name", screenshot)?; Ok(text) } @@ -112,82 +98,6 @@ impl OCREngine { changes } - fn preprocess_image( - &self, - img: &DynamicImage, - config: &PreprocessingConfig - ) -> DynamicImage { - let mut processed = img.clone(); - - // Convert to grayscale if needed - if config.grayscale { - processed = DynamicImage::ImageLuma8(processed.to_luma8()); - } - - // Apply contrast and brightness - if config.contrast != 1.0 || config.brightness != 0.0 { - processed = self.adjust_contrast_brightness( - &processed, - config.contrast, - config.brightness - ); - } - - // Apply threshold if specified - if let Some(thresh) = config.threshold { - if let DynamicImage::ImageLuma8(gray) = &processed { - let thresholded = threshold(gray, thresh); - processed = DynamicImage::ImageLuma8(thresholded); - } - } - - // Invert if needed - if config.invert { - processed.invert(); - } - - processed - } - - fn adjust_contrast_brightness( - &self, - img: &DynamicImage, - contrast: f32, - brightness: f32 - ) -> DynamicImage { - // Apply contrast stretch - if let DynamicImage::ImageLuma8(gray) = img { - let adjusted = stretch_contrast( - gray, - (brightness * 255.0) as u8, - ((1.0 + contrast) * 255.0) as u8 - ); - DynamicImage::ImageLuma8(adjusted) - } else { - img.clone() - } - } - - fn perform_ocr(&mut self, img: &DynamicImage) -> Result { - // Convert to bytes for Tesseract - let mut bytes: Vec = Vec::new(); - img.write_to( - &mut std::io::Cursor::new(&mut bytes), - image::ImageOutputFormat::Png - ) - .map_err(|e| format!("Failed to encode image: {}", e))?; - - // Set image in Tesseract - self.tesseract.set_image_from_mem(&bytes) - .map_err(|e| format!("Failed to set image: {:?}", e))?; - - // Get text - let text = self.tesseract.get_utf8_text() - .map_err(|e| format!("OCR failed: {:?}", e))?; - - Ok(text) - } - fn parse_hp_text(text: &str) -> Option<(f32, f32)> { // Parse formats: "1234/5678", "1234 / 5678", "1,234/5,678" let cleaned: String = text.chars()