Add inference pipeline

inference.py

@@ -0,0 +1,142 @@
+"""
+Inference Pipeline for Multimodal Deepfake Detection
+=====================================================
+Supports:
+- Single image classification with confidence + GradCAM heatmap
+- Video classification (frame-by-frame → aggregated score)
+- Text classification (human vs AI-generated)
+- Multimodal (image + text combined)
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+import json
+import os
+
+
+def load_model(checkpoint_path, device='cpu'):
+    from model import MultimodalDeepfakeDetector
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    config = checkpoint['config']
+    model = MultimodalDeepfakeDetector(visual_pretrained=False, text_model_name=config['text_model_name'], dropout=0.0)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model = model.to(device)
+    model.eval()
+    return model, config
+
+
+def classify_image(model, image_path_or_pil, device='cpu', return_gradcam=True):
+    from preprocessing import get_image_transforms
+    from model import GradCAM
+    if isinstance(image_path_or_pil, str):
+        image = Image.open(image_path_or_pil).convert('RGB')
+    else:
+        image = image_path_or_pil.convert('RGB')
+    transform = get_image_transforms('eval', 224)
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        results = model(images=image_tensor, modality='visual')
+    confidence = results['confidence'].item()
+    prediction = 'fake' if confidence > 0.5 else 'real'
+    output = {'prediction': prediction, 'confidence': confidence, 'visual_score': results['modality_scores']['visual'].item()}
+    if return_gradcam:
+        model.visual_branch.eval()
+        gradcam = GradCAM(model.visual_branch, model.visual_branch.get_gradcam_target_layer())
+        image_tensor_grad = image_tensor.clone().requires_grad_(True)
+        cam = gradcam.generate(image_tensor_grad, class_idx=1)
+        output['gradcam'] = cam.squeeze().cpu().numpy()
+        gradcam.remove_hooks()
+    return output
+
+
+def classify_video(model, video_path, device='cpu', num_frames=32, aggregation='mean'):
+    from preprocessing import extract_video_frames, get_image_transforms
+    from model import aggregate_video_predictions
+    frames = extract_video_frames(video_path, num_frames=num_frames)
+    transform = get_image_transforms('eval', 224)
+    frame_scores = []
+    model.eval()
+    with torch.no_grad():
+        for frame in frames:
+            image_tensor = transform(frame.convert('RGB')).unsqueeze(0).to(device)
+            results = model(images=image_tensor, modality='visual')
+            frame_scores.append(results['confidence'].item())
+    video_confidence = aggregate_video_predictions(torch.tensor(frame_scores), method=aggregation)
+    return {
+        'prediction': 'fake' if video_confidence > 0.5 else 'real',
+        'confidence': video_confidence,
+        'frame_scores': frame_scores,
+        'num_frames_analyzed': len(frames),
+        'aggregation_method': aggregation,
+    }
+
+
+def classify_text(model, text, tokenizer=None, device='cpu', max_length=512):
+    from transformers import AutoTokenizer
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained('roberta-base')
+    encoding = tokenizer(text, max_length=max_length, padding='max_length', truncation=True, return_tensors='pt')
+    input_ids = encoding['input_ids'].to(device)
+    attention_mask = encoding['attention_mask'].to(device)
+    model.eval()
+    with torch.no_grad():
+        results = model(input_ids=input_ids, attention_mask=attention_mask, modality='text')
+    confidence = results['confidence'].item()
+    return {
+        'prediction': 'ai_generated' if confidence > 0.5 else 'human',
+        'confidence': confidence,
+        'text_score': results['modality_scores']['text'].item(),
+    }
+
+
+def classify_multimodal(model, image_path_or_pil=None, text=None, tokenizer=None, device='cpu'):
+    from preprocessing import get_image_transforms
+    from transformers import AutoTokenizer
+    images = input_ids = attention_mask = None
+    if image_path_or_pil is not None:
+        if isinstance(image_path_or_pil, str):
+            image = Image.open(image_path_or_pil).convert('RGB')
+        else:
+            image = image_path_or_pil.convert('RGB')
+        transform = get_image_transforms('eval', 224)
+        images = transform(image).unsqueeze(0).to(device)
+    if text is not None:
+        if tokenizer is None:
+            tokenizer = AutoTokenizer.from_pretrained('roberta-base')
+        encoding = tokenizer(text, max_length=512, padding='max_length', truncation=True, return_tensors='pt')
+        input_ids = encoding['input_ids'].to(device)
+        attention_mask = encoding['attention_mask'].to(device)
+    model.eval()
+    with torch.no_grad():
+        results = model(images=images, input_ids=input_ids, attention_mask=attention_mask, modality='auto')
+    confidence = results['confidence'].item()
+    output = {
+        'prediction': 'fake/ai_generated' if confidence > 0.5 else 'real/human',
+        'confidence': confidence,
+        'modality_scores': {k: v.item() for k, v in results['modality_scores'].items()},
+    }
+    with torch.no_grad():
+        fusion_weights = F.softmax(model.fusion_weights, dim=0)
+        output['fusion_weights'] = {'visual': fusion_weights[0].item(), 'text': fusion_weights[1].item()}
+    return output
+
+
+def visualize_gradcam(image_path, gradcam_heatmap, confidence, save_path=None):
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+    image = Image.open(image_path).convert('RGB')
+    image_np = np.array(image.resize((224, 224))) / 255.0
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+    axes[0].imshow(image_np); axes[0].set_title('Original'); axes[0].axis('off')
+    axes[1].imshow(gradcam_heatmap, cmap='jet'); axes[1].set_title('GradCAM Heatmap'); axes[1].axis('off')
+    axes[2].imshow(image_np); axes[2].imshow(gradcam_heatmap, cmap='jet', alpha=0.4)
+    axes[2].set_title('Overlay (Explanation)'); axes[2].axis('off')
+    label = "FAKE" if confidence > 0.5 else "REAL"
+    color = 'red' if confidence > 0.5 else 'green'
+    fig.suptitle(f'{label} — Confidence: {confidence:.2%}', fontsize=16, fontweight='bold', color=color)
+    plt.tight_layout()
+    if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight')
+    plt.close()