File size: 2,322 Bytes
9e14838 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import torch
from huggingface_hub import hf_hub_download
from lightning.fabric import Fabric
from PIL import Image
from src.config import Config
from src.model.dfdet import DeepfakeDetectionModel
DEVICES = [0]
torch.set_float32_matmul_precision("high")
# Check if weights/model.ckpt exists, if not, download it from huggingface
repo_id = "yermandy/deepfake-detection"
filename = "model.ckpt"
model_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir="weights")
# Load checkpoint
ckpt = torch.load(model_path, map_location="cpu")
run_name = ckpt["hyper_parameters"]["run_name"]
print(run_name)
# Initialize model from config
model = DeepfakeDetectionModel(Config(**ckpt["hyper_parameters"]))
model.eval()
# Load model state dict
model.load_state_dict(ckpt["state_dict"])
# Get preprocessing function
preprocessing = model.get_preprocessing()
# Load some images
paths = [
"datasets/CDFv2/Celeb-synthesis/id0_id1_0000/000.png",
"datasets/CDFv2/Celeb-synthesis/id0_id1_0000/045.png",
"datasets/CDFv2/Celeb-synthesis/id0_id1_0000/030.png",
"datasets/CDFv2/Celeb-synthesis/id0_id1_0000/015.png",
"datasets/CDFv2/YouTube-real/00000/000.png",
"datasets/CDFv2/YouTube-real/00000/014.png",
"datasets/CDFv2/YouTube-real/00000/028.png",
"datasets/CDFv2/YouTube-real/00000/043.png",
"datasets/CDFv2/Celeb-real/id0_0000/045.png",
"datasets/CDFv2/Celeb-real/id0_0000/030.png",
"datasets/CDFv2/Celeb-real/id0_0000/015.png",
"datasets/CDFv2/Celeb-real/id0_0000/000.png",
]
# To pillow images
pillow_images = [Image.open(image) for image in paths]
# To tensors
batch_images = torch.stack([preprocessing(image) for image in pillow_images])
precision = ckpt["hyper_parameters"]["precision"]
fabric = Fabric(accelerator="cuda", devices=DEVICES, precision=precision)
fabric.launch()
model = fabric.setup_module(model)
# perform inference
with torch.no_grad():
# Move batch_images to the correct device and dtype
batch_images = batch_images.to(fabric.device).to(model.dtype)
# Forward pass
output = model(batch_images)
# logits to probabilities
softmax_output = output.logits_labels.softmax(dim=1).cpu().numpy()
for path, (p_real, p_fake) in zip(paths, softmax_output):
print(f"p(real) = {p_real:.4f}, p(fake) = {p_fake:.4f}, image: {path}")
|