""" OCR Processing Engine Core OCR functionality with vision and traditional OCR """ import io import time import base64 import hashlib import numpy as np import pytesseract from PIL import Image, ImageEnhance import openai from typing import Dict from config import config from logger import ProcessingLogger from corruption_detector import CorruptionDetector from text_processor import ContentFormatter class OCREngine: """Core OCR processing engine with vision and traditional OCR.""" def __init__(self, logger: ProcessingLogger): self.logger = logger self.vision_cache: Dict[str, str] = {} self.vision_calls_used = 0 self.vision_enabled = bool(config.openai_api_key) def preprocess_image(self, img: Image.Image) -> Image.Image: """Preprocess image for better OCR results.""" # Convert to grayscale and enhance img_gray = img.convert('L') enhancer = ImageEnhance.Contrast(img_gray) img_enhanced = enhancer.enhance(1.5) img_array = np.array(img_enhanced) threshold = np.mean(img_array) * 0.85 img_binary = np.where(img_array > threshold, 255, 0).astype(np.uint8) return Image.fromarray(img_binary) # want to try this later # def preprocess_image_advanced(self, img: Image.Image) -> Image.Image: # """Enhanced preprocessing with additional options.""" # # Convert to grayscale # img_gray = img.convert('L') # # # Optional: Denoise before enhancement # img_array = np.array(img_gray) # from scipy.ndimage import median_filter # img_denoised = median_filter(img_array, size=3) # # # Enhance contrast # img_pil = Image.fromarray(img_denoised) # enhancer = ImageEnhance.Contrast(img_pil) # img_enhanced = enhancer.enhance(1.5) # # # Optional: Sharpen text # enhancer_sharp = ImageEnhance.Sharpness(img_enhanced) # img_sharp = enhancer_sharp.enhance(1.2) # # # Binary conversion with Otsu's method (alternative) # from skimage.filters import threshold_otsu # img_array = np.array(img_sharp) # threshold = threshold_otsu(img_array) # More sophisticated than mean # img_binary = np.where(img_array > threshold, 255, 0).astype(np.uint8) # # return Image.fromarray(img_binary) def extract_with_vision(self, page, page_no: int, pdf_text: str) -> tuple[str, bool]: """Extract text using OpenAI Vision API with caching. Returns: (text, success_flag) """ if not self.vision_enabled: self.logger.log_step(f"Page {page_no}", "Vision OCR disabled (no API key)") return "", False text_hash = hashlib.md5(pdf_text.encode()).hexdigest()[:16] if text_hash in self.vision_cache: self.logger.log_step(f"Page {page_no}", "Using cached vision result") return self.vision_cache[text_hash], True self.logger.log_step(f"Page {page_no}", "Attempting vision OCR") start_time = time.time() try: pix = page.get_pixmap(dpi=config.dpi) img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) buffered = io.BytesIO() img.save(buffered, format="PNG") img_base64 = base64.b64encode(buffered.getvalue()).decode() # Vision prompt prompt = """ Extract ALL text from this document maintaining its layout. For regular text: - All headers, body text, footnotes, numbers, dates - Legal text, contact information, disclaimers For tables: - Keep column headers clearly separated from data rows - For multi-line cells, keep lines together with clear cell boundaries - Empty cells should be represented with appropriate spacing - Maintain visual column structure so data aligns under headers Output text exactly as it appears with spatial relationships intact. """ client = openai.OpenAI(api_key=config.openai_api_key) response = client.chat.completions.create( model=config.openai_model, messages=[ {"role": "system", "content": "You are an AI vision specialist focused on complete, accurate text recognition from document images. Capture all content exactly as it appears and provide preserved, clean text output."}, { "role": "user", "content": [ {"type": "text", "text": "Please extract all text from this document image, preserving structure and accuracy. Do not add labels or append processed date."}, { "type": "image_url", "image_url": { "url": f"data:image/png;base64,{img_base64}" } } ] } ], temperature=config.temperature, ) result = response.choices[0].message.content.strip() processing_time = time.time() - start_time # Cache result self.vision_cache[text_hash] = result self.logger.log_success(f"Page {page_no} vision OCR completed in {processing_time:.1f}s - {len(result)} chars") return result, True except Exception as e: self.logger.log_error(f"Page {page_no} vision OCR failed: {e}") return "", False def extract_with_traditional_ocr(self, page, page_no: int) -> str: """Extract text using traditional OCR (Tesseract).""" try: self.logger.log_step(f"Page {page_no}", "Using traditional OCR") pix = page.get_pixmap(dpi=config.dpi) img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) processed_img = self.preprocess_image(img) ocr_text = pytesseract.image_to_string(processed_img, config='--oem 3 --psm 3') result = ContentFormatter.basic_cleanup(ocr_text) self.logger.log_success(f"Page {page_no} traditional OCR completed - {len(result)} chars") return result except Exception as e: self.logger.log_error(f"Page {page_no} traditional OCR failed: {e}") return f"OCR extraction failed for page {page_no}" def extract_page_text(self, page, page_no: int) -> str: """Main text extraction method with intelligent OCR selection.""" try: # Try PDF text extraction first pdf_text = page.get_text("text") if pdf_text and len(pdf_text.strip()) > 30: cleaned_text = ContentFormatter.basic_cleanup(pdf_text.strip()) should_use_vision, reason = CorruptionDetector.should_use_vision( cleaned_text, self.vision_calls_used ) self.logger.log_step( f"Page {page_no}", f"Text length: {len(cleaned_text)}, Vision decision: {should_use_vision} ({reason})" ) if should_use_vision: vision_result, vision_success = self.extract_with_vision(page, page_no, cleaned_text) # If vision succeeded and has good result if vision_success and len(vision_result.strip()) > 30: self.vision_calls_used += 1 self.logger.log_success(f"Page {page_no} using vision result ({len(vision_result)} chars)") return vision_result # If vision failed, fall back to traditional OCR elif not vision_success: self.logger.log_warning(f"Page {page_no} vision failed, falling back to traditional OCR") return self.extract_with_traditional_ocr(page, page_no) # Vision succeeded but result minimal else: self.logger.log_warning(f"Page {page_no} vision result too minimal, using PDF text") return cleaned_text return cleaned_text except Exception as e: self.logger.log_error(f"Page {page_no} PDF extraction failed: {e}") # Fallback to traditional OCR return self.extract_with_traditional_ocr(page, page_no) def get_vision_calls_used(self) -> int: """Get the number of vision API calls used.""" return self.vision_calls_used def reset_vision_counter(self) -> None: """Reset the vision calls counter.""" self.vision_calls_used = 0