utils/pdf_ocr.py

#!/usr/bin/env python3
"""
PDFOCR - Module for processing PDF files with OCR and extracting structured data.
Provides robust PDF to image conversion before OCR processing.
"""

import json
import os
import tempfile
import logging
from pathlib import Path
from typing import Optional, Dict, List, Union, Tuple, Any

# Configure logging
logging.basicConfig(level=logging.INFO, 
                   format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("pdf_ocr")

# Import StructuredOCR for OCR processing
from structured_ocr import StructuredOCR

class PDFConversionResult:
    """Class to hold results of PDF to image conversion."""
    
    def __init__(self, 
                 success: bool, 
                 images: List[Path] = None, 
                 error: str = None, 
                 page_count: int = 0,
                 temp_files: List[str] = None):
        """Initialize the conversion result.
        
        Args:
            success: Whether the conversion was successful
            images: List of paths to the converted images
            error: Error message if conversion failed
            page_count: Total number of pages in the PDF
            temp_files: List of temporary files that should be cleaned up
        """
        self.success = success
        self.images = images or []
        self.error = error
        self.page_count = page_count
        self.temp_files = temp_files or []
    
    def __bool__(self):
        """Enable boolean evaluation of the result."""
        return self.success
    
    def cleanup(self):
        """Clean up any temporary files created during conversion."""
        for temp_file in self.temp_files:
            try:
                if os.path.exists(temp_file):
                    os.unlink(temp_file)
                    logger.debug(f"Removed temporary file: {temp_file}")
            except Exception as e:
                logger.warning(f"Failed to remove temporary file {temp_file}: {e}")
        self.temp_files = []


class PDFOCR:
    """Class for processing PDF files with OCR and extracting structured data."""
    
    def __init__(self, api_key=None):
        """Initialize the PDF OCR processor."""
        self.processor = StructuredOCR(api_key=api_key)
        self.temp_files = []
    
    def __del__(self):
        """Clean up resources when object is destroyed."""
        self.cleanup()
    
    def cleanup(self):
        """Clean up any temporary files."""
        for temp_file in self.temp_files:
            try:
                if os.path.exists(temp_file):
                    os.unlink(temp_file)
                    logger.debug(f"Removed temporary file: {temp_file}")
            except Exception as e:
                logger.warning(f"Failed to remove temporary file {temp_file}: {e}")
        self.temp_files = []
    
    def convert_pdf_to_images(self, 
                              pdf_path: Union[str, Path], 
                              dpi: int = 200,
                              max_pages: Optional[int] = None,
                              page_numbers: Optional[List[int]] = None) -> PDFConversionResult:
        """
        Convert a PDF file to images.
        
        Args:
            pdf_path: Path to the PDF file
            dpi: DPI for the output images
            max_pages: Maximum number of pages to convert (None for all)
            page_numbers: Specific page numbers to convert (1-based indexing)
            
        Returns:
            PDFConversionResult object with conversion results
        """
        pdf_path = Path(pdf_path)
        if not pdf_path.exists():
            return PDFConversionResult(
                success=False, 
                error=f"PDF file not found: {pdf_path}"
            )
        
        # Check file size
        file_size_mb = pdf_path.stat().st_size / (1024 * 1024)
        logger.info(f"PDF size: {file_size_mb:.2f} MB")
        
        try:
            # Import pdf2image for conversion
            import pdf2image
            
            # Initialize list for temporary files
            temp_files = []
            
            # Optimize conversion parameters based on file size
            thread_count = min(4, os.cpu_count() or 2)
            
            # First, determine total pages in the document
            logger.info("Determining PDF page count...")
            try:
                # Use a lightweight approach with multi-threading for faster processing
                pdf_info = pdf2image.convert_from_path(
                    pdf_path, 
                    dpi=72,  # Low DPI just for info
                    first_page=1, 
                    last_page=1,
                    size=(100, 100),  # Tiny image to save memory
                    fmt="jpeg", 
                    thread_count=thread_count,
                    output_file=None
                )
                
                # Get page count from poppler info if available
                if hasattr(pdf_info, 'n_pages'):
                    total_pages = pdf_info.n_pages
                else:
                    # Try a different approach to get page count
                    try:
                        from pypdf import PdfReader
                        reader = PdfReader(pdf_path)
                        total_pages = len(reader.pages)
                    except:
                        total_pages = 1
                        logger.warning("Could not determine total page count, assuming 1 page")
            except Exception as e:
                logger.warning(f"Failed to determine page count: {e}")
                total_pages = 1
            
            logger.info(f"PDF has {total_pages} total pages")
            
            # Determine which pages to process
            pages_to_process = []
            
            # If specific pages are requested, use those
            if page_numbers and any(1 <= p <= total_pages for p in page_numbers):
                pages_to_process = [p for p in page_numbers if 1 <= p <= total_pages]
                logger.info(f"Converting {len(pages_to_process)} specified pages: {pages_to_process}")
            # If max_pages is set, limit to that number
            elif max_pages and max_pages < total_pages:
                pages_to_process = list(range(1, max_pages + 1))
                logger.info(f"Converting first {max_pages} pages of {total_pages} total")
            # Otherwise convert all pages if reasonable count
            else:
                pages_to_process = list(range(1, total_pages + 1))
                logger.info(f"Converting all {total_pages} pages")
            
            # Convert PDF to images
            converted_images = []
            
            # Process in batches for better memory management
            batch_size = min(5, len(pages_to_process))  # Process up to 5 pages at a time
            for i in range(0, len(pages_to_process), batch_size):
                batch_pages = pages_to_process[i:i+batch_size]
                logger.info(f"Converting batch of pages {batch_pages}")
                
                # Convert this batch of pages
                try:
                    batch_images = pdf2image.convert_from_path(
                        pdf_path, 
                        dpi=dpi,
                        first_page=min(batch_pages),
                        last_page=max(batch_pages),
                        thread_count=thread_count,
                        fmt="jpeg"
                    )
                    
                    # Map converted images to requested page numbers
                    for idx, page_num in enumerate(range(min(batch_pages), max(batch_pages) + 1)):
                        if page_num in pages_to_process and idx < len(batch_images):
                            # Save the image to a temporary file
                            img_temp_path = tempfile.NamedTemporaryFile(suffix=f'_page{page_num}.jpg', delete=False).name
                            batch_images[idx].save(img_temp_path, format='JPEG', quality=95)
                            
                            # Add to results and track the temp file
                            converted_images.append((page_num, Path(img_temp_path)))
                            temp_files.append(img_temp_path)
                except Exception as e:
                    logger.error(f"Failed to convert batch {batch_pages}: {e}")
                    # Continue with other batches
            
            # Sort by page number to ensure correct order
            converted_images.sort(key=lambda x: x[0])
            
            # Extract just the image paths in correct page order
            image_paths = [img_path for _, img_path in converted_images]
            
            if not image_paths:
                # No images were successfully converted
                return PDFConversionResult(
                    success=False,
                    error="Failed to convert PDF to images",
                    page_count=total_pages,
                    temp_files=temp_files
                )
            
            # Store temp files for later cleanup
            self.temp_files.extend(temp_files)
            
            # Return successful result
            return PDFConversionResult(
                success=True,
                images=image_paths,
                page_count=total_pages,
                temp_files=temp_files
            )
            
        except ImportError:
            return PDFConversionResult(
                success=False,
                error="pdf2image module not available. Please install with: pip install pdf2image"
            )
        except Exception as e:
            logger.error(f"PDF conversion error: {str(e)}")
            return PDFConversionResult(
                success=False,
                error=f"Failed to convert PDF to images: {str(e)}"
            )
    
    def process_pdf(self, pdf_path, use_vision=True, max_pages=None, custom_pages=None, custom_prompt=None):
        """
        Process a PDF file with OCR and extract structured data.
        
        Args:
            pdf_path: Path to the PDF file
            use_vision: Whether to use vision model for improved analysis
            max_pages: Maximum number of pages to process
            custom_pages: Specific page numbers to process (1-based indexing)
            custom_prompt: Custom instructions for processing
            
        Returns:
            Dictionary with structured OCR results
        """
        pdf_path = Path(pdf_path)
        if not pdf_path.exists():
            raise FileNotFoundError(f"PDF file not found: {pdf_path}")
        
        # Convert page numbers to list if provided
        page_numbers = None
        if custom_pages:
            if isinstance(custom_pages, (list, tuple)):
                page_numbers = custom_pages
            else:
                try:
                    # Try to parse as comma-separated string
                    page_numbers = [int(p.strip()) for p in str(custom_pages).split(',')]
                except:
                    logger.warning(f"Invalid custom_pages format: {custom_pages}. Should be list or comma-separated string.")
        
        # First try our optimized PDF to image conversion
        conversion_result = self.convert_pdf_to_images(
            pdf_path=pdf_path,
            max_pages=max_pages,
            page_numbers=page_numbers
        )
        
        if conversion_result.success and conversion_result.images:
            logger.info(f"Successfully converted PDF to {len(conversion_result.images)} images")
            
            # Determine if we need to add PDF-specific context to the prompt
            modified_prompt = custom_prompt
            if not modified_prompt:
                modified_prompt = f"This is a multi-page PDF document with {conversion_result.page_count} total pages, of which {len(conversion_result.images)} were processed."
            elif "pdf" not in modified_prompt.lower() and "multi-page" not in modified_prompt.lower():
                modified_prompt += f" This is a multi-page PDF document with {conversion_result.page_count} total pages, of which {len(conversion_result.images)} were processed."
            
            try:
                # First process the first page with vision if requested
                first_page_result = self.processor.process_file(
                    file_path=conversion_result.images[0],
                    file_type="image",
                    use_vision=use_vision,
                    custom_prompt=modified_prompt
                )
                
                # Process additional pages if available
                all_pages_text = []
                all_languages = set()
                
                # Extract text from first page
                if 'ocr_contents' in first_page_result and 'raw_text' in first_page_result['ocr_contents']:
                    all_pages_text.append(first_page_result['ocr_contents']['raw_text'])
                
                # Track languages from first page
                if 'languages' in first_page_result:
                    for lang in first_page_result['languages']:
                        all_languages.add(str(lang))
                
                # Process additional pages if any
                for i, img_path in enumerate(conversion_result.images[1:], 1):
                    try:
                        # Simple text extraction for additional pages
                        page_result = self.processor.process_file(
                            file_path=img_path,
                            file_type="image",
                            use_vision=False,  # Use simpler processing for additional pages
                            custom_prompt=f"This is page {i+1} of a {conversion_result.page_count}-page document."
                        )
                        
                        # Extract text
                        if 'ocr_contents' in page_result and 'raw_text' in page_result['ocr_contents']:
                            all_pages_text.append(page_result['ocr_contents']['raw_text'])
                        
                        # Track languages
                        if 'languages' in page_result:
                            for lang in page_result['languages']:
                                all_languages.add(str(lang))
                    except Exception as e:
                        logger.warning(f"Error processing page {i+1}: {e}")
                
                # Combine all text into a single document
                combined_text = "\n\n".join(all_pages_text)
                
                # Update the first page result with combined data
                if 'ocr_contents' in first_page_result:
                    first_page_result['ocr_contents']['raw_text'] = combined_text
                
                # Update languages with all detected languages
                if all_languages:
                    first_page_result['languages'] = list(all_languages)
                
                # Add PDF metadata
                first_page_result['file_name'] = pdf_path.name
                first_page_result['file_type'] = "pdf"
                first_page_result['total_pages'] = conversion_result.page_count
                first_page_result['processed_pages'] = len(conversion_result.images)
                
                # Add conversion info
                first_page_result['pdf_conversion'] = {
                    "method": "pdf2image",
                    "pages_converted": len(conversion_result.images),
                    "pages_requested": len(page_numbers) if page_numbers else (max_pages or conversion_result.page_count)
                }
                
                return first_page_result
            except Exception as e:
                logger.error(f"Error processing converted images: {e}")
                # Fall back to direct processing via StructuredOCR
                
            finally:
                # Clean up temporary files
                conversion_result.cleanup()
        
        # If conversion failed or processing the images failed, fall back to direct processing
        logger.info(f"Using direct StructuredOCR processing for PDF")
        return self.processor.process_file(
            file_path=pdf_path, 
            file_type="pdf", 
            use_vision=use_vision,
            max_pages=max_pages,
            custom_pages=custom_pages,
            custom_prompt=custom_prompt
        )
    
    def save_json_output(self, pdf_path, output_path, use_vision=True, max_pages=None, custom_pages=None, custom_prompt=None):
        """
        Process a PDF file and save the structured output as JSON.
        
        Args:
            pdf_path: Path to the PDF file
            output_path: Path where to save the JSON output
            use_vision: Whether to use vision model for improved analysis
            max_pages: Maximum number of pages to process
            custom_pages: Specific page numbers to process (1-based indexing)
            custom_prompt: Custom instructions for processing
            
        Returns:
            Path to the saved JSON file
        """
        # Process the PDF
        result = self.process_pdf(
            pdf_path, 
            use_vision=use_vision,
            max_pages=max_pages,
            custom_pages=custom_pages,
            custom_prompt=custom_prompt
        )
        
        # Save the result to JSON
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        with open(output_path, 'w') as f:
            json.dump(result, f, indent=2)
        
        return output_path

# For testing directly
if __name__ == "__main__":
    import sys
    import argparse
    
    parser = argparse.ArgumentParser(description="Process PDF files with OCR.")
    parser.add_argument("pdf_path", help="Path to the PDF file to process")
    parser.add_argument("--output", "-o", help="Path to save the output JSON")
    parser.add_argument("--no-vision", dest="use_vision", action="store_false", 
                        help="Disable vision model for processing")
    parser.add_argument("--max-pages", type=int, help="Maximum number of pages to process")
    parser.add_argument("--pages", help="Specific pages to process (comma-separated)")
    parser.add_argument("--prompt", help="Custom prompt for processing")
    
    args = parser.parse_args()
    
    processor = PDFOCR()
    
    # Parse custom pages if provided
    custom_pages = None
    if args.pages:
        try:
            custom_pages = [int(p.strip()) for p in args.pages.split(',')]
        except:
            print(f"Error parsing pages: {args.pages}. Should be comma-separated list of numbers.")
            sys.exit(1)
    
    if args.output:
        result_path = processor.save_json_output(
            args.pdf_path, 
            args.output,
            use_vision=args.use_vision,
            max_pages=args.max_pages,
            custom_pages=custom_pages,
            custom_prompt=args.prompt
        )
        print(f"Results saved to: {result_path}")
    else:
        result = processor.process_pdf(
            args.pdf_path,
            use_vision=args.use_vision,
            max_pages=args.max_pages,
            custom_pages=custom_pages,
            custom_prompt=args.prompt
        )
        print(json.dumps(result, indent=2))