Spaces:
Runtime error
Runtime error
File size: 4,843 Bytes
410242f | 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 | """
OCR Engine using Tesseract and OpenCV for document text extraction.
Supports multiple languages and document types.
"""
import pytesseract
import cv2
import numpy as np
from PIL import Image
import logging
from pathlib import Path
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
logger = logging.getLogger(__name__)
@dataclass
class OCRResult:
"""Result from OCR processing"""
text: str
confidence: float
language: str
raw_text: str
bounding_boxes: List[Dict]
class OCREngine:
"""Advanced OCR engine with preprocessing and multi-language support"""
SUPPORTED_LANGUAGES = {
'eng': 'English',
'fra': 'French',
'deu': 'German',
'spa': 'Spanish',
'jpn': 'Japanese',
'chi_sim': 'Chinese (Simplified)',
}
def __init__(self, languages: Optional[List[str]] = None):
"""Initialize OCR engine"""
self.languages = languages or ['eng']
self.lang_string = '+'.join(self.languages)
def preprocess_image(self, image_path: str) -> np.ndarray:
"""Preprocess image for better OCR accuracy"""
img = cv2.imread(image_path)
if img is None:
raise ValueError(f"Cannot read image: {image_path}")
# Convert to grayscale
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
enhanced = clahe.apply(gray)
# Denoise
denoised = cv2.fastNlMeansDenoising(enhanced)
# Threshold
_, binary = cv2.threshold(denoised, 150, 255, cv2.THRESH_BINARY)
# Dilate to connect nearby characters
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))
dilated = cv2.dilate(binary, kernel, iterations=1)
return dilated
def extract_text(self, image_path: str) -> OCRResult:
"""Extract text from image using OCR"""
try:
# Preprocess image
processed = self.preprocess_image(image_path)
# Convert to PIL Image for pytesseract
pil_image = Image.fromarray(processed)
# Extract text with Tesseract
raw_text = pytesseract.image_to_string(
pil_image,
lang=self.lang_string,
config='--psm 3'
)
# Get confidence scores
data = pytesseract.image_to_data(
pil_image,
lang=self.lang_string,
output_type=pytesseract.Output.DICT
)
# Calculate average confidence
confidences = [int(conf) for conf in data['confidence'] if int(conf) > 0]
avg_confidence = np.mean(confidences) / 100 if confidences else 0.0
# Extract bounding boxes
bounding_boxes = []
for i in range(len(data['text'])):
if int(data['conf'][i]) > 0:
bounding_boxes.append({
'text': data['text'][i],
'x': data['left'][i],
'y': data['top'][i],
'width': data['width'][i],
'height': data['height'][i],
'confidence': int(data['conf'][i])
})
# Clean text
cleaned_text = self._clean_text(raw_text)
return OCRResult(
text=cleaned_text,
confidence=avg_confidence,
language=self.lang_string,
raw_text=raw_text,
bounding_boxes=bounding_boxes
)
except Exception as e:
logger.error(f"OCR extraction failed: {str(e)}")
raise
@staticmethod
def _clean_text(text: str) -> str:
"""Clean extracted text"""
# Remove extra whitespace
text = ' '.join(text.split())
# Remove special characters but keep punctuation
text = text.replace('\x00', '')
return text.strip()
def extract_tables(self, image_path: str) -> List[Dict]:
"""Detect and extract table structures from image"""
img = cv2.imread(image_path)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Detect lines
edges = cv2.Canny(gray, 100, 200)
lines = cv2.HoughLinesP(edges, 1, np.pi/180, 100, minLineLength=100, maxLineGap=10)
if lines is None:
return []
# Find table cells
tables = []
# This is a simplified implementation
return tables
|