| 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 |
|
|
