mve-micro-docs/app/services/ocr_service.py

"""OCR service using Tesseract for scanned PDFs."""
import pytesseract
from pdf2image import convert_from_bytes
from PIL import Image
import logging
from typing import Tuple, List, Optional
from io import BytesIO

logger = logging.getLogger(__name__)


class OCRError(Exception):
    """Custom exception for OCR processing errors."""
    pass


def is_pdf_scanned(text: str, min_char_threshold: int = 50) -> bool:
    """
    Determine if a PDF is scanned (no selectable text) or has text.
    
    Args:
        text: Extracted text from PDF
        min_char_threshold: Minimum characters to consider as "has text"
        
    Returns:
        True if PDF appears to be scanned (no text), False otherwise
    """
    # Remove whitespace and count actual characters
    cleaned_text = text.replace(" ", "").replace("\n", "").replace("\t", "")
    
    if len(cleaned_text) < min_char_threshold:
        logger.info(f"PDF appears to be scanned (only {len(cleaned_text)} characters found)")
        return True
    
    logger.info(f"PDF has selectable text ({len(cleaned_text)} characters)")
    return False


def extract_text_with_ocr(
    pdf_bytes: bytes, 
    lang: str = "spa",
    dpi: int = 300,
    timeout: int = 300
) -> Tuple[str, int]:
    """
    Extract text from scanned PDF using Tesseract OCR.
    
    Args:
        pdf_bytes: PDF file content as bytes
        lang: Language for OCR (default: "spa" for Spanish, use "eng" for English)
        dpi: DPI for image conversion (higher = better quality but slower)
        timeout: Maximum time in seconds for OCR processing
        
    Returns:
        Tuple of (extracted_text, page_count)
        
    Raises:
        OCRError: If OCR processing fails
    """
    try:
        logger.info(f"Starting OCR extraction with language={lang}, dpi={dpi}")
        
        # Convert PDF to images
        images = convert_from_bytes(
            pdf_bytes,
            dpi=dpi,
            fmt='png',
            thread_count=2,
            timeout=timeout
        )
        
        page_count = len(images)
        logger.info(f"Converted PDF to {page_count} images")
        
        # Extract text from each page
        text_parts = []
        for i, image in enumerate(images, 1):
            try:
                # Perform OCR on the image
                page_text = pytesseract.image_to_string(
                    image,
                    lang=lang,
                    timeout=timeout // page_count  # Distribute timeout across pages
                )
                
                if page_text.strip():
                    text_parts.append(f"--- Página {i} ---\n{page_text}")
                    logger.debug(f"Page {i}: Extracted {len(page_text)} characters")
                else:
                    logger.warning(f"Page {i}: No text extracted")
                    text_parts.append(f"--- Página {i} ---\n[Sin texto detectado]\n")
                    
            except Exception as e:
                logger.error(f"OCR failed on page {i}: {str(e)}")
                text_parts.append(f"--- Página {i} ---\n[Error en OCR: {str(e)}]\n")
        
        full_text = "\n\n".join(text_parts)
        
        logger.info(f"OCR extraction completed: {len(full_text)} characters from {page_count} pages")
        return full_text, page_count
        
    except Exception as e:
        logger.error(f"OCR extraction failed: {str(e)}")
        raise OCRError(f"OCR extraction failed: {str(e)}")


def optimize_image_for_ocr(image: Image.Image) -> Image.Image:
    """
    Optimize image for better OCR results.
    
    Args:
        image: PIL Image object
        
    Returns:
        Optimized PIL Image
    """
    # Convert to grayscale
    image = image.convert('L')
    
    # Optional: Apply threshold to make text clearer
    # This can be adjusted based on your PDFs
    # from PIL import ImageEnhance
    # enhancer = ImageEnhance.Contrast(image)
    # image = enhancer.enhance(2)
    
    return image


def extract_text_with_ocr_optimized(
    pdf_bytes: bytes,
    lang: str = "spa",
    dpi: int = 300,
    timeout: int = 300,
    optimize: bool = True
) -> Tuple[str, int]:
    """
    Extract text from scanned PDF with image optimization.
    
    Args:
        pdf_bytes: PDF file content as bytes
        lang: Language for OCR
        dpi: DPI for image conversion
        timeout: Maximum time in seconds
        optimize: Whether to optimize images before OCR
        
    Returns:
        Tuple of (extracted_text, page_count)
        
    Raises:
        OCRError: If OCR processing fails
    """
    try:
        logger.info(f"Starting optimized OCR extraction")
        
        # Convert PDF to images
        images = convert_from_bytes(
            pdf_bytes,
            dpi=dpi,
            fmt='png',
            thread_count=2,
            timeout=timeout
        )
        
        page_count = len(images)
        text_parts = []
        
        for i, image in enumerate(images, 1):
            try:
                # Optimize image if requested
                if optimize:
                    image = optimize_image_for_ocr(image)
                
                # Perform OCR
                page_text = pytesseract.image_to_string(
                    image,
                    lang=lang,
                    config='--psm 1',  # Automatic page segmentation with OSD
                    timeout=timeout // page_count
                )
                
                if page_text.strip():
                    text_parts.append(f"--- Página {i} ---\n{page_text}")
                else:
                    text_parts.append(f"--- Página {i} ---\n[Sin texto detectado]\n")
                    
            except Exception as e:
                logger.error(f"OCR failed on page {i}: {str(e)}")
                text_parts.append(f"--- Página {i} ---\n[Error en OCR]\n")
        
        full_text = "\n\n".join(text_parts)
        
        logger.info(f"Optimized OCR completed: {len(full_text)} characters")
        return full_text, page_count
        
    except Exception as e:
        logger.error(f"Optimized OCR failed: {str(e)}")
        raise OCRError(f"Optimized OCR failed: {str(e)}")