Update app.py

app.py

@@ -1,6 +1,7 @@
 """
-Document Forgery Detection – Professional Gradio Dashboard
-Hugging Face Spaces Deployment
+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
@@ -8,14 +9,11 @@ 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
 
-# -------------------------------------------------
-# PATH SETUP
-# -------------------------------------------------
+# Add src to path
 sys.path.insert(0, str(Path(__file__).parent))
 
 from src.models import get_model
@@ -26,181 +24,253 @@ 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"}
+# Class names
+CLASS_NAMES = {0: 'Copy-Move', 1: 'Splicing', 2: 'Generation'}
 CLASS_COLORS = {
-    0: (255, 0, 0),
-    1: (0, 255, 0),
-    2: (0, 0, 255),
+    0: (255, 0, 0),      # Red for Copy-Move
+    1: (0, 255, 0),      # Green for Splicing
+    2: (0, 0, 255)       # Blue for Generation
 }
 
-# -------------------------------------------------
-# FORGERY DETECTOR (UNCHANGED CORE LOGIC)
-# -------------------------------------------------
+
 class ForgeryDetector:
+    """Main forgery detection pipeline"""
+    
     def __init__(self):
         print("Loading models...")
-
-        self.config = get_config("config.yaml")
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
+        
+        # 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"])
+        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")
-
-        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")
-
+        
+        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)
-
-        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 to 2D array (1, n_features) for classifier
             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:
+            
+            # 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({
-                    "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
                 })
-
-        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
+        json_results = {
+            'num_detections': len(results),
+            'detections': results,
+            'model_info': {
+                'segmentation_dice': '75%',
+                'classifier_accuracy': '92%'
+            }
         }
-
-    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
+        
+        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()
 
 
-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).
+    """
+)
 
-# -------------------------------------------------
-# 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()
+