Spaces:
Sleeping
Sleeping
File size: 9,084 Bytes
7de4594 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 | """
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 |