Update app.py

app.py

@@ -1,7 +1,6 @@
 """
-Document Forgery Detection - Gradio Interface for Hugging Face Spaces
-
-This app provides a web interface for detecting and classifying document forgeries.
+Document Forgery Detection – Professional Gradio Dashboard
+Hugging Face Spaces Deployment
 """
 
 import gradio as gr
@@ -9,11 +8,14 @@ import torch
 import cv2
 import numpy as np
 from PIL import Image
-import json
+import plotly.graph_objects as go
 from pathlib import Path
 import sys
+import json
 
-# Add src to path
+# -------------------------------------------------
+# PATH SETUP
+# -------------------------------------------------
 sys.path.insert(0, str(Path(__file__).parent))
 
 from src.models import get_model
@@ -24,253 +26,181 @@ 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
 
-# Class names
-CLASS_NAMES = {0: 'Copy-Move', 1: 'Splicing', 2: 'Generation'}
+# -------------------------------------------------
+# CONSTANTS
+# -------------------------------------------------
+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
+    0: (255, 0, 0),
+    1: (0, 255, 0),
+    2: (0, 0, 255),
 }
 
-
+# -------------------------------------------------
+# FORGERY DETECTOR (UNCHANGED CORE LOGIC)
+# -------------------------------------------------
 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.config = get_config("config.yaml")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
         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
+
         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.classifier.load("models/classifier")
+
+        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!")
-    
+
+        print("✓ Models loaded")
+
     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:
+
+        if image.ndim == 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
+
+        original = image.copy()
+
         preprocessed, _ = self.preprocessor(image, None)
-        
-        # Augment
         augmented = self.augmentation(preprocessed, None)
-        image_tensor = augmented['image'].unsqueeze(0).to(self.device)
-        
-        # Run localization
+        image_tensor = augmented["image"].unsqueeze(0).to(self.device)
+
         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
+
+        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)
+
         results = []
-        for region in regions:
-            # Extract features
+        for r in regions:
             features = self.feature_extractor.extract(
-                preprocessed,
-                region['region_mask'],
-                [f.cpu() for f in decoder_features]
+                preprocessed, r["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
+
+            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:
                 results.append({
-                    'region_id': region['region_id'],
-                    'bounding_box': region['bounding_box'],
-                    'forgery_type': CLASS_NAMES[forgery_type],
-                    'confidence': confidence
+                    "bounding_box": r["bounding_box"],
+                    "forgery_type": CLASS_NAMES[int(pred[0])],
+                    "confidence": float(conf[0]),
                 })
-        
-        # 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%'
-            }
+
+        overlay = self._draw_overlay(original, results)
+
+        return overlay, {
+            "num_detections": len(results),
+            "detections": results,
         }
-        
-        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 _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
 
-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).
-    """
-)
 
+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 (%)"),
+    )
+
+# -------------------------------------------------
+# UI
+# -------------------------------------------------
+with gr.Blocks(theme=gr.themes.Soft(), title="Document Forgery Detection") as demo:
+
+    gr.Markdown("# 📄 Document Forgery Detection System")
+
+    with gr.Row():
+        file_input = gr.File(label="Upload Document (Image/PDF)")
+        detect_btn = gr.Button("Run Detection", variant="primary")
+
+    output_img = gr.Image(label="Forgery Localization Result", type="numpy")
+
+    with gr.Tabs():
+        with gr.Tab("📊 Metrics"):
+            with gr.Row():
+                dice_plot = gr.Plot()
+                acc_plot = gr.Plot()
+                conf_plot = gr.Plot()
+
+        with gr.Tab("🧾 Details"):
+            json_out = gr.JSON()
+
+        with gr.Tab("👥 Team"):
+            gr.Markdown("""
+            **Document Forgery Detection Project**
+
+            - Krishnanandhaa — Model & Training  
+            - Teammate 1 — Feature Engineering  
+            - Teammate 2 — Evaluation  
+            - Teammate 3 — Deployment  
+
+            *Collaborators are added via Hugging Face Space settings.*
+            """)
+
+    detect_btn.click(
+        run_detection,
+        inputs=file_input,
+        outputs=[output_img, json_out, dice_plot, acc_plot, conf_plot]
+    )
 
 if __name__ == "__main__":
     demo.launch()
-