File size: 2,069 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 | import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import CLIPProcessor
DEVICE = "cuda:0"
DTYPE = torch.bfloat16
torch.set_float32_matmul_precision("high")
# Check if weights/model.torchscript exists, if not, download it from huggingface
repo_id = "yermandy/deepfake-detection"
filename = "model.torchscript"
model_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir="weights")
# Load checkpoint
model = torch.jit.load(model_path, map_location=DEVICE)
# Load preprocessing function
preprocess = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
# 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(
[preprocess(images=image, return_tensors="pt")["pixel_values"][0] for image in pillow_images]
)
# Set model to evaluation mode
model.eval()
# Move model to the correct device and dtype
model = model.to(DEVICE).to(DTYPE)
# Move inputs to the correct device and dtype
batch_images = batch_images.to(DEVICE).to(DTYPE)
with torch.no_grad():
with torch.autocast(device_type="cuda", dtype=DTYPE):
# Forward pass
output = model(batch_images)
softmax_output = output.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}")
|