analyze_document.py

"""
Aadhaar Document Classification Pipeline
=========================================
This script runs Part 1: Aadhaar Classification.
It checks whether an uploaded document is an Aadhaar card.

Usage:
    python analyze_document.py path/to/image.jpg

Output:
    JSON result with classification result
"""

import os
import sys
import json
import cv2
import numpy as np
import threading


# ====================================================================================
# THREAD-SAFE MODEL SINGLETON
# ====================================================================================

_model_lock = threading.Lock()
_model = None
_model_path = None


def load_model_once():
    """
    Thread-safe singleton pattern for loading the Keras model.
    Model is loaded only once at import time and reused across all requests.
    
    Returns:
        Loaded Keras model
    """
    global _model, _model_path
    
    if _model is None:
        with _model_lock:
            # Double-check pattern
            if _model is None:
                from tensorflow.keras.models import load_model
                
                _model_path = os.path.join(
                    os.path.dirname(__file__), 
                    "part1", "model", "aadhar_model.keras"
                )
                
                if not os.path.exists(_model_path):
                    raise FileNotFoundError(
                        f"Model not found at {_model_path}. "
                        "Please train the model first by running: python part1/src/train_model.py"
                    )
                
                try:
                    # Load without compiling to save memory
                    _model = load_model(_model_path, compile=False)
                    print(f"✓ Classification model loaded from: {_model_path}")
                except Exception as e:
                    # Clear the global variable so it can retry on next call
                    _model = None
                    raise RuntimeError(f"Failed to load classification model: {str(e)}")
    
    return _model


# ====================================================================================
# STEP 1: AADHAAR CLASSIFICATION (from part1)
# ====================================================================================

class AadhaarClassifier:
    """
    Wrapper for Part 1 classification model.
    Uses singleton pattern to load model once and reuse across requests.
    """
    
    def __init__(self):
        """Initialize classifier configuration."""
        self.img_size = (224, 224)  # Standard size from config
        self.prediction_threshold = 0.5
    
    def predict(self, image_path):
        """
        Predict if the image is an Aadhaar card.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            tuple: (is_aadhaar: bool, confidence: float, raw_score: float)
        """
        # Get singleton model instance (loads once, reuses thereafter)
        model = load_model_once()
        
        # Validate image path
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"Image not found: {image_path}")
        
        try:
            # Load and preprocess image
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError(f"Could not read image file: {image_path}")
            
            # Convert BGR to RGB
            img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            # Resize to model input size
            img_resized = cv2.resize(img_rgb, self.img_size)
            
            # Normalize and add batch dimension
            img_array = np.expand_dims(img_resized / 255.0, axis=0)
            
            # Predict using singleton model
            pred_prob = model.predict(img_array, verbose=0)[0][0]
            
            # Determine class based on threshold
            # Sigmoid output: < 0.5 = Aadhaar (class 0), > 0.5 = Not Aadhaar (class 1)
            is_aadhaar = pred_prob <= self.prediction_threshold
            confidence = (1 - pred_prob) if is_aadhaar else pred_prob
            
            return is_aadhaar, float(confidence), float(pred_prob)
            
        except Exception as e:
            raise RuntimeError(f"Prediction failed: {str(e)}")



# ====================================================================================
# CLASSIFICATION PIPELINE
# ====================================================================================

def analyze_document(image_path, verbose=True):
    """
    Aadhaar classification pipeline (Part 1 only).
    
    Pipeline:
    1. Classify if document is an Aadhaar card
    2. Return classification result
    
    Args:
        image_path: Path to the document image
        verbose: If True, print progress messages
        
    Returns:
        dict: Classification result in JSON format (always returns a dict, never None)
    """
    if verbose:
        print("\n" + "="*70)
        print("AADHAAR DOCUMENT CLASSIFICATION")
        print("="*70)
        print(f"Image: {image_path}")
        print("="*70)
    
    # Validate image path
    if not os.path.exists(image_path):
        return {
            "error": "Image file not found",
            "image_path": image_path,
            "is_aadhaar": None
        }
    
    try:
        # -----------------------------------------------------------------------
        # STEP 1: AADHAAR CLASSIFICATION
        # -----------------------------------------------------------------------
        if verbose:
            print("\n[STEP 1] Classifying document type...")
        
        # Initialize classifier (loads model once)
        classifier = AadhaarClassifier()
        
        # Predict with defensive error handling
        try:
            prediction_result = classifier.predict(image_path)
            
            # Handle different return types
            if prediction_result is None:
                raise ValueError("Classifier returned None - prediction failed")
            
            if isinstance(prediction_result, tuple) and len(prediction_result) >= 2:
                is_aadhaar = prediction_result[0]
                confidence = prediction_result[1]
                raw_score = prediction_result[2] if len(prediction_result) > 2 else confidence
            else:
                raise ValueError(f"Unexpected classifier return type: {type(prediction_result)}")
            
            # Ensure correct types
            is_aadhaar = bool(is_aadhaar)
            confidence = float(confidence)
            raw_score = float(raw_score)
            
        except Exception as classifier_error:
            print(f"✗ Classifier error: {str(classifier_error)}")
            import traceback
            traceback.print_exc()
            return {
                "error": f"Classification failed: {str(classifier_error)}",
                "image_path": image_path,
                "is_aadhaar": None
            }
        
        if verbose:
            classification = "✓ AADHAAR CARD" if is_aadhaar else "✗ NOT AADHAAR"
            print(f"  Result: {classification}")
            print(f"  Confidence: {confidence * 100:.2f}%")
            print(f"  Raw Score: {raw_score:.4f} (< 0.5 = Aadhaar, > 0.5 = Not Aadhaar)")
            print("\n" + "="*70)
            print(f"FINAL RESULT: {'AADHAAR CARD' if is_aadhaar else 'NOT AN AADHAAR CARD'}")
            print("="*70)
        else:
            label = "✓ Aadhaar card" if is_aadhaar else "✗ NOT an Aadhaar card"
            print(f"\n{label} (confidence: {confidence*100:.2f}%)")
        
        return {
            "is_aadhaar": is_aadhaar,
            "classification_confidence": confidence,
            "classification_raw_score": raw_score,
            "message": (
                f"Document is an Aadhaar card (confidence: {confidence*100:.2f}%)."
                if is_aadhaar else
                f"Document is not an Aadhaar card (confidence: {confidence*100:.2f}%)."
            )
        }
        
    except Exception as e:
        import traceback
        print(f"\n✗ ERROR: {str(e)}")
        traceback.print_exc()
        
        return {
            "error": str(e),
            "error_traceback": traceback.format_exc(),
            "image_path": image_path,
            "is_aadhaar": None
        }


# ====================================================================================
# CLI INTERFACE
# ====================================================================================

def main():
    """Command-line interface for the unified pipeline."""
    
    # Check arguments
    if len(sys.argv) < 2:
        print("Usage: python analyze_document.py <image_path>")
        print("\nExample:")
        print("  python analyze_document.py path/to/aadhaar.jpg")
        sys.exit(1)
    
    image_path = sys.argv[1]
    
    # Run analysis
    result = analyze_document(image_path, verbose=True)
    
    # Print JSON result
    print("\n" + "="*70)
    print("JSON OUTPUT")
    print("="*70)
    print(json.dumps(result, indent=2))
    print("="*70)
    
    # Return appropriate exit code
    if "error" in result:
        sys.exit(1)
    else:
        sys.exit(0)


# ====================================================================================
# ENTRY POINT
# ====================================================================================

if __name__ == "__main__":
    main()