Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,314 @@
+"""
+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'):
+    """Load the trained multimodal ensemble model.
+    
+    Args:
+        checkpoint_path: Path to multimodal_ensemble.pt
+        device: 'cpu' or 'cuda'
+        
+    Returns:
+        model: MultimodalDeepfakeDetector
+        config: training configuration dict
+    """
+    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,  # No dropout at inference
+    )
+    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):
+    """Classify a single image as real or AI-generated/deepfake.
+    
+    Args:
+        model: MultimodalDeepfakeDetector
+        image_path_or_pil: Path to image or PIL Image
+        device: computation device
+        return_gradcam: whether to generate explainability map
+        
+    Returns:
+        dict with:
+            - prediction: 'real' or 'fake'
+            - confidence: float [0, 1] (probability of being fake)
+            - gradcam: numpy array (H, W) heatmap if 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:
+        # Enable gradients for 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)  # Heatmap for "fake" class
+        output['gradcam'] = cam.squeeze().cpu().numpy()
+        gradcam.remove_hooks()
+
+    return output
+
+
+def classify_video(model, video_path, device='cpu', num_frames=32, aggregation='mean'):
+    """Classify a video as real or deepfake.
+    
+    Extracts frames uniformly, classifies each, and aggregates.
+    
+    Args:
+        model: MultimodalDeepfakeDetector
+        video_path: Path to video file
+        device: computation device
+        num_frames: number of frames to sample
+        aggregation: 'mean', 'max', or 'voting'
+        
+    Returns:
+        dict with video-level prediction and per-frame scores
+    """
+    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):
+    """Classify text as human-written or AI-generated.
+    
+    Args:
+        model: MultimodalDeepfakeDetector
+        text: input text string
+        tokenizer: optional pre-loaded tokenizer
+        device: computation device
+        max_length: max sequence length
+        
+    Returns:
+        dict with prediction and confidence
+    """
+    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'):
+    """Combined multimodal classification.
+    
+    Uses both image and text when available, with learned fusion weights.
+    
+    Returns:
+        dict with combined prediction, confidence, and per-modality scores
+    """
+    from preprocessing import get_image_transforms
+    from transformers import AutoTokenizer
+
+    images = None
+    input_ids = None
+    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()},
+    }
+
+    # Show fusion weights
+    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):
+    """Visualize GradCAM overlay on the original image.
+    
+    Args:
+        image_path: Path to original image
+        gradcam_heatmap: (H, W) numpy array from classify_image
+        confidence: fake confidence score
+        save_path: optional path to save visualization
+    """
+    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.show()
+    plt.close()
+
+
+# ============================================================
+# Demo / Usage Example
+# ============================================================
+if __name__ == '__main__':
+    print("=" * 60)
+    print("Multimodal Deepfake Detection - Inference Demo")
+    print("=" * 60)
+    print()
+    print("Usage:")
+    print("  from inference import load_model, classify_image, classify_text, classify_multimodal")
+    print()
+    print("  # Load model")
+    print("  model, config = load_model('output/multimodal_ensemble.pt', device='cuda')")
+    print()
+    print("  # Image classification")
+    print("  result = classify_image(model, 'face.jpg', device='cuda')")
+    print("  print(f'Prediction: {result[\"prediction\"]} (confidence: {result[\"confidence\"]:.2%})')")
+    print()
+    print("  # Text classification")
+    print("  result = classify_text(model, 'This text was generated by AI...')")
+    print("  print(f'Prediction: {result[\"prediction\"]} (confidence: {result[\"confidence\"]:.2%})')")
+    print()
+    print("  # Video classification")
+    print("  result = classify_video(model, 'video.mp4', device='cuda')")
+    print("  print(f'Prediction: {result[\"prediction\"]} (confidence: {result[\"confidence\"]:.2%})')")
+    print()
+    print("  # Multimodal (image + text)")
+    print("  result = classify_multimodal(model, image_path_or_pil='face.jpg', text='Caption...')")
+    print("  print(f'Combined: {result[\"prediction\"]} — Scores: {result[\"modality_scores\"]}')")