Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |
| from typing import Dict, List, Tuple | |
| # 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: 'Text Substitution'} | |
| CLASS_COLORS = { | |
| 0: (217, 83, 79), # #d9534f - Muted red | |
| 1: (92, 184, 92), # #5cb85c - Muted green | |
| 2: (65, 105, 225) # #4169E1 - Royal blue | |
| } | |
| # Actual model performance metrics | |
| MODEL_METRICS = { | |
| 'segmentation': { | |
| 'dice': 0.6212, | |
| 'iou': 0.4506, | |
| 'precision': 0.7077, | |
| 'recall': 0.5536 | |
| }, | |
| 'classification': { | |
| 'overall_accuracy': 0.8897, | |
| 'per_class': { | |
| 'copy_move': 0.92, | |
| 'splicing': 0.85, | |
| 'generation': 0.90 | |
| } | |
| } | |
| } | |
| 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: | |
| original_image: Original uploaded image | |
| overlay_image: Image with detection overlay | |
| results_html: Detection results as HTML | |
| """ | |
| # 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) | |
| # 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 | |
| if features.ndim == 1: | |
| features = features.reshape(1, -1) | |
| # Pad/truncate features to match classifier | |
| 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({ | |
| '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 HTML response | |
| results_html = self._create_html_report(results) | |
| return original_image, overlay, results_html | |
| def _create_overlay(self, image, results): | |
| """Create overlay visualization""" | |
| overlay = image.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 = 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%}" | |
| font = cv2.FONT_HERSHEY_SIMPLEX | |
| font_scale = 0.5 | |
| thickness = 1 | |
| (label_w, label_h), baseline = cv2.getTextSize(label, font, font_scale, thickness) | |
| cv2.rectangle(overlay, (x, y-label_h-8), (x+label_w+4, y), color, -1) | |
| cv2.putText(overlay, label, (x+2, y-4), font, font_scale, (255, 255, 255), thickness) | |
| return overlay | |
| def _create_html_report(self, results): | |
| """Create HTML report with detection results""" | |
| num_detections = len(results) | |
| if num_detections == 0: | |
| return """ | |
| <div style='padding:12px; border:1px solid #5cb85c; border-radius:8px;'> | |
| ✓ <b>No forgery detected.</b><br> | |
| The document appears to be authentic. | |
| </div> | |
| """ | |
| # Calculate statistics | |
| avg_confidence = sum(r['confidence'] for r in results) / num_detections | |
| type_counts = {} | |
| for r in results: | |
| ft = r['forgery_type'] | |
| type_counts[ft] = type_counts.get(ft, 0) + 1 | |
| html = f""" | |
| <div style='padding:12px; border:1px solid #d9534f; border-radius:8px;'> | |
| <b>⚠️ Forgery Detected</b><br><br> | |
| <b>Summary:</b><br> | |
| • Regions detected: {num_detections}<br> | |
| • Average confidence: {avg_confidence*100:.1f}%<br><br> | |
| <b>Model Performance:</b><br> | |
| • Segmentation Dice: {MODEL_METRICS['segmentation']['dice']*100:.1f}%<br> | |
| • Classification Accuracy: {MODEL_METRICS['classification']['overall_accuracy']*100:.1f}%<br><br> | |
| <b>Detections:</b><br> | |
| """ | |
| for i, result in enumerate(results, 1): | |
| forgery_type = result['forgery_type'] | |
| confidence = result['confidence'] | |
| bbox = result['bounding_box'] | |
| forgery_id = [k for k, v in CLASS_NAMES.items() if v == forgery_type][0] | |
| color_rgb = CLASS_COLORS[forgery_id] | |
| color_hex = f"#{color_rgb[0]:02x}{color_rgb[1]:02x}{color_rgb[2]:02x}" | |
| html += f""" | |
| <div style='margin:8px 0; padding:8px; border-left:3px solid {color_hex}; background:#f9f9f9;'> | |
| <b>Region {i}:</b> {forgery_type} ({confidence*100:.1f}%)<br> | |
| <small>Location: ({bbox[0]}, {bbox[1]}) | Size: {bbox[2]}×{bbox[3]}px</small> | |
| </div> | |
| """ | |
| html += """ | |
| </div> | |
| """ | |
| return html | |
| # Initialize detector | |
| detector = ForgeryDetector() | |
| def detect_forgery(file): | |
| """Gradio interface function""" | |
| try: | |
| if file is None: | |
| return None, None, "<div style='padding:12px; border:1px solid #d9534f; border-radius:8px;'>❌ <b>No file uploaded.</b></div>" | |
| # Get file path | |
| file_path = file.name if hasattr(file, 'name') else file | |
| # Check if PDF | |
| if file_path.lower().endswith('.pdf'): | |
| original, overlay, results_html = detector.detect(file_path) | |
| else: | |
| image = Image.open(file_path) | |
| original, overlay, results_html = detector.detect(image) | |
| return original, overlay, results_html | |
| except Exception as e: | |
| import traceback | |
| error_details = traceback.format_exc() | |
| print(f"Error: {error_details}") | |
| error_html = f""" | |
| <div style='padding:12px; border:1px solid #d9534f; border-radius:8px;'> | |
| ❌ <b>Error:</b> {str(e)} | |
| </div> | |
| """ | |
| return None, None, error_html | |
| # Custom CSS - subtle styling | |
| custom_css = """ | |
| .predict-btn { | |
| background-color: #4169E1 !important; | |
| color: white !important; | |
| } | |
| .clear-btn { | |
| background-color: #6A89A7 !important; | |
| color: white !important; | |
| } | |
| """ | |
| # Create Gradio interface | |
| with gr.Blocks(css=custom_css) as demo: | |
| gr.Markdown( | |
| """ | |
| # 📄 Document Forgery Detection | |
| Upload a document image or PDF to detect and classify forgeries. | |
| """ | |
| ) | |
| 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" | |
| ) | |
| with gr.Row(): | |
| clear_btn = gr.Button("🧹 Clear", elem_classes="clear-btn") | |
| analyze_btn = gr.Button("🔍 Analyze", elem_classes="predict-btn") | |
| gr.Markdown( | |
| """ | |
| **Supported formats:** | |
| - Images: JPG, PNG, BMP, TIFF, WebP | |
| - PDF: First page analyzed | |
| **Forgery types:** | |
| - Copy-Move: Duplicated regions | |
| - Splicing: Mixed sources | |
| - Text Substitution: Modified text | |
| """ | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Original Image") | |
| original_image = gr.Image(label="Uploaded Document", type="numpy") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Detection Result") | |
| output_image = gr.Image(label="Annotated Document", type="numpy") | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Analysis Report") | |
| output_html = gr.HTML( | |
| value="<i>No analysis yet. Upload a document and click Analyze.</i>" | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| **Model Architecture:** | |
| - **Localization:** MobileNetV3-Small + UNet (Dice: 62.1%, IoU: 45.1%) | |
| - **Classification:** LightGBM with 526 features (Accuracy: 88.97%) | |
| - **Training:** 140K samples (DocTamper + SCD + FCD datasets) | |
| """ | |
| ) | |
| # Event handlers | |
| analyze_btn.click( | |
| fn=detect_forgery, | |
| inputs=[input_file], | |
| outputs=[original_image, output_image, output_html] | |
| ) | |
| clear_btn.click( | |
| fn=lambda: (None, None, None, "<i>No analysis yet. Upload a document and click Analyze.</i>"), | |
| inputs=None, | |
| outputs=[input_file, original_image, output_image, output_html] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |