DocForg/Document_Forgery_Detection
0
1"""
2Complete Document Forgery Detection Pipeline
3Implements Full Algorithm Steps 1-11
4
5Features:
6- ✅ Localization (WHERE is forgery?)
7- ✅ Classification (WHAT type of forgery?)
8- ✅ Confidence filtering
9- ✅ Visualizations (heatmaps, overlays, bounding boxes)
10- ✅ JSON output with detailed results
11- ✅ Actual vs Predicted comparison (if ground truth available)
12
13Usage:
14 python scripts/inference_pipeline.py --image path/to/document.jpg
15 python scripts/inference_pipeline.py --image path/to/document.jpg --ground_truth path/to/mask.png
16"""
17
18import sys
19from pathlib import Path
20import argparse
21import numpy as np
22import cv2
23import torch
24import json
25from datetime import datetime
26import matplotlib.pyplot as plt
27import matplotlib.patches as patches
28
29sys.path.insert(0, str(Path(__file__).parent.parent))
30
31from src.config import get_config
32from src.models import get_model
33from src.features import get_feature_extractor, get_mask_refiner, get_region_extractor
34from src.training.classifier import ForgeryClassifier
35from src.data.preprocessing import DocumentPreprocessor
36
37# Class mapping
38CLASS_NAMES = {
39 0: 'Copy-Move',
40 1: 'Splicing',
41 2: 'Generation'
42}
43
44CLASS_COLORS = {
45 0: (255, 0, 0), # Red for Copy-Move
46 1: (0, 255, 0), # Green for Splicing
47 2: (0, 0, 255) # Blue for Generation
48}
49
50
51class ForgeryDetectionPipeline:
52 """
53 Complete forgery detection pipeline
54 Implements Algorithm Steps 1-11
55 """
56
57 def __init__(self, config_path='config.yaml'):
58 """Initialize pipeline with models"""
59 print("="*70)
60 print("Initializing Forgery Detection Pipeline")
61 print("="*70)
62
63 self.config = get_config(config_path)
64 self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
65
66 # Load localization model (Steps 1-6)
67 print("\n1. Loading localization model...")
68 self.localization_model = get_model(self.config).to(self.device)
69 checkpoint = torch.load('outputs/checkpoints/best_doctamper.pth',
70 map_location=self.device)
71 self.localization_model.load_state_dict(checkpoint['model_state_dict'])
72 self.localization_model.eval()
73 print(f" ✓ Loaded (Val Dice: {checkpoint.get('best_metric', 0):.2%})")
74
75 # Load classifier (Step 8)
76 print("\n2. Loading forgery type classifier...")
77 self.classifier = ForgeryClassifier(self.config)
78 self.classifier.load('outputs/classifier')
79 print(" ✓ Loaded")
80
81 # Initialize components
82 print("\n3. Initializing components...")
83 self.preprocessor = DocumentPreprocessor(self.config, 'doctamper')
84
85 # Initialize augmentation for inference
86 from src.data.augmentation import DatasetAwareAugmentation
87 self.augmentation = DatasetAwareAugmentation(self.config, 'doctamper', is_training=False)
88
89 self.feature_extractor = get_feature_extractor(self.config, is_text_document=True)
90 self.mask_refiner = get_mask_refiner(self.config)
91 self.region_extractor = get_region_extractor(self.config)
92 print(" ✓ Ready")
93
94 print("\n" + "="*70)
95 print("Pipeline Initialized Successfully!")
96 print("="*70 + "\n")
97
98 def detect(self, image_path, ground_truth_path=None, output_dir='outputs/inference'):
99 """
100 Run complete detection pipeline
101
102 Args:
103 image_path: Path to input document image
104 ground_truth_path: Optional path to ground truth mask
105 output_dir: Directory to save outputs
106
107 Returns:
108 results: Dictionary with detection results
109 """
110 print(f"\n{'='*70}")
111 print(f"Processing: {image_path}")
112 print(f"{'='*70}\n")
113
114 # Create output directory
115 output_path = Path(output_dir)
116 output_path.mkdir(parents=True, exist_ok=True)
117
118 # Get base filename
119 base_name = Path(image_path).stem
120
121 # Step 1-2: Load and preprocess image (EXACTLY like dataset)
122 print("Step 1-2: Loading and preprocessing...")
123 image = cv2.imread(str(image_path))
124 if image is None:
125 raise ValueError(f"Could not load image: {image_path}")
126
127 image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
128
129 # Create dummy mask for preprocessing
130 dummy_mask = np.zeros(image_rgb.shape[:2], dtype=np.uint8)
131
132 # Step 1: Preprocess (like dataset line: image, mask = self.preprocessor(image, mask))
133 preprocessed_img, preprocessed_mask = self.preprocessor(image_rgb, dummy_mask)
134
135 # Step 2: Augment (like dataset line: augmented = self.augmentation(image, mask))
136 augmented = self.augmentation(preprocessed_img, preprocessed_mask)
137
138 # Step 3: Extract tensor (like dataset line: image = augmented['image'])
139 image_tensor = augmented['image']
140
141 print(f" ✓ Image shape: {image_rgb.shape}")
142 print(f" ✓ Preprocessed tensor shape: {image_tensor.shape}")
143 print(f" ✓ Tensor range: [{image_tensor.min():.4f}, {image_tensor.max():.4f}]")
144
145 # Load ground truth if provided
146 ground_truth = None
147 if ground_truth_path:
148 ground_truth = cv2.imread(str(ground_truth_path), cv2.IMREAD_GRAYSCALE)
149 if ground_truth is not None:
150 # Resize to match preprocessed size
151 target_size = (image_tensor.shape[2], image_tensor.shape[1]) # (W, H)
152 ground_truth = cv2.resize(ground_truth, target_size)
153 print(f" ✓ Ground truth loaded")
154
155 # Step 3-4: Localization (WHERE is forgery?)
156 print("\nStep 3-4: Forgery localization...")
157 image_batch = image_tensor.unsqueeze(0).to(self.device)
158
159 with torch.no_grad():
160 logits, decoder_features = self.localization_model(image_batch)
161 prob_map = torch.sigmoid(logits).cpu().numpy()[0, 0]
162
163 print(f" ✓ Probability map generated")
164 print(f" ✓ Prob map range: [{prob_map.min():.4f}, {prob_map.max():.4f}]")
165
166 # Step 5: Binary mask generation
167 print("\nStep 5: Generating binary mask...")
168 binary_mask = (prob_map > 0.5).astype(np.uint8)
169 refined_mask = self.mask_refiner.refine(binary_mask)
170 print(f" ✓ Mask refined")
171
172 # Step 6: Region extraction
173 print("\nStep 6: Extracting forgery regions...")
174 # Convert tensor to numpy for region extraction and feature extraction
175 preprocessed_numpy = image_tensor.permute(1, 2, 0).cpu().numpy()
176 regions = self.region_extractor.extract(refined_mask, prob_map, preprocessed_numpy)
177 print(f" ✓ Found {len(regions)} regions")
178
179 if len(regions) == 0:
180 print("\n⚠ No forgery regions detected!")
181 # Still create visualizations if ground truth exists
182 if ground_truth is not None:
183 print("\nCreating comparison with ground truth...")
184 self._create_comparison_visualization(
185 image_rgb, prob_map, refined_mask, ground_truth,
186 base_name, output_path
187 )
188 return self._create_clean_result(image_rgb, base_name, output_path, ground_truth)
189
190 # Step 7-8: Feature extraction and classification
191 print("\nStep 7-8: Classifying forgery types...")
192 region_results = []
193
194 for i, region in enumerate(regions):
195 # Extract features (Step 7)
196 features = self.feature_extractor.extract(
197 preprocessed_numpy,
198 region['region_mask'],
199 [f.cpu() for f in decoder_features]
200 )
201
202 # Ensure correct dimension (526)
203 expected_dim = 526
204 if len(features) < expected_dim:
205 features = np.pad(features, (0, expected_dim - len(features)))
206 elif len(features) > expected_dim:
207 features = features[:expected_dim]
208
209 features = features.reshape(1, -1)
210
211 # Classify (Step 8)
212 predictions, confidences = self.classifier.predict(features)
213 forgery_type = int(predictions[0])
214 confidence = float(confidences[0])
215
216 region_results.append({
217 'region_id': i + 1,
218 'bounding_box': region['bounding_box'],
219 'area': int(region['area']),
220 'forgery_type': CLASS_NAMES[forgery_type],
221 'forgery_type_id': forgery_type,
222 'confidence': confidence,
223 'mask_probability_mean': float(prob_map[region['region_mask'] > 0].mean())
224 })
225
226 print(f" Region {i+1}: {CLASS_NAMES[forgery_type]} "
227 f"(confidence: {confidence:.2%})")
228
229 # Step 9: False positive removal
230 print("\nStep 9: Filtering low-confidence regions...")
231 confidence_threshold = self.config.get('classification.confidence_threshold', 0.6)
232 filtered_results = [r for r in region_results if r['confidence'] >= confidence_threshold]
233 print(f" ✓ Kept {len(filtered_results)}/{len(region_results)} regions "
234 f"(threshold: {confidence_threshold:.0%})")
235
236 # Step 10-11: Generate outputs
237 print("\nStep 10-11: Generating outputs...")
238
239 # Calculate scale factors for coordinate conversion
240 # Bounding boxes are in preprocessed coordinates (384x384)
241 # Need to scale to original image coordinates
242 orig_h, orig_w = image_rgb.shape[:2]
243 prep_h, prep_w = prob_map.shape
244 scale_x = orig_w / prep_w
245 scale_y = orig_h / prep_h
246
247 # Create visualizations
248 self._create_visualizations(
249 image_rgb, prob_map, refined_mask, filtered_results,
250 ground_truth, base_name, output_path, scale_x, scale_y
251 )
252
253 # Create JSON output
254 results = self._create_json_output(
255 image_path, filtered_results, ground_truth, base_name, output_path
256 )
257
258 print(f"\n{'='*70}")
259 print("✅ Detection Complete!")
260 print(f"{'='*70}")
261 print(f"Output directory: {output_path}")
262 print(f"Detected {len(filtered_results)} forgery regions")
263 print(f"{'='*70}\n")
264
265 return results
266
267 def _create_visualizations(self, image, prob_map, mask, results,
268 ground_truth, base_name, output_path, scale_x, scale_y):
269 """Create all visualizations"""
270
271 # 1. Probability heatmap
272 plt.figure(figsize=(15, 5))
273
274 plt.subplot(1, 3, 1)
275 plt.imshow(image)
276 plt.title('Original Document')
277 plt.axis('off')
278
279 plt.subplot(1, 3, 2)
280 plt.imshow(prob_map, cmap='hot', vmin=0, vmax=1)
281 plt.colorbar(label='Forgery Probability')
282 plt.title('Probability Heatmap')
283 plt.axis('off')
284
285 plt.subplot(1, 3, 3)
286 plt.imshow(mask, cmap='gray')
287 plt.title('Binary Mask')
288 plt.axis('off')
289
290 plt.tight_layout()
291 plt.savefig(output_path / f'{base_name}_heatmap.png', dpi=150, bbox_inches='tight')
292 plt.close()
293 print(f" ✓ Saved heatmap")
294
295 # 2. Overlay with bounding boxes and labels
296 overlay = image.copy()
297 alpha = 0.4
298
299 # Create colored mask overlay (scale mask to original size)
300 mask_scaled = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
301 colored_mask = np.zeros_like(image)
302
303 for result in results:
304 bbox = result['bounding_box']
305 forgery_type = result['forgery_type_id']
306 color = CLASS_COLORS[forgery_type]
307
308 # Scale bounding box to original image coordinates
309 x, y, w, h = bbox
310 x_scaled = int(x * scale_x)
311 y_scaled = int(y * scale_y)
312 w_scaled = int(w * scale_x)
313 h_scaled = int(h * scale_y)
314
315 # Color the region
316 colored_mask[y_scaled:y_scaled+h_scaled, x_scaled:x_scaled+w_scaled] = color
317
318 # Blend with original
319 overlay = cv2.addWeighted(overlay, 1-alpha, colored_mask, alpha, 0)
320
321 # Draw bounding boxes and labels
322 fig, ax = plt.subplots(1, figsize=(12, 8))
323 ax.imshow(overlay)
324
325 for result in results:
326 bbox = result['bounding_box']
327 x, y, w, h = bbox # bbox is [x, y, w, h] in preprocessed coordinates
328
329 # Scale to original image coordinates
330 x_scaled = x * scale_x
331 y_scaled = y * scale_y
332 w_scaled = w * scale_x
333 h_scaled = h * scale_y
334
335 forgery_type = result['forgery_type']
336 confidence = result['confidence']
337 color_rgb = tuple(c/255 for c in CLASS_COLORS[result['forgery_type_id']])
338
339 # Draw rectangle
340 rect = patches.Rectangle((x_scaled, y_scaled), w_scaled, h_scaled,
341 linewidth=2, edgecolor=color_rgb,
342 facecolor='none')
343 ax.add_patch(rect)
344
345 # Add label
346 label = f"{forgery_type}\n{confidence:.1%}"
347 ax.text(x_scaled, y_scaled-10, label, color='white', fontsize=10,
348 bbox=dict(boxstyle='round', facecolor=color_rgb, alpha=0.8))
349
350 ax.axis('off')
351 ax.set_title('Forgery Detection Results', fontsize=14, fontweight='bold')
352 plt.tight_layout()
353 plt.savefig(output_path / f'{base_name}_overlay.png', dpi=150, bbox_inches='tight')
354 plt.close()
355 print(f" ✓ Saved overlay")
356
357 # 3. Comparison with ground truth (if available)
358 if ground_truth is not None:
359 fig, axes = plt.subplots(1, 3, figsize=(18, 6))
360
361 axes[0].imshow(image)
362 axes[0].set_title('Original Document', fontsize=12)
363 axes[0].axis('off')
364
365 axes[1].imshow(ground_truth, cmap='gray')
366 axes[1].set_title('Ground Truth', fontsize=12)
367 axes[1].axis('off')
368
369 axes[2].imshow(mask, cmap='gray')
370 axes[2].set_title('Predicted Mask', fontsize=12)
371 axes[2].axis('off')
372
373 # Calculate metrics
374 intersection = np.logical_and(ground_truth > 127, mask > 0).sum()
375 union = np.logical_or(ground_truth > 127, mask > 0).sum()
376 iou = intersection / (union + 1e-8)
377 dice = 2 * intersection / (ground_truth.sum() + mask.sum() + 1e-8)
378
379 fig.suptitle(f'Actual vs Predicted (IoU: {iou:.2%}, Dice: {dice:.2%})',
380 fontsize=14, fontweight='bold')
381
382 plt.tight_layout()
383 plt.savefig(output_path / f'{base_name}_comparison.png', dpi=150, bbox_inches='tight')
384 plt.close()
385 print(f" ✓ Saved comparison (IoU: {iou:.2%}, Dice: {dice:.2%})")
386
387 # 4. Per-region visualization
388 if len(results) > 0:
389 n_regions = len(results)
390 cols = min(4, n_regions)
391 rows = (n_regions + cols - 1) // cols
392
393 fig, axes = plt.subplots(rows, cols, figsize=(4*cols, 4*rows))
394 if n_regions == 1:
395 axes = [axes]
396 else:
397 axes = axes.flatten()
398
399 for i, result in enumerate(results):
400 bbox = result['bounding_box']
401 x, y, w, h = bbox # bbox is [x, y, w, h] in preprocessed coordinates
402
403 # Scale to original image coordinates
404 x_scaled = int(x * scale_x)
405 y_scaled = int(y * scale_y)
406 w_scaled = int(w * scale_x)
407 h_scaled = int(h * scale_y)
408
409 region_img = image[y_scaled:y_scaled+h_scaled, x_scaled:x_scaled+w_scaled]
410
411 axes[i].imshow(region_img)
412 axes[i].set_title(f"Region {i+1}: {result['forgery_type']}\n"
413 f"Confidence: {result['confidence']:.1%}",
414 fontsize=10)
415 axes[i].axis('off')
416
417 # Hide unused subplots
418 for i in range(n_regions, len(axes)):
419 axes[i].axis('off')
420
421 plt.tight_layout()
422 plt.savefig(output_path / f'{base_name}_regions.png', dpi=150, bbox_inches='tight')
423 plt.close()
424 print(f" ✓ Saved region details")
425
426 def _create_json_output(self, image_path, results, ground_truth, base_name, output_path):
427 """Create JSON output with results"""
428
429 output = {
430 'image_path': str(image_path),
431 'timestamp': datetime.now().isoformat(),
432 'num_regions_detected': len(results),
433 'regions': results
434 }
435
436 # Add ground truth comparison if available
437 if ground_truth is not None:
438 output['has_ground_truth'] = True
439
440 # Save JSON
441 json_path = output_path / f'{base_name}_results.json'
442 with open(json_path, 'w') as f:
443 json.dump(output, f, indent=2)
444
445 print(f" ✓ Saved JSON results")
446
447 return output
448
449 def _create_comparison_visualization(self, image, prob_map, mask, ground_truth,
450 base_name, output_path):
451 """Create comparison visualization between actual and predicted"""
452
453 fig, axes = plt.subplots(2, 2, figsize=(16, 12))
454
455 # Original image
456 axes[0, 0].imshow(image)
457 axes[0, 0].set_title('Original Document', fontsize=14, fontweight='bold')
458 axes[0, 0].axis('off')
459
460 # Ground truth
461 axes[0, 1].imshow(ground_truth, cmap='gray')
462 axes[0, 1].set_title('Ground Truth (Actual)', fontsize=14, fontweight='bold')
463 axes[0, 1].axis('off')
464
465 # Predicted mask
466 axes[1, 0].imshow(mask, cmap='gray')
467 axes[1, 0].set_title('Predicted Mask', fontsize=14, fontweight='bold')
468 axes[1, 0].axis('off')
469
470 # Probability heatmap
471 im = axes[1, 1].imshow(prob_map, cmap='hot', vmin=0, vmax=1)
472 axes[1, 1].set_title('Probability Heatmap', fontsize=14, fontweight='bold')
473 axes[1, 1].axis('off')
474 plt.colorbar(im, ax=axes[1, 1], fraction=0.046, pad=0.04)
475
476 # Calculate metrics
477 intersection = np.logical_and(ground_truth > 127, mask > 0).sum()
478 union = np.logical_or(ground_truth > 127, mask > 0).sum()
479 gt_sum = (ground_truth > 127).sum()
480 pred_sum = (mask > 0).sum()
481
482 iou = intersection / (union + 1e-8)
483 dice = 2 * intersection / (gt_sum + pred_sum + 1e-8)
484 precision = intersection / (pred_sum + 1e-8) if pred_sum > 0 else 0
485 recall = intersection / (gt_sum + 1e-8) if gt_sum > 0 else 0
486
487 fig.suptitle(f'Actual vs Predicted Comparison\n'
488 f'IoU: {iou:.2%} | Dice: {dice:.2%} | '
489 f'Precision: {precision:.2%} | Recall: {recall:.2%}',
490 fontsize=16, fontweight='bold')
491
492 plt.tight_layout()
493 plt.savefig(output_path / f'{base_name}_comparison.png', dpi=150, bbox_inches='tight')
494 plt.close()
495 print(f" ✓ Saved comparison (IoU: {iou:.2%}, Dice: {dice:.2%})")
496
497 def _create_clean_result(self, image, base_name, output_path, ground_truth=None):
498 """Create result for clean (no forgery) document"""
499
500 # Save original image
501 plt.figure(figsize=(10, 8))
502 plt.imshow(image)
503 plt.title('No Forgery Detected', fontsize=14, fontweight='bold', color='green')
504 plt.axis('off')
505 plt.tight_layout()
506 plt.savefig(output_path / f'{base_name}_clean.png', dpi=150, bbox_inches='tight')
507 plt.close()
508
509 # Create JSON
510 output = {
511 'timestamp': datetime.now().isoformat(),
512 'num_regions_detected': 0,
513 'regions': [],
514 'status': 'clean'
515 }
516
517 json_path = output_path / f'{base_name}_results.json'
518 with open(json_path, 'w') as f:
519 json.dump(output, f, indent=2)
520
521 return output
522
523
524def main():
525 parser = argparse.ArgumentParser(description='Document Forgery Detection Pipeline')
526 parser.add_argument('--image', type=str, required=True,
527 help='Path to input document image')
528 parser.add_argument('--ground_truth', type=str, default=None,
529 help='Path to ground truth mask (optional)')
530 parser.add_argument('--output_dir', type=str, default='outputs/inference',
531 help='Output directory for results')
532 parser.add_argument('--config', type=str, default='config.yaml',
533 help='Path to config file')
534
535 args = parser.parse_args()
536
537 # Initialize pipeline
538 pipeline = ForgeryDetectionPipeline(args.config)
539
540 # Run detection
541 results = pipeline.detect(
542 args.image,
543 ground_truth_path=args.ground_truth,
544 output_dir=args.output_dir
545 )
546
547 # Print summary
548 print("\nDetection Summary:")
549 print(f" Regions detected: {results['num_regions_detected']}")
550 if results['num_regions_detected'] > 0:
551 for region in results['regions']:
552 print(f" - {region['forgery_type']}: {region['confidence']:.1%} confidence")
553
554
555if __name__ == '__main__':
556 main()
557