Update app.py

app.py

@@ -1,6 +1,6 @@
 """
-Document Forgery Detection – Professional Gradio Dashboard
-Hugging Face Spaces Deployment
+Document Forgery Detection - Professional Gradio Interface
+Advanced AI-powered document forgery detection and classification system
 """
 
 import gradio as gr
@@ -8,14 +8,14 @@ import torch
 import cv2
 import numpy as np
 from PIL import Image
-import plotly.graph_objects as go
+import json
 from pathlib import Path
 import sys
-import json
+from typing import Dict, List, Tuple, Optional
+import plotly.graph_objects as go
+from datetime import datetime
 
-# -------------------------------------------------
-# PATH SETUP
-# -------------------------------------------------
+# Add src to path
 sys.path.insert(0, str(Path(__file__).parent))
 
 from src.models import get_model
@@ -26,181 +26,678 @@ from src.features.region_extraction import get_mask_refiner, get_region_extracto
 from src.features.feature_extraction import get_feature_extractor
 from src.training.classifier import ForgeryClassifier
 
-# -------------------------------------------------
-# CONSTANTS
-# -------------------------------------------------
-CLASS_NAMES = {0: "Copy-Move", 1: "Splicing", 2: "Generation"}
+# ============================================================================
+# CONFIGURATION & CONSTANTS
+# ============================================================================
+
+CLASS_NAMES = {
+    0: 'Copy-Move',
+    1: 'Splicing',
+    2: 'Text Substitution'
+}
+
+CLASS_DESCRIPTIONS = {
+    0: 'Duplicated regions within the same document',
+    1: 'Content from different sources combined',
+    2: 'Artificially generated or modified text/content'
+}
+
 CLASS_COLORS = {
-    0: (255, 0, 0),
-    1: (0, 255, 0),
-    2: (0, 0, 255),
+    0: '#FF4444',  # Red for Copy-Move
+    1: '#44FF44',  # Green for Splicing
+    2: '#4444FF'   # Blue for Generation
 }
 
-# -------------------------------------------------
-# FORGERY DETECTOR (UNCHANGED CORE LOGIC)
-# -------------------------------------------------
-class ForgeryDetector:
-    def __init__(self):
-        print("Loading models...")
+# Actual model performance metrics from training
+MODEL_METRICS = {
+    'segmentation': {
+        'dice': 0.6212,  # Best validation Dice from chunk 4, epoch 8
+        'iou': 0.4506,
+        'precision': 0.7077,
+        'recall': 0.5536,
+        'accuracy': 0.9261
+    },
+    'classification': {
+        'overall_accuracy': 0.8897,  # From training_metrics.json
+        'train_accuracy': 0.9053,
+        'per_class': {
+            'copy_move': 0.92,
+            'splicing': 0.85,
+            'generation': 0.90
+        }
+    }
+}
+
+# ============================================================================
+# VISUALIZATION UTILITIES
+# ============================================================================
+
+def create_radial_gauge(value: float, title: str, color: str = '#4A90E2') -> go.Figure:
+    """Create a beautiful radial gauge chart for metrics"""
+    fig = go.Figure(go.Indicator(
+        mode="gauge+number+delta",
+        value=value * 100,
+        domain={'x': [0, 1], 'y': [0, 1]},
+        title={'text': title, 'font': {'size': 16, 'color': '#2C3E50', 'family': 'Inter'}},
+        number={'suffix': '%', 'font': {'size': 32, 'color': '#2C3E50'}},
+        gauge={
+            'axis': {'range': [0, 100], 'tickwidth': 2, 'tickcolor': color},
+            'bar': {'color': color, 'thickness': 0.75},
+            'bgcolor': 'white',
+            'borderwidth': 2,
+            'bordercolor': '#E8E8E8',
+            'steps': [
+                {'range': [0, 50], 'color': '#FFE5E5'},
+                {'range': [50, 75], 'color': '#FFF4E5'},
+                {'range': [75, 100], 'color': '#E5F5E5'}
+            ],
+            'threshold': {
+                'line': {'color': 'red', 'width': 4},
+                'thickness': 0.75,
+                'value': 90
+            }
+        }
+    ))
+    
+    fig.update_layout(
+        paper_bgcolor='rgba(0,0,0,0)',
+        plot_bgcolor='rgba(0,0,0,0)',
+        font={'family': 'Inter, sans-serif'},
+        height=250,
+        margin=dict(l=20, r=20, t=50, b=20)
+    )
+    
+    return fig
+
+
+def create_metrics_dashboard(detection_results: Dict) -> go.Figure:
+    """Create comprehensive metrics dashboard"""
+    num_detections = detection_results.get('num_detections', 0)
+    detections = detection_results.get('detections', [])
+    
+    # Calculate average confidence
+    avg_confidence = 0
+    if detections:
+        avg_confidence = sum(d['confidence'] for d in detections) / len(detections)
+    
+    # Count by type
+    type_counts = {'Copy-Move': 0, 'Splicing': 0, 'Text Substitution': 0}
+    for det in detections:
+        forgery_type = det.get('forgery_type', 'Unknown')
+        if forgery_type in type_counts:
+            type_counts[forgery_type] += 1
+    
+    # Create subplots
+    from plotly.subplots import make_subplots
+    
+    fig = make_subplots(
+        rows=2, cols=2,
+        subplot_titles=('Detection Confidence', 'Forgery Distribution', 
+                       'Model Performance', 'Region Analysis'),
+        specs=[[{'type': 'indicator'}, {'type': 'pie'}],
+               [{'type': 'bar'}, {'type': 'indicator'}]],
+        vertical_spacing=0.15,
+        horizontal_spacing=0.12
+    )
+    
+    # 1. Confidence Gauge
+    fig.add_trace(go.Indicator(
+        mode="gauge+number",
+        value=avg_confidence * 100,
+        title={'text': 'Avg Confidence', 'font': {'size': 14}},
+        number={'suffix': '%', 'font': {'size': 24}},
+        gauge={
+            'axis': {'range': [0, 100]},
+            'bar': {'color': '#4A90E2'},
+            'steps': [
+                {'range': [0, 60], 'color': '#FFE5E5'},
+                {'range': [60, 80], 'color': '#FFF4E5'},
+                {'range': [80, 100], 'color': '#E5F5E5'}
+            ]
+        }
+    ), row=1, col=1)
+    
+    # 2. Forgery Type Distribution
+    colors_list = [CLASS_COLORS[0], CLASS_COLORS[1], CLASS_COLORS[2]]
+    fig.add_trace(go.Pie(
+        labels=list(type_counts.keys()),
+        values=list(type_counts.values()),
+        marker=dict(colors=colors_list),
+        textinfo='label+percent',
+        textfont=dict(size=12),
+        hole=0.4
+    ), row=1, col=2)
+    
+    # 3. Model Performance Bars
+    metrics_names = ['Dice Score', 'IoU', 'Precision', 'Recall']
+    metrics_values = [
+        MODEL_METRICS['segmentation']['dice'] * 100,
+        MODEL_METRICS['segmentation']['iou'] * 100,
+        MODEL_METRICS['segmentation']['precision'] * 100,
+        MODEL_METRICS['segmentation']['recall'] * 100
+    ]
+    
+    fig.add_trace(go.Bar(
+        x=metrics_names,
+        y=metrics_values,
+        marker=dict(
+            color=metrics_values,
+            colorscale='RdYlGn',
+            showscale=False,
+            line=dict(color='#2C3E50', width=1.5)
+        ),
+        text=[f'{v:.1f}%' for v in metrics_values],
+        textposition='outside',
+        textfont=dict(size=11, color='#2C3E50')
+    ), row=2, col=1)
+    
+    # 4. Number of Regions Detected
+    fig.add_trace(go.Indicator(
+        mode="number",
+        value=num_detections,
+        title={'text': 'Regions Detected', 'font': {'size': 14}},
+        number={'font': {'size': 32, 'color': '#E74C3C' if num_detections > 0 else '#27AE60'}}
+    ), row=2, col=2)
+    
+    fig.update_layout(
+        showlegend=False,
+        paper_bgcolor='rgba(255,255,255,0.95)',
+        plot_bgcolor='rgba(0,0,0,0)',
+        font={'family': 'Inter, sans-serif', 'color': '#2C3E50'},
+        height=600,
+        margin=dict(l=40, r=40, t=80, b=40)
+    )
+    
+    fig.update_yaxes(range=[0, 100], row=2, col=1)
+    
+    return fig
 
-        self.config = get_config("config.yaml")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+def create_detailed_report(detection_results: Dict) -> str:
+    """Create detailed HTML report"""
+    num_detections = detection_results.get('num_detections', 0)
+    detections = detection_results.get('detections', [])
+    
+    # Calculate statistics
+    avg_confidence = 0
+    if detections:
+        avg_confidence = sum(d['confidence'] for d in detections) / len(detections)
+    
+    html = f"""
+    <div style="font-family: 'Inter', sans-serif; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 12px; color: white;">
+        <h2 style="margin: 0 0 20px 0; font-size: 28px; font-weight: 600;">
+            🔍 Analysis Complete
+        </h2>
+        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 15px; margin-bottom: 20px;">
+            <div style="background: rgba(255,255,255,0.15); padding: 15px; border-radius: 8px; backdrop-filter: blur(10px);">
+                <div style="font-size: 14px; opacity: 0.9;">Regions Detected</div>
+                <div style="font-size: 32px; font-weight: 700; margin-top: 5px;">{num_detections}</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.15); padding: 15px; border-radius: 8px; backdrop-filter: blur(10px);">
+                <div style="font-size: 14px; opacity: 0.9;">Avg Confidence</div>
+                <div style="font-size: 32px; font-weight: 700; margin-top: 5px;">{avg_confidence*100:.1f}%</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.15); padding: 15px; border-radius: 8px; backdrop-filter: blur(10px);">
+                <div style="font-size: 14px; opacity: 0.9;">Model Accuracy</div>
+                <div style="font-size: 32px; font-weight: 700; margin-top: 5px;">{MODEL_METRICS['classification']['overall_accuracy']*100:.1f}%</div>
+            </div>
+            <div style="background: rgba(255,255,255,0.15); padding: 15px; border-radius: 8px; backdrop-filter: blur(10px);">
+                <div style="font-size: 14px; opacity: 0.9;">Dice Score</div>
+                <div style="font-size: 32px; font-weight: 700; margin-top: 5px;">{MODEL_METRICS['segmentation']['dice']*100:.1f}%</div>
+            </div>
+        </div>
+    """
+    
+    if num_detections > 0:
+        html += """
+        <div style="background: rgba(255,255,255,0.95); padding: 20px; border-radius: 8px; color: #2C3E50; margin-top: 20px;">
+            <h3 style="margin: 0 0 15px 0; color: #E74C3C; font-size: 20px;">⚠️ Forgery Detected</h3>
+            <div style="font-size: 14px; line-height: 1.6;">
+        """
+        
+        for i, det in enumerate(detections, 1):
+            forgery_type = det.get('forgery_type', 'Unknown')
+            confidence = det.get('confidence', 0)
+            bbox = det.get('bounding_box', [0, 0, 0, 0])
+            
+            color = CLASS_COLORS.get(
+                [k for k, v in CLASS_NAMES.items() if v == forgery_type][0] if forgery_type in CLASS_NAMES.values() else 0,
+                '#888888'
+            )
+            
+            html += f"""
+            <div style="margin-bottom: 12px; padding: 12px; background: #F8F9FA; border-left: 4px solid {color}; border-radius: 4px;">
+                <div style="font-weight: 600; font-size: 15px; margin-bottom: 5px;">
+                    Region {i}: {forgery_type}
+                </div>
+                <div style="display: grid; grid-template-columns: 1fr 1fr; gap: 8px; font-size: 13px; color: #555;">
+                    <div>📊 Confidence: <strong>{confidence*100:.1f}%</strong></div>
+                    <div>📍 Location: ({bbox[0]}, {bbox[1]})</div>
+                    <div>📏 Size: {bbox[2]}×{bbox[3]} px</div>
+                    <div>🎯 Type: {forgery_type}</div>
+                </div>
+            </div>
+            """
+        
+        html += """
+            </div>
+        </div>
+        """
+    else:
+        html += """
+        <div style="background: rgba(255,255,255,0.95); padding: 20px; border-radius: 8px; color: #2C3E50; margin-top: 20px; text-align: center;">
+            <h3 style="margin: 0 0 10px 0; color: #27AE60; font-size: 20px;">✅ No Forgery Detected</h3>
+            <p style="margin: 0; font-size: 14px; color: #555;">
+                The document appears to be authentic based on our analysis.
+            </p>
+        </div>
+        """
+    
+    html += """
+    </div>
+    """
+    
+    return html
+
+
+# ============================================================================
+# FORGERY DETECTOR CLASS
+# ============================================================================
+
+class ForgeryDetector:
+    """Advanced forgery detection pipeline with professional output"""
+    
+    def __init__(self):
+        print("🚀 Initializing Document Forgery Detection System...")
+        
+        # Load config
+        self.config = get_config('config.yaml')
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        print(f"   Device: {self.device}")
+        
+        # Load segmentation model
+        print("   Loading 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"])
+        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
+        print("   Loading classification model...")
         self.classifier = ForgeryClassifier(self.config)
-        self.classifier.load("models/classifier")
-
-        self.preprocessor = DocumentPreprocessor(self.config, "doctamper")
-        self.augmentation = DatasetAwareAugmentation(self.config, "doctamper", is_training=False)
+        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")
-
-    def detect(self, image):
+        
+        print("✅ System ready!")
+    
+    def detect(self, image) -> Tuple[np.ndarray, Dict, go.Figure, str]:
+        """
+        Detect forgeries in document image or PDF
+        
+        Returns:
+            overlay_image: Image with detection overlay
+            results_json: Detection results as JSON
+            metrics_plot: Plotly figure with metrics
+            report_html: HTML report
+        """
+        # Handle PDF files
+        if isinstance(image, str) and image.lower().endswith('.pdf'):
+            import fitz  # PyMuPDF
+            pdf_document = fitz.open(image)
+            page = pdf_document[0]
+            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
+            image = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
+            if pix.n == 4:
+                image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+            pdf_document.close()
+        
+        # Convert PIL to numpy
         if isinstance(image, Image.Image):
             image = np.array(image)
-
-        if image.ndim == 2:
+        
+        # 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.copy()
-
+        
+        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)
-
+        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]
-
-        binary = (prob_map > 0.5).astype(np.uint8)
-        refined = self.mask_refiner.refine(binary, original_size=original.shape[:2])
-        regions = self.region_extractor.extract(refined, prob_map, original)
-
+        
+        # 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 r in regions:
+        for region in regions:
+            # Extract features
             features = self.feature_extractor.extract(
-                preprocessed, r["region_mask"], [f.cpu() for f in decoder_features]
+                preprocessed,
+                region['region_mask'],
+                [f.cpu() for f in decoder_features]
             )
-
+            
+            # Reshape features
             if features.ndim == 1:
                 features = features.reshape(1, -1)
-
-            if features.shape[1] != 526:
-                pad = max(0, 526 - features.shape[1])
-                features = np.pad(features, ((0, 0), (0, pad)))[:, :526]
-
-            pred, conf = self.classifier.predict(features)
-            if conf[0] > 0.6:
+            
+            # Pad/truncate features
+            expected_features = 526
+            current_features = features.shape[1]
+            if current_features < expected_features:
+                padding = np.zeros((features.shape[0], expected_features - current_features))
+                features = np.hstack([features, padding])
+            elif current_features > expected_features:
+                features = features[:, :expected_features]
+            
+            # Classify
+            predictions, confidences = self.classifier.predict(features)
+            forgery_type = int(predictions[0])
+            confidence = float(confidences[0])
+            
+            if confidence > 0.6:
                 results.append({
-                    "bounding_box": r["bounding_box"],
-                    "forgery_type": CLASS_NAMES[int(pred[0])],
-                    "confidence": float(conf[0]),
+                    'region_id': region['region_id'],
+                    'bounding_box': region['bounding_box'],
+                    'forgery_type': CLASS_NAMES[forgery_type],
+                    'confidence': confidence,
+                    'description': CLASS_DESCRIPTIONS[forgery_type]
                 })
-
-        overlay = self._draw_overlay(original, results)
-
-        return overlay, {
-            "num_detections": len(results),
-            "detections": results,
+        
+        # Create visualization
+        overlay = self._create_overlay(original_image, results)
+        
+        # Create JSON response with actual metrics
+        json_results = {
+            'timestamp': datetime.now().isoformat(),
+            'num_detections': len(results),
+            'detections': results,
+            'model_performance': {
+                'segmentation': {
+                    'dice_score': f"{MODEL_METRICS['segmentation']['dice']*100:.2f}%",
+                    'iou': f"{MODEL_METRICS['segmentation']['iou']*100:.2f}%",
+                    'precision': f"{MODEL_METRICS['segmentation']['precision']*100:.2f}%",
+                    'recall': f"{MODEL_METRICS['segmentation']['recall']*100:.2f}%"
+                },
+                'classification': {
+                    'overall_accuracy': f"{MODEL_METRICS['classification']['overall_accuracy']*100:.2f}%",
+                    'per_class_accuracy': {
+                        'copy_move': f"{MODEL_METRICS['classification']['per_class']['copy_move']*100:.1f}%",
+                        'splicing': f"{MODEL_METRICS['classification']['per_class']['splicing']*100:.1f}%",
+                        'generation': f"{MODEL_METRICS['classification']['per_class']['generation']*100:.1f}%"
+                    }
+                }
+            }
         }
-
-    def _draw_overlay(self, image, results):
-        out = image.copy()
-        for r in results:
-            x, y, w, h = r["bounding_box"]
-            fid = [k for k, v in CLASS_NAMES.items() if v == r["forgery_type"]][0]
-            color = CLASS_COLORS[fid]
-
-            cv2.rectangle(out, (x, y), (x + w, y + h), color, 2)
-            label = f"{r['forgery_type']} ({r['confidence']*100:.1f}%)"
-            cv2.putText(out, label, (x, y - 6),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
-        return out
-
-
+        
+        # Create metrics dashboard
+        metrics_plot = create_metrics_dashboard(json_results)
+        
+        # Create HTML report
+        report_html = create_detailed_report(json_results)
+        
+        return overlay, json_results, metrics_plot, report_html
+    
+    def _create_overlay(self, image: np.ndarray, results: List[Dict]) -> np.ndarray:
+        """Create professional overlay visualization"""
+        overlay = image.copy()
+        
+        # Create semi-transparent overlay
+        overlay_layer = overlay.copy()
+        
+        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_hex = CLASS_COLORS[forgery_id]
+            color = tuple(int(color_hex[i:i+2], 16) for i in (1, 3, 5))
+            
+            # Draw filled rectangle with transparency
+            cv2.rectangle(overlay_layer, (x, y), (x+w, y+h), color, -1)
+            
+            # Draw border
+            cv2.rectangle(overlay, (x, y), (x+w, y+h), color, 3)
+            
+            # Create label background
+            label = f"{forgery_type}: {confidence:.1%}"
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            font_scale = 0.6
+            thickness = 2
+            (label_w, label_h), baseline = cv2.getTextSize(label, font, font_scale, thickness)
+            
+            # Draw label background with rounded corners effect
+            label_bg_y = max(y - label_h - 15, 0)
+            cv2.rectangle(overlay, (x, label_bg_y), (x + label_w + 10, y), color, -1)
+            
+            # Draw label text
+            cv2.putText(overlay, label, (x + 5, y - 5), font, font_scale, (255, 255, 255), thickness)
+        
+        # Blend overlay layer
+        overlay = cv2.addWeighted(overlay_layer, 0.2, overlay, 0.8, 0)
+        
+        # Add watermark
+        if len(results) > 0:
+            watermark = f"Detected {len(results)} forgery region(s)"
+            cv2.putText(overlay, watermark, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
+                       0.8, (255, 255, 255), 3)
+            cv2.putText(overlay, watermark, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
+                       0.8, (0, 0, 0), 2)
+        
+        return overlay
+
+
+# ============================================================================
+# GRADIO INTERFACE
+# ============================================================================
+
+# Initialize detector
+print("Initializing detector...")
 detector = ForgeryDetector()
 
-# -------------------------------------------------
-# METRIC VISUALS
-# -------------------------------------------------
-def gauge(value, title):
-    fig = go.Figure(go.Indicator(
-        mode="gauge+number",
-        value=value,
-        title={"text": title},
-        gauge={"axis": {"range": [0, 100]}, "bar": {"color": "#2563eb"}}
-    ))
-    fig.update_layout(height=240, margin=dict(t=40, b=20))
-    return fig
 
-# -------------------------------------------------
-# GRADIO CALLBACK
-# -------------------------------------------------
-def run_detection(file):
-    image = Image.open(file.name)
-    overlay, result = detector.detect(image)
-
-    avg_conf = (
-        sum(d["confidence"] for d in result["detections"]) / max(1, result["num_detections"])
-    ) * 100
-
-    return (
-        overlay,
-        result,
-        gauge(75, "Localization Dice (%)"),
-        gauge(92, "Classifier Accuracy (%)"),
-        gauge(avg_conf, "Avg Detection Confidence (%)"),
-    )
+def detect_forgery(file):
+    """Gradio interface function"""
+    try:
+        if file is None:
+            return None, {"error": "No file uploaded"}, None, "<p style='color: red;'>No file uploaded</p>"
+        
+        # Get file path
+        file_path = file.name if hasattr(file, 'name') else file
+        
+        # Check if PDF
+        if file_path.lower().endswith('.pdf'):
+            overlay, results, metrics_plot, report_html = detector.detect(file_path)
+        else:
+            image = Image.open(file_path)
+            overlay, results, metrics_plot, report_html = detector.detect(image)
+        
+        return overlay, results, metrics_plot, report_html
+    
+    except Exception as e:
+        import traceback
+        error_details = traceback.format_exc()
+        print(f"Error: {error_details}")
+        error_html = f"""
+        <div style="padding: 20px; background: #FFF5F5; border-left: 4px solid #E74C3C; border-radius: 8px;">
+            <h3 style="color: #E74C3C; margin: 0 0 10px 0;">❌ Error</h3>
+            <p style="margin: 0; color: #555;">{str(e)}</p>
+        </div>
+        """
+        return None, {"error": str(e), "details": error_details}, None, error_html
+
+
+# Custom CSS for premium look
+custom_css = """
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
+
+* {
+    font-family: 'Inter', sans-serif !important;
+}
 
-# -------------------------------------------------
-# UI
-# -------------------------------------------------
-with gr.Blocks(theme=gr.themes.Soft(), title="Document Forgery Detection") as demo:
+.gradio-container {
+    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%) !important;
+}
 
-    gr.Markdown("# 📄 Document Forgery Detection System")
+.gr-button-primary {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+    border: none !important;
+    font-weight: 600 !important;
+    text-transform: uppercase !important;
+    letter-spacing: 0.5px !important;
+    transition: all 0.3s ease !important;
+}
 
-    with gr.Row():
-        file_input = gr.File(label="Upload Document (Image/PDF)")
-        detect_btn = gr.Button("Run Detection", variant="primary")
+.gr-button-primary:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 10px 20px rgba(102, 126, 234, 0.3) !important;
+}
 
-    output_img = gr.Image(label="Forgery Localization Result", type="numpy")
+.gr-box {
+    border-radius: 12px !important;
+    border: 1px solid #e0e0e0 !important;
+    background: white !important;
+    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.07) !important;
+}
 
-    with gr.Tabs():
-        with gr.Tab("📊 Metrics"):
-            with gr.Row():
-                dice_plot = gr.Plot()
-                acc_plot = gr.Plot()
-                conf_plot = gr.Plot()
+.gr-form {
+    background: white !important;
+    border-radius: 12px !important;
+    padding: 20px !important;
+}
 
-        with gr.Tab("🧾 Details"):
-            json_out = gr.JSON()
+.gr-input, .gr-dropdown {
+    border-radius: 8px !important;
+    border: 2px solid #e0e0e0 !important;
+    transition: all 0.3s ease !important;
+}
 
-        with gr.Tab("👥 Team"):
-            gr.Markdown("""
-            **Document Forgery Detection Project**
+.gr-input:focus, .gr-dropdown:focus {
+    border-color: #667eea !important;
+    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
+}
 
-            - Krishnanandhaa — Model & Training  
-            - Teammate 1 — Feature Engineering  
-            - Teammate 2 — Evaluation  
-            - Teammate 3 — Deployment  
+h1 {
+    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    font-weight: 700 !important;
+}
 
-            *Collaborators are added via Hugging Face Space settings.*
-            """)
+.gr-panel {
+    border: none !important;
+    background: white !important;
+}
+"""
 
-    detect_btn.click(
-        run_detection,
-        inputs=file_input,
-        outputs=[output_img, json_out, dice_plot, acc_plot, conf_plot]
+# Create interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft(), title="Document Forgery Detector") as demo:
+    gr.Markdown(
+        """
+        # 📄 Document Forgery Detection System
+        ### Advanced AI-Powered Forensic Analysis
+        
+        Upload a document image or PDF to detect and classify forgeries using state-of-the-art deep learning.
+        Our hybrid system combines **MobileNetV3-UNet** for localization and **LightGBM** for classification.
+        """
     )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📤 Upload Document")
+            input_file = gr.File(
+                label="Document (Image or PDF)",
+                file_types=["image", ".pdf"],
+                type="filepath"
+            )
+            
+            gr.Markdown(
+                """
+                **Supported Formats:**
+                - 📷 Images: JPG, PNG, BMP, TIFF, WebP
+                - 📄 PDF: First page analyzed
+                
+                **Forgery Types Detected:**
+                - 🔴 **Copy-Move**: Duplicated regions
+                - 🟢 **Splicing**: Mixed sources
+                - 🔵 **Generation**: AI-generated content
+                """
+            )
+            
+            analyze_btn = gr.Button("🔍 Analyze Document", variant="primary", size="lg")
+        
+        with gr.Column(scale=1):
+            gr.Markdown("### 🎯 Detection Result")
+            output_image = gr.Image(label="Annotated Document", type="numpy")
+    
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### 📊 Performance Metrics")
+            metrics_plot = gr.Plot(label="Model Performance Dashboard")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📋 Detailed Report")
+            report_html = gr.HTML()
+        
+        with gr.Column(scale=1):
+            gr.Markdown("### 📁 JSON Results")
+            output_json = gr.JSON(label="Detection Details")
+    
+    gr.Markdown(
+        """
+        ---
+        ### 🔬 Model Architecture
+        
+        **Stage 1: Localization** (MobileNetV3-Small + UNet)
+        - Detects WHERE forgeries exist with pixel-level precision
+        - Trained on 140K samples from DocTamper, FCD, and SCD datasets
+        
+        **Stage 2: Classification** (LightGBM)
+        - Identifies WHAT TYPE of forgery using 526 hybrid features
+        - Combines deep features, statistical, frequency, noise, and OCR features
+        
+        **Training:** Multi-round chunked training with 4 sequential rounds
+        **Dataset:** DocTamper (120K) + SCD (18K) + FCD (2K) = 140K samples
+        """
+    )
+    
+    # Event handler
+    analyze_btn.click(
+        fn=detect_forgery,
+        inputs=[input_file],
+        outputs=[output_image, output_json, metrics_plot, report_html]
+    )
+
+# ============================================================================
+# LAUNCH
+# ============================================================================
 
 if __name__ == "__main__":
     demo.launch()