DocForg/Document_Forgery_Detection
0
1"""
2LightGBM Classifier Training - DocTamper with Tampering Labels
3FIXED VERSION with proper checkpointing and feature dimension handling
4
5Implements Algorithm Steps 7-8:
6 7. Hybrid Feature Extraction
7 8. Region-wise Forgery Classification
8
9Uses:
10- Localization: best_doctamper.pth (Steps 1-6 complete)
11- Training: DocTamper TrainingSet + tampering/DocTamperV1-TrainingSet.pk
12- Testing: DocTamper TestingSet + tampering/DocTamperV1-TestingSet.pk
13- Classes: Copy-Move (CM), Splicing (SP), Generation (GE)
14
15Features:
16- ✅ Checkpoint saving every 1000 samples
17- ✅ Resume from checkpoint if interrupted
18- ✅ Fixed feature dimension mismatch
19- ✅ Robust error handling
20
21Usage:
22 python scripts/train_classifier_doctamper_fixed.py
23"""
24
25import sys
26from pathlib import Path
27import numpy as np
28import pickle
29import lmdb
30import cv2
31import torch
32from tqdm import tqdm
33import json
34
35sys.path.insert(0, str(Path(__file__).parent.parent))
36
37from src.config import get_config
38from src.models import get_model
39from src.features import get_feature_extractor
40from src.training.classifier import get_classifier
41
42# Configuration
43MODEL_PATH = 'outputs/checkpoints/best_doctamper.pth'
44OUTPUT_DIR = 'outputs/classifier'
45MAX_SAMPLES = 999999 # Use all available samples
46
47# Label mapping (Algorithm Step 8.2) - 3 classes
48LABEL_MAP = {
49 'CM': 0, # Copy-Move
50 'SP': 1, # Splicing
51 'GE': 2, # Generation (AI-generated, separate from Splicing)
52}
53
54
55def load_tampering_labels(label_file):
56 """Load forgery type labels from tampering folder"""
57 with open(label_file, 'rb') as f:
58 labels = pickle.load(f)
59
60 print(f"Loaded {len(labels)} labels from {label_file}")
61 return labels
62
63
64def load_sample_from_lmdb(lmdb_env, index):
65 """Load image and mask from LMDB"""
66 txn = lmdb_env.begin()
67
68 # Get image
69 img_key = f'image-{index:09d}'.encode('utf-8')
70 img_data = txn.get(img_key)
71 if not img_data:
72 return None, None
73
74 # Get mask (DocTamper uses 'label-' not 'mask-')
75 mask_key = f'label-{index:09d}'.encode('utf-8')
76 mask_data = txn.get(mask_key)
77 if not mask_data:
78 return None, None
79
80 # Decode
81 img_array = np.frombuffer(img_data, dtype=np.uint8)
82 image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
83 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
84
85 mask_array = np.frombuffer(mask_data, dtype=np.uint8)
86 mask = cv2.imdecode(mask_array, cv2.IMREAD_GRAYSCALE)
87
88 return image, mask
89
90
91def extract_features(config, model, lmdb_path, tampering_labels,
92 max_samples, device, split_name):
93 """
94 Extract hybrid features with checkpointing and resume capability
95 """
96
97 print(f"\n{'='*60}")
98 print(f"Extracting features from {split_name}")
99 print(f"{'='*60}")
100
101 # Setup checkpoint directory
102 checkpoint_dir = Path(OUTPUT_DIR)
103 checkpoint_dir.mkdir(parents=True, exist_ok=True)
104
105 # Check for existing checkpoint to resume
106 checkpoints = list(checkpoint_dir.glob(f'checkpoint_{split_name}_*.npz'))
107 if checkpoints:
108 latest_checkpoint = max(checkpoints, key=lambda p: int(p.stem.split('_')[-1]))
109 print(f"✓ Found checkpoint: {latest_checkpoint.name}")
110
111 data = np.load(latest_checkpoint, allow_pickle=True)
112 all_features = data['features'].tolist()
113 all_labels = data['labels'].tolist()
114 expected_dim = int(data['feature_dim'])
115 start_idx = len(all_features)
116
117 print(f"✓ Resuming from sample {start_idx}, feature_dim={expected_dim}")
118 else:
119 all_features = []
120 all_labels = []
121 expected_dim = None
122 start_idx = 0
123
124 # Open LMDB
125 env = lmdb.open(lmdb_path, readonly=True, lock=False)
126
127 # Initialize feature extractor
128 feature_extractor = get_feature_extractor(config, is_text_document=True)
129
130 # Process samples
131 num_processed = start_idx
132 dim_mismatch_count = 0
133
134 for i in tqdm(range(start_idx, min(len(tampering_labels), max_samples)),
135 desc=f"Processing {split_name}", initial=start_idx,
136 total=min(len(tampering_labels), max_samples)):
137 try:
138 # Skip if no label
139 if i not in tampering_labels:
140 continue
141
142 # Get forgery type label
143 forgery_type = tampering_labels[i]
144 if forgery_type not in LABEL_MAP:
145 continue
146
147 label = LABEL_MAP[forgery_type]
148
149 # Load image and mask
150 image, mask = load_sample_from_lmdb(env, i)
151 if image is None or mask is None:
152 continue
153
154 # Skip if no forgery
155 if mask.max() == 0:
156 continue
157
158 # Prepare for model
159 image_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
160 image_tensor = image_tensor.unsqueeze(0).to(device)
161
162 # Get deep features from localization model
163 with torch.no_grad():
164 logits, decoder_features = model(image_tensor)
165
166 # Use ground truth mask for feature extraction
167 mask_binary = (mask > 127).astype(np.uint8)
168
169 # Extract hybrid features
170 features = feature_extractor.extract(
171 image / 255.0,
172 mask_binary,
173 [f.cpu() for f in decoder_features]
174 )
175
176 # Set expected dimension from first valid sample
177 if expected_dim is None:
178 expected_dim = len(features)
179 print(f"\n✓ Feature dimension set to: {expected_dim}")
180
181 # Ensure consistent feature dimension
182 if len(features) != expected_dim:
183 if len(features) < expected_dim:
184 features = np.pad(features, (0, expected_dim - len(features)), mode='constant')
185 else:
186 features = features[:expected_dim]
187 dim_mismatch_count += 1
188
189 all_features.append(features)
190 all_labels.append(label)
191 num_processed += 1
192
193 # Save checkpoint every 10,000 samples (only 12 checkpoints total)
194 if num_processed % 10000 == 0:
195 checkpoint_path = checkpoint_dir / f'checkpoint_{split_name}_{num_processed}.npz'
196 features_array = np.array(all_features, dtype=np.float32)
197 labels_array = np.array(all_labels, dtype=np.int32)
198
199 np.savez_compressed(checkpoint_path,
200 features=features_array,
201 labels=labels_array,
202 feature_dim=expected_dim)
203 print(f"\n✓ Checkpoint: {num_processed} samples (dim={expected_dim}, mismatches={dim_mismatch_count})")
204
205 # Delete old checkpoints to save space (keep only last 2)
206 old_checkpoints = sorted(checkpoint_dir.glob(f'checkpoint_{split_name}_*.npz'))
207 if len(old_checkpoints) > 2:
208 for old_cp in old_checkpoints[:-2]:
209 old_cp.unlink()
210 print(f" Cleaned up: {old_cp.name}")
211
212 except Exception as e:
213 print(f"\n⚠ Error at sample {i}: {str(e)[:80]}")
214 continue
215
216 env.close()
217
218 print(f"\n✓ Extracted {num_processed} samples")
219 if dim_mismatch_count > 0:
220 print(f"⚠ Fixed {dim_mismatch_count} dimension mismatches")
221
222 # Save final features
223 final_path = checkpoint_dir / f'features_{split_name}_final.npz'
224 if len(all_features) > 0:
225 features_array = np.array(all_features, dtype=np.float32)
226 labels_array = np.array(all_labels, dtype=np.int32)
227
228 np.savez_compressed(final_path,
229 features=features_array,
230 labels=labels_array,
231 feature_dim=expected_dim)
232 print(f"✓ Final features saved: {final_path}")
233 print(f" Shape: features={features_array.shape}, labels={labels_array.shape}")
234
235 return features_array, labels_array
236
237 return None, None
238
239
240def main():
241 config = get_config('config.yaml')
242 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
243
244 print("\n" + "="*60)
245 print("LightGBM Classifier Training - DocTamper (FIXED)")
246 print("Implements Algorithm Steps 7-8")
247 print("="*60)
248 print(f"Model: {MODEL_PATH}")
249 print(f"Device: {device}")
250 print(f"Max samples: {MAX_SAMPLES}")
251 print("="*60)
252 print("\nForgery Type Classes (Step 8.2):")
253 print(" 0: Copy-Move (CM)")
254 print(" 1: Splicing (SP)")
255 print(" 2: Generation (GE)")
256 print("="*60)
257
258 # Load localization model
259 print("\nLoading localization model...")
260 model = get_model(config).to(device)
261 checkpoint = torch.load(MODEL_PATH, map_location=device)
262 model.load_state_dict(checkpoint['model_state_dict'])
263 model.eval()
264 print(f"✓ Model loaded (Val Dice: {checkpoint.get('best_metric', 0):.4f})")
265
266 # Load tampering labels
267 train_labels = load_tampering_labels(
268 'datasets/DocTamper/tampering/DocTamperV1-TrainingSet.pk'
269 )
270 test_labels = load_tampering_labels(
271 'datasets/DocTamper/tampering/DocTamperV1-TestingSet.pk'
272 )
273
274 # Extract features from TrainingSet
275 X_train, y_train = extract_features(
276 config, model,
277 'datasets/DocTamper/DocTamperV1-TrainingSet',
278 train_labels,
279 MAX_SAMPLES,
280 device,
281 'TrainingSet'
282 )
283
284 # Extract features from TestingSet
285 X_test, y_test = extract_features(
286 config, model,
287 'datasets/DocTamper/DocTamperV1-TestingSet',
288 test_labels,
289 MAX_SAMPLES // 4,
290 device,
291 'TestingSet'
292 )
293
294 if X_train is None or X_test is None:
295 print("\n❌ No features extracted!")
296 return
297
298 # Summary
299 print("\n" + "="*60)
300 print("Dataset Summary")
301 print("="*60)
302 print(f"Training samples: {len(X_train):,}")
303 print(f"Testing samples: {len(X_test):,}")
304 print(f"Feature dimension: {X_train.shape[1]}")
305
306 print(f"\nTraining class distribution:")
307 train_counts = np.bincount(y_train)
308 class_names = ['Copy-Move', 'Splicing', 'Generation']
309 for i, count in enumerate(train_counts):
310 if i < len(class_names):
311 print(f" {class_names[i]}: {count:,} ({count/len(y_train)*100:.1f}%)")
312
313 print(f"\nTesting class distribution:")
314 test_counts = np.bincount(y_test)
315 for i, count in enumerate(test_counts):
316 if i < len(class_names):
317 print(f" {class_names[i]}: {count:,} ({count/len(y_test)*100:.1f}%)")
318
319 # Train classifier
320 print("\n" + "="*60)
321 print("Training LightGBM Classifier (Step 8.1)")
322 print("="*60)
323
324 output_dir = Path(OUTPUT_DIR)
325 output_dir.mkdir(parents=True, exist_ok=True)
326
327 classifier = get_classifier(config)
328 feature_names = get_feature_extractor(config, is_text_document=True).get_feature_names()
329
330 # Combine train and test for sklearn train_test_split
331 X_combined = np.vstack([X_train, X_test])
332 y_combined = np.concatenate([y_train, y_test])
333
334 metrics = classifier.train(X_combined, y_combined, feature_names=feature_names)
335
336 # Save results
337 classifier.save(str(output_dir))
338 print(f"\n✓ Classifier saved to: {output_dir}")
339
340 # Save metrics
341 metrics_path = output_dir / 'training_metrics.json'
342 with open(metrics_path, 'w') as f:
343 json.dump(metrics, f, indent=2)
344
345 # Save class mapping
346 class_mapping = {
347 0: 'Copy-Move',
348 1: 'Splicing',
349 2: 'Generation'
350 }
351 mapping_path = output_dir / 'class_mapping.json'
352 with open(mapping_path, 'w') as f:
353 json.dump(class_mapping, f, indent=2)
354
355 print("\n" + "="*60)
356 print("✅ Classifier Training Complete!")
357 print("Algorithm Steps 7-8: DONE")
358 print("="*60)
359 print(f"\nResults:")
360 print(f" Test Accuracy: {metrics.get('test_accuracy', 'N/A')}")
361 print(f" Test F1 Score: {metrics.get('test_f1', 'N/A')}")
362 print(f"\nOutput: {output_dir}")
363 print("\nNext: Implement Steps 9-11 in inference pipeline")
364 print("="*60 + "\n")
365
366
367if __name__ == '__main__':
368 main()
369