app.py

"""
Document Forgery Detection - Gradio Interface for Hugging Face Spaces

This app provides a web interface for detecting and classifying document forgeries.
"""

import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
import json
from pathlib import Path
import sys

# Add src to path
sys.path.insert(0, str(Path(__file__).parent))

from src.models import get_model
from src.config import get_config
from src.data.preprocessing import DocumentPreprocessor
from src.data.augmentation import DatasetAwareAugmentation
from src.features.region_extraction import get_mask_refiner, get_region_extractor
from src.features.feature_extraction import get_feature_extractor
from src.training.classifier import ForgeryClassifier

# Class names
CLASS_NAMES = {0: 'Copy-Move', 1: 'Splicing', 2: 'Generation'}
CLASS_COLORS = {
    0: (255, 0, 0),      # Red for Copy-Move
    1: (0, 255, 0),      # Green for Splicing
    2: (0, 0, 255)       # Blue for Generation
}


class ForgeryDetector:
    """Main forgery detection pipeline"""
    
    def __init__(self):
        print("Loading models...")
        
        # Load config
        self.config = get_config('config.yaml')
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Load segmentation model
        self.model = get_model(self.config).to(self.device)
        checkpoint = torch.load('models/best_doctamper.pth', map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.eval()
        
        # Load classifier
        self.classifier = ForgeryClassifier(self.config)
        self.classifier.load('models/classifier')
        
        # Initialize components
        self.preprocessor = DocumentPreprocessor(self.config, 'doctamper')
        self.augmentation = DatasetAwareAugmentation(self.config, 'doctamper', is_training=False)
        self.mask_refiner = get_mask_refiner(self.config)
        self.region_extractor = get_region_extractor(self.config)
        self.feature_extractor = get_feature_extractor(self.config, is_text_document=True)
        
        print("✓ Models loaded successfully!")
    
    def detect(self, image):
        """
        Detect forgeries in document image or PDF
        
        Args:
            image: PIL Image, numpy array, or path to PDF file
            
        Returns:
            overlay_image: Image with detection overlay
            results_json: Detection results as JSON
        """
        # Handle PDF files
        if isinstance(image, str) and image.lower().endswith('.pdf'):
            import fitz  # PyMuPDF
            # Open PDF and convert first page to image
            pdf_document = fitz.open(image)
            page = pdf_document[0]  # First page
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))  # 2x scale for better quality
            image = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
            if pix.n == 4:  # RGBA
                image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
            pdf_document.close()
        
        # Convert PIL to numpy
        if isinstance(image, Image.Image):
            image = np.array(image)
        
        # Convert to RGB
        if len(image.shape) == 2:
            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
        elif image.shape[2] == 4:
            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        
        original_image = image.copy()
        
        # Preprocess
        preprocessed, _ = self.preprocessor(image, None)
        
        # Augment
        augmented = self.augmentation(preprocessed, None)
        image_tensor = augmented['image'].unsqueeze(0).to(self.device)
        
        # Run localization
        with torch.no_grad():
            logits, decoder_features = self.model(image_tensor)
            prob_map = torch.sigmoid(logits).cpu().numpy()[0, 0]
        
        # Refine mask
        binary_mask = (prob_map > 0.5).astype(np.uint8)
        refined_mask = self.mask_refiner.refine(binary_mask, original_size=original_image.shape[:2])
        
        # Extract regions
        regions = self.region_extractor.extract(refined_mask, prob_map, original_image)
        
        # Classify regions
        results = []
        for region in regions:
            # Extract features
            features = self.feature_extractor.extract(
                preprocessed,
                region['region_mask'],
                [f.cpu() for f in decoder_features]
            )
            
            # Reshape features to 2D array (1, n_features) for classifier
            if features.ndim == 1:
                features = features.reshape(1, -1)
            
            # TEMPORARY FIX: Pad features to match classifier's expected count
            expected_features = 526
            current_features = features.shape[1]
            if current_features < expected_features:
                # Pad with zeros
                padding = np.zeros((features.shape[0], expected_features - current_features))
                features = np.hstack([features, padding])
                print(f"Warning: Padded features from {current_features} to {expected_features}")
            elif current_features > expected_features:
                # Truncate
                features = features[:, :expected_features]
                print(f"Warning: Truncated features from {current_features} to {expected_features}")
            
            # Classify
            predictions, confidences = self.classifier.predict(features)
            forgery_type = int(predictions[0])
            confidence = float(confidences[0])
            
            if confidence > 0.6:  # Confidence threshold
                results.append({
                    'region_id': region['region_id'],
                    'bounding_box': region['bounding_box'],
                    'forgery_type': CLASS_NAMES[forgery_type],
                    'confidence': confidence
                })
        
        # Create visualization
        overlay = self._create_overlay(original_image, results)
        
        # Create JSON response
        json_results = {
            'num_detections': len(results),
            'detections': results,
            'model_info': {
                'segmentation_dice': '75%',
                'classifier_accuracy': '92%'
            }
        }
        
        return overlay, json_results
    
    def _create_overlay(self, image, results):
        """Create overlay visualization"""
        overlay = image.copy()
        
        # Draw bounding boxes and labels
        for result in results:
            bbox = result['bounding_box']
            x, y, w, h = bbox
            
            forgery_type = result['forgery_type']
            confidence = result['confidence']
            
            # Get color
            forgery_id = [k for k, v in CLASS_NAMES.items() if v == forgery_type][0]
            color = CLASS_COLORS[forgery_id]
            
            # Draw rectangle
            cv2.rectangle(overlay, (x, y), (x+w, y+h), color, 2)
            
            # Draw label
            label = f"{forgery_type}: {confidence:.1%}"
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
            cv2.rectangle(overlay, (x, y-label_size[1]-10), (x+label_size[0], y), color, -1)
            cv2.putText(overlay, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        
        # Add legend
        if len(results) > 0:
            legend_y = 30
            cv2.putText(overlay, f"Detected {len(results)} forgery region(s)", 
                       (10, legend_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
        
        return overlay


# Initialize detector
detector = ForgeryDetector()


def detect_forgery(file):
    """Gradio interface function"""
    try:
        if file is None:
            return None, {"error": "No file uploaded"}
        
        # Get file path
        file_path = file.name if hasattr(file, 'name') else file
        
        # Check if PDF
        if file_path.lower().endswith('.pdf'):
            # Pass PDF path directly to detector
            overlay, results = detector.detect(file_path)
        else:
            # Load image and pass to detector
            image = Image.open(file_path)
            overlay, results = detector.detect(image)
        
        return overlay, results  # Return dict directly, not json.dumps
    except Exception as e:
        import traceback
        error_details = traceback.format_exc()
        print(f"Error: {error_details}")
        return None, {"error": str(e), "details": error_details}


# Create Gradio interface
demo = gr.Interface(
    fn=detect_forgery,
    inputs=gr.File(label="Upload Document (Image or PDF)", file_types=["image", ".pdf"]),
    outputs=[
        gr.Image(type="numpy", label="Detection Result"),
        gr.JSON(label="Detection Details")
    ],
    title="📄 Document Forgery Detector",
    description="""
    Upload a document image or PDF to detect and classify forgeries.
    
    **Supported Formats:**
    - 📷 Images: JPG, PNG, BMP, TIFF, WebP
    - 📄 PDF: First page will be analyzed
    
    **Supported Forgery Types:**
    - 🔴 Copy-Move: Duplicated regions within the document
    - 🟢 Splicing: Content from different sources
    - 🔵 Generation: AI-generated or synthesized content
    
    **Model Performance:**
    - Localization: 75% Dice Score
    - Classification: 92% Accuracy
    """,
    article="""
    ### About
    This model uses a hybrid deep learning approach:
    1. **Localization**: MobileNetV3-Small + UNet-Lite (detects WHERE)
    2. **Classification**: LightGBM with hybrid features (detects WHAT)
    
    Trained on DocTamper dataset (140K samples).
    """
)


if __name__ == "__main__":
    demo.launch()