ocr_processor.py

"""
med-ocr-toolkit: OCR Document Processor for Healthcare
Extracts structured data from scanned clinical documents.

Author: Monika Sonnad Math
License: MIT
"""

import re
import logging
from pathlib import Path
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger(__name__)


@dataclass
class ExtractionResult:
    """Holds the result of an OCR extraction."""
    raw_text: str
    urns: list[str] = field(default_factory=list)
    dates: list[str] = field(default_factory=list)
    confidence: float = 0.0
    page_count: int = 0
    errors: list[str] = field(default_factory=list)

    @property
    def success(self) -> bool:
        return len(self.errors) == 0 and len(self.urns) > 0


class OCRProcessor:
    """
    Extracts patient URNs and clinical metadata from scanned PDF documents.

    Designed to handle real-world messy scans, not just clean test inputs.
    Uses adaptive preprocessing to improve recognition on low-quality documents.

    Example:
        processor = OCRProcessor()
        result = processor.process("patient_referral.pdf")
        print(result.urns)  # ['URN-2024-00123', ...]
    """

    # Regex patterns for common clinical document formats
    URN_PATTERNS = [
        r'\bURN[-\s]?\d{4}[-\s]?\d{4,6}\b',       # URN-2024-001234
        r'\b[A-Z]{2,3}[-\s]?\d{6,8}\b',             # MR-123456 or NHS-12345678
        r'\bPatient\s*(?:ID|No\.?|Number)[\s:]+(\w{6,12})\b',  # Patient ID: ABC123456
    ]

    DATE_PATTERNS = [
        r'\b\d{1,2}[\/\-\.]\d{1,2}[\/\-\.]\d{2,4}\b',  # 01/01/2024
        r'\b\d{1,2}\s+(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\w*\s+\d{4}\b',
    ]

    def __init__(self, dpi: int = 300, lang: str = "eng"):
        """
        Args:
            dpi: Resolution for PDF-to-image conversion. 300 works well for most scans.
                 Use 400+ for very poor quality documents.
            lang: Tesseract language code. Default English.
        """
        self.dpi = dpi
        self.lang = lang
        self._check_dependencies()

    def _check_dependencies(self):
        """Warn if optional dependencies are missing."""
        try:
            import pytesseract  # noqa
            import pdf2image    # noqa
        except ImportError:
            logger.warning(
                "pytesseract or pdf2image not installed. "
                "Run: pip install pytesseract pdf2image"
            )

    def process(self, pdf_path: str | Path) -> ExtractionResult:
        """
        Main entry point. Converts PDF pages to images, runs OCR,
        and extracts structured fields.

        Args:
            pdf_path: Path to the scanned PDF file.

        Returns:
            ExtractionResult with extracted URNs, dates, and confidence score.
        """
        path = Path(pdf_path)
        if not path.exists():
            return ExtractionResult(raw_text="", errors=[f"File not found: {path}"])

        try:
            images = self._pdf_to_images(path)
            all_text = self._run_ocr(images)
            return self._extract_fields(all_text, page_count=len(images))
        except Exception as e:
            logger.error(f"Processing failed for {path}: {e}")
            return ExtractionResult(raw_text="", errors=[str(e)])

    def _pdf_to_images(self, path: Path):
        """Convert PDF pages to PIL images for OCR."""
        from pdf2image import convert_from_path
        return convert_from_path(str(path), dpi=self.dpi)

    def _run_ocr(self, images) -> str:
        """Run Tesseract OCR across all pages, with basic preprocessing."""
        import pytesseract
        from PIL import ImageFilter, ImageEnhance

        pages_text = []
        for i, img in enumerate(images):
            # Preprocessing: sharpen + increase contrast for messy scans
            img = img.convert("L")                          # grayscale
            img = img.filter(ImageFilter.SHARPEN)
            img = ImageEnhance.Contrast(img).enhance(1.5)

            config = f"--oem 3 --psm 6 -l {self.lang}"
            text = pytesseract.image_to_string(img, config=config)
            pages_text.append(text)
            logger.debug(f"Page {i+1}: extracted {len(text)} characters")

        return "\n".join(pages_text)

    def _extract_fields(self, text: str, page_count: int) -> ExtractionResult:
        """Apply regex patterns to extract URNs, dates, and estimate confidence."""
        urns = []
        for pattern in self.URN_PATTERNS:
            matches = re.findall(pattern, text, re.IGNORECASE)
            urns.extend(matches)

        dates = []
        for pattern in self.DATE_PATTERNS:
            matches = re.findall(pattern, text, re.IGNORECASE)
            dates.extend(matches)

        # Simple confidence heuristic based on text density
        words = len(text.split())
        confidence = min(1.0, words / (page_count * 150)) if page_count > 0 else 0.0

        return ExtractionResult(
            raw_text=text,
            urns=list(set(urns)),       # deduplicate
            dates=list(set(dates)),
            confidence=round(confidence, 2),
            page_count=page_count,
        )

    def process_batch(self, folder: str | Path) -> dict[str, ExtractionResult]:
        """
        Process all PDFs in a folder.

        Args:
            folder: Path to folder containing PDF files.

        Returns:
            Dict mapping filename -> ExtractionResult
        """
        folder = Path(folder)
        results = {}
        pdf_files = list(folder.glob("*.pdf"))
        logger.info(f"Processing {len(pdf_files)} files in {folder}")

        for pdf in pdf_files:
            results[pdf.name] = self.process(pdf)

        return results