Deepfake Detection

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, models
from PIL import Image

class DeepfakeDetector(nn.Module):
    """
    Kiến trúc detector 2 nhánh:
    - Nhánh 1: EfficientNet-B4 đọc ảnh RGB gốc
    - Nhánh 2: CNN đọc phổ FFT (high-freq features)
    - Fusion: concat + FC head
    """
    def __init__(self, num_classes: int = 2):
        super().__init__()

        # Nhánh RGB
        backbone = models.efficientnet_b4(weights=None)
        backbone.classifier = nn.Identity()
        self.rgb_branch = backbone
        self.rgb_dim = 1792

        # Nhánh tần số
        self.freq_branch = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
        )
        self.freq_dim = 128

        # Fusion head
        self.head = nn.Sequential(
            nn.Linear(self.rgb_dim + self.freq_dim, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes),
        )

    @staticmethod
    def compute_spectrum(x: torch.Tensor) -> torch.Tensor:
        """x: [B, 3, H, W] -> log-magnitude spectrum [B, 1, H, W]."""
        gray = x.mean(dim=1, keepdim=True)
        fft = torch.fft.fft2(gray, norm="ortho")
        mag = torch.abs(torch.fft.fftshift(fft, dim=(-2, -1)))
        return torch.log1p(mag)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rgb_feat = self.rgb_branch(x)
        spectrum = self.compute_spectrum(x)
        freq_feat = self.freq_branch(spectrum)
        fused = torch.cat([rgb_feat, freq_feat], dim=1)
        return self.head(fused)


# ============== Inference pipeline ==============
def predict_deepfake(model, image_path: str, device: str = "cuda"):
    transform = transforms.Compose([
        transforms.Resize((380, 380)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        ),
    ])
    img = Image.open(image_path).convert("RGB")
    x = transform(img).unsqueeze(0).to(device)

    model.eval()
    with torch.no_grad():
        logits = model(x)
        probs = F.softmax(logits, dim=-1)
        p_fake = probs[0, 1].item()

    if p_fake > 0.5:
        print(f"[CẢNH BÁO] Deepfake — confidence {p_fake:.1%}")
        print("Không nên tin tưởng. Xác minh qua kênh khác.")
    else:
        print(f"[OK] Có vẻ thật — confidence {1 - p_fake:.1%}")
        print("Vẫn nên xác minh nếu có yêu cầu tài chính/pháp lý.")
    return p_fake


if __name__ == "__main__":
    model = DeepfakeDetector(num_classes=2)
    model.load_state_dict(torch.load("checkpoints/detector_ff++_cdf.pth"))
    model.to("cuda")
    predict_deepfake(model, "evidence/zalo_call_frame_001.jpg")