File size: 1,292 Bytes
734b5b4 | 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 | import pytesseract
from PIL import Image
import io
from typing import Dict, List, Any
class TextExtractionAgent:
"""PDF text extraction using OCR via Tesseract."""
def __init__(self, tessdata_path: str = None):
self.tessdata_path = tessdata_path
if tessdata_path:
pytesseract.pytesseract.tesseract_cmd = tessdata_path
def process(self, pdf_page: Image.Image) -> Dict[str, Any]:
"""Extract text from a PDF page image."""
text = pytesseract.image_to_string(pdf_page)
return {
"text": text,
"confidence": self._get_confidence(pdf_page),
"blocks": self._extract_blocks(text),
}
def _get_confidence(self, image: Image.Image) -> float:
"""Get OCR confidence score."""
data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
confidences = [int(conf) for conf in data["conf"] if conf != "-1"]
return sum(confidences) / len(confidences) if confidences else 0.0
def _extract_blocks(self, text: str) -> List[Dict]:
"""Extract text blocks."""
blocks = []
for i, line in enumerate(text.split("\n")):
if line.strip():
blocks.append({"id": i, "text": line})
return blocks
|