inference_pipeline.py

"""
Unified Inference Pipeline for IDP
Orchestrates preprocessing, OCR, classification, NER, and post-processing
Optimized for low-latency CPU inference
"""

# Fix for TensorFlow/PaddlePaddle mutex warnings on macOS
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['VECLIB_MAXIMUM_THREADS'] = '1'
os.environ['NUMEXPR_NUM_THREADS'] = '1'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

import warnings
warnings.filterwarnings('ignore')

import cv2
import numpy as np
from typing import Dict, List, Optional, Union, Tuple
import logging
import time
from pathlib import Path
import tempfile

from preprocessing import DocumentPreprocessor
from ocr_engine import LightweightOCR, PDFtoImageConverter
from postprocessing import PostProcessor
# Import demo mode by default; try to import trained-model wrappers lazily
from demo_mode import DemoClassifier, DemoNER

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class IDPPipeline:
    """
    Complete IDP pipeline
    Processes document from image/PDF to structured JSON output
    """
    
    def __init__(
        self,
        classifier_model_path: str = "models/classifier/best_classifier.pt",
        ner_model_path: str = "models/ner/best_ner.pt",
        use_gpu: bool = False,
        ocr_confidence_threshold: float = 0.5
    ):
        """
        Initialize IDP pipeline with all components
        
        Args:
            classifier_model_path: Path to trained classifier weights
            ner_model_path: Path to trained NER weights
            use_gpu: Use GPU for inference
            ocr_confidence_threshold: Minimum OCR confidence
        """
        logger.info("Initializing IDP Pipeline...")
        
        self.use_gpu = use_gpu
        self.ocr_confidence_threshold = ocr_confidence_threshold
        
        # Initialize components
        logger.info("Loading preprocessing...")
        self.preprocessor = DocumentPreprocessor(
            max_width=2048,
            enable_deskew=True,
            enable_denoise=True,
            enable_contrast=True
        )
        
        logger.info("Loading OCR engine...")
        self.ocr = LightweightOCR(
            lang='en',
            use_gpu=use_gpu
        )
        
        logger.info("Loading document classifier (lazy import)")
        try:
            # Lazy import to avoid top-level failures when torch isn't installed
            from classifier_model import DocumentClassifierInference

            self.classifier = DocumentClassifierInference(
                model_path=classifier_model_path,
                device='cuda' if use_gpu else 'cpu',
                use_fp16=use_gpu
            )
            logger.info("✓ Loaded trained classifier model")
        except Exception as e:
            logger.warning(f"Could not import/load trained classifier: {str(e)}")
            logger.warning("⚠ Using demo mode (rule-based) classifier")
            self.classifier = DemoClassifier()
            self.classifier.is_demo = True
        
        logger.info("Loading NER model (lazy import)")
        try:
            # Lazy import to avoid top-level failures when torch isn't installed
            from ner_model import DocumentNERInference

            self.ner = DocumentNERInference(
                model_path=ner_model_path,
                device='cuda' if use_gpu else 'cpu',
                use_fp16=use_gpu
            )
            logger.info("✓ Loaded trained NER model")
        except Exception as e:
            logger.warning(f"Could not import/load trained NER model: {str(e)}")
            logger.warning("⚠ Using demo mode (rule-based) NER")
            self.ner = DemoNER()
            self.ner.is_demo = True
        
        logger.info("Loading post-processor...")
        self.postprocessor = PostProcessor()
        
        logger.info("IDP Pipeline initialized successfully!")
    
    def process_document(
        self,
        file_path: str,
        adaptive_threshold: bool = False,
        page_number: Optional[int] = None
    ) -> Dict:
        """
        Process a document (image or PDF)
        
        Args:
            file_path: Path to document file
            adaptive_threshold: Apply adaptive thresholding for poor quality scans
            page_number: Process specific page (for PDFs), None for all pages
            
        Returns:
            Structured document data with extracted fields
        """
        start_time = time.time()
        logger.info(f"Processing document: {file_path}")
        
        # Check file extension
        ext = Path(file_path).suffix.lower()
        
        if ext == '.pdf':
            results = self._process_pdf(file_path, adaptive_threshold, page_number)
        elif ext in ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']:
            results = self._process_image(file_path, adaptive_threshold)
        else:
            raise ValueError(f"Unsupported file format: {ext}")
        
        total_time = time.time() - start_time
        logger.info(f"Document processing complete in {total_time:.2f}s")
        
        return results
    
    def _process_pdf(
        self,
        pdf_path: str,
        adaptive_threshold: bool,
        page_number: Optional[int]
    ) -> Dict:
        """Process PDF document"""
        # Convert PDF to images
        images = PDFtoImageConverter.pdf_to_images(pdf_path, dpi=200)
        
        if page_number is not None:
            # Process single page
            if page_number < 1 or page_number > len(images):
                raise ValueError(f"Invalid page number: {page_number}")
            
            image = images[page_number - 1]
            result = self._process_single_image(image, adaptive_threshold)
            result['page_number'] = page_number
            
            return {
                'file_type': 'pdf',
                'total_pages': len(images),
                'processed_pages': 1,
                'pages': [result]
            }
        else:
            # Process all pages
            results = []
            for idx, image in enumerate(images):
                logger.info(f"Processing page {idx + 1}/{len(images)}")
                result = self._process_single_image(image, adaptive_threshold)
                result['page_number'] = idx + 1
                results.append(result)
            
            return {
                'file_type': 'pdf',
                'total_pages': len(images),
                'processed_pages': len(results),
                'pages': results
            }
    
    def _process_image(
        self,
        image_path: str,
        adaptive_threshold: bool
    ) -> Dict:
        """Process image file"""
        # Load image
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not load image: {image_path}")
        
        result = self._process_single_image(image, adaptive_threshold)
        
        return {
            'file_type': 'image',
            'total_pages': 1,
            'processed_pages': 1,
            'pages': [result]
        }
    
    def _process_single_image(
        self,
        image: np.ndarray,
        adaptive_threshold: bool
    ) -> Dict:
        """
        Process a single image through the complete pipeline
        
        Returns:
            Structured document data
        """
        # Step 1: Preprocessing
        step_start = time.time()
        preprocessed = self.preprocessor.preprocess(
            image,
            adaptive_threshold=adaptive_threshold
        )
        preprocess_time = time.time() - step_start
        logger.info(f"Preprocessing: {preprocess_time:.3f}s")
        
        # Step 2: OCR
        step_start = time.time()
        ocr_result = self.ocr.extract_text(
            preprocessed,
            return_boxes=True,
            confidence_threshold=self.ocr_confidence_threshold
        )
        ocr_time = time.time() - step_start
        logger.info(f"OCR: {ocr_time:.3f}s, extracted {len(ocr_result['lines'])} lines")
        
        # Step 3: Document Classification
        step_start = time.time()
        # Check if using demo mode or trained model
        if hasattr(self.classifier, 'is_demo') and self.classifier.is_demo:
            # Demo classifier doesn't need return_probabilities parameter
            classification = self.classifier.predict(ocr_result['text'])
        else:
            # Trained classifier
            classification = self.classifier.predict(
                ocr_result['text'],
                return_probabilities=True
            )
        classification_time = time.time() - step_start
        
        # Refine classification
        refined_class, refined_conf = self._refine_classification(
            classification['predicted_class'],
            ocr_result['text'],
            classification['confidence']
        )
        classification['predicted_class'] = refined_class
        classification['confidence'] = refined_conf
        
        logger.info(f"Classification: {classification_time:.3f}s, "
                   f"type={classification['predicted_class']}, "
                   f"conf={classification['confidence']:.3f}")
        
        # Step 4: NER
        step_start = time.time()
        # Check if using demo mode or trained model
        if hasattr(self.ner, 'is_demo') and self.ner.is_demo:
            # Demo NER doesn't need return_token_scores parameter
            ner_result = self.ner.predict(ocr_result['text'])
        else:
            # Trained NER model
            ner_result = self.ner.predict(
                ocr_result['text'],
                return_token_scores=False
            )
        ner_time = time.time() - step_start
        logger.info(f"NER: {ner_time:.3f}s, extracted {len(ner_result['entities'])} entities")
        
        # Step 5: Post-processing
        step_start = time.time()
        final_result = self.postprocessor.process_document(
            document_type=classification['predicted_class'],
            classification_confidence=classification['confidence'],
            ocr_text=ocr_result['text'],
            ocr_boxes=ocr_result['boxes'],
            ner_entities=ner_result['entities']
        )
        postprocess_time = time.time() - step_start
        logger.info(f"Post-processing: {postprocess_time:.3f}s")
        
        # Add timing information
        final_result['processing_time'] = {
            'preprocessing': preprocess_time,
            'ocr': ocr_time,
            'classification': classification_time,
            'ner': ner_time,
            'postprocessing': postprocess_time,
            'total': preprocess_time + ocr_time + classification_time + ner_time + postprocess_time
        }
        
        # Add classification probabilities
        final_result['classification_probabilities'] = classification.get('probabilities', {})
        
        return final_result
    
    def _refine_classification(self, predicted_class: str, text: str, confidence: float) -> Tuple[str, float]:
        """Refine classification based on strong keywords"""
        text_lower = text.lower()
        
        # Strong keywords for Invoice
        if 'invoice' in text_lower:
            # If model says RECEIPT but text says INVOICE, trust text
            if predicted_class == 'RECEIPT':
                return 'INVOICE', float(max(confidence, 0.9))
        
        # Strong keywords for Bank Statement
        if any(k in text_lower for k in ['bank statement', 'account statement', 'statement of account', 'account summary']):
            return 'BANK_STATEMENT', 0.95
            
        return predicted_class, float(confidence)
    
    def process_image_bytes(
        self,
        image_bytes: bytes,
        adaptive_threshold: bool = False
    ) -> Dict:
        """
        Process image from bytes (useful for API endpoints)
        
        Args:
            image_bytes: Image data as bytes
            adaptive_threshold: Apply adaptive thresholding
            
        Returns:
            Structured document data
        """
        # Save to temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp:
            tmp.write(image_bytes)
            tmp_path = tmp.name
        
        try:
            # Process image
            result = self._process_image(tmp_path, adaptive_threshold)
            return result
        finally:
            # Clean up
            if os.path.exists(tmp_path):
                os.remove(tmp_path)


def benchmark_pipeline(
    pipeline: IDPPipeline,
    test_image_path: str,
    num_runs: int = 10
):
    """
    Benchmark pipeline performance
    
    Args:
        pipeline: Initialized pipeline
        test_image_path: Path to test image
        num_runs: Number of benchmark runs
    """
    logger.info(f"Benchmarking pipeline with {num_runs} runs...")
    
    times = []
    for i in range(num_runs):
        start = time.time()
        result = pipeline.process_document(test_image_path)
        elapsed = time.time() - start
        times.append(elapsed)
        logger.info(f"Run {i+1}/{num_runs}: {elapsed:.3f}s")
    
    logger.info(f"\nBenchmark Results:")
    logger.info(f"Mean: {np.mean(times):.3f}s")
    logger.info(f"Median: {np.median(times):.3f}s")
    logger.info(f"Min: {np.min(times):.3f}s")
    logger.info(f"Max: {np.max(times):.3f}s")
    logger.info(f"Std: {np.std(times):.3f}s")


if __name__ == "__main__":
    import sys
    import json
    
    if len(sys.argv) < 2:
        print("Usage: python inference_pipeline.py <image_or_pdf_path>")
        sys.exit(1)
    
    input_path = sys.argv[1]
    
    # Initialize pipeline
    pipeline = IDPPipeline(use_gpu=False)
    
    # Process document
    result = pipeline.process_document(input_path)
    
    # Save result
    output_path = "pipeline_output.json"
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(result, f, indent=2, ensure_ascii=False)
    
    print(f"\nResults saved to {output_path}")
    print(f"\nDocument Type: {result['pages'][0]['document_type']}")
    print(f"Confidence: {result['pages'][0]['classification_confidence']:.3f}")
    print(f"Processing Time: {result['pages'][0]['processing_time']['total']:.3f}s")
    print(f"\nExtracted Fields:")
    for field_name, field_data in result['pages'][0]['fields'].items():
        print(f"  {field_name}: {field_data['value']} (conf: {field_data['confidence']:.3f})")