File size: 3,058 Bytes
0a19a23 |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
"""
Example: Using the model for deepfake detection
"""
import torch
from torchvision import transforms
from PIL import Image
from model import load_model
import json
# Load model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = load_model('model_best_checkpoint.ckpt', device=device)
# Load calibrated thresholds
with open('thresholds_calibrated.json', 'r') as f:
config = json.load(f)
threshold = config['reconstruction_thresholds']['thresholds']['balanced']['value']
print(f"Using threshold: {threshold:.6f}")
# Prepare image preprocessing
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
def detect_deepfake(image_path, model, threshold, device):
"""
Detect if an image is likely a deepfake based on reconstruction error.
Args:
image_path: Path to image file
model: Loaded autoencoder model
threshold: MSE threshold for detection
device: torch device
Returns:
is_fake: Boolean indicating if image is likely fake
error: Reconstruction error value
confidence: Confidence score (0-1)
"""
# Load and preprocess image
image = Image.open(image_path).convert('RGB')
input_tensor = transform(image).unsqueeze(0).to(device)
# Calculate reconstruction error
with torch.no_grad():
error = model.reconstruction_error(input_tensor, reduction='none')
error_value = error.item()
is_fake = error_value > threshold
# Calculate confidence (normalized error relative to threshold)
confidence = min(abs(error_value - threshold) / threshold, 1.0)
return is_fake, error_value, confidence
# Example usage
image_path = "test_image.jpg"
is_fake, error, confidence = detect_deepfake(image_path, model, threshold, device)
print(f"\nResults for: {image_path}")
print(f"Reconstruction Error: {error:.6f}")
print(f"Threshold: {threshold:.6f}")
print(f"Classification: {'FAKE' if is_fake else 'REAL'}")
print(f"Confidence: {confidence:.2%}")
# Batch processing example
def batch_detect(image_paths, model, threshold, device):
"""Process multiple images at once"""
images = []
for path in image_paths:
img = Image.open(path).convert('RGB')
images.append(transform(img))
batch = torch.stack(images).to(device)
with torch.no_grad():
errors = model.reconstruction_error(batch, reduction='none')
results = []
for i, error in enumerate(errors):
is_fake = error.item() > threshold
results.append({
'path': image_paths[i],
'error': error.item(),
'is_fake': is_fake
})
return results
# Example batch processing
# image_paths = ["img1.jpg", "img2.jpg", "img3.jpg"]
# results = batch_detect(image_paths, model, threshold, device)
# for r in results:
# print(f"{r['path']}: {'FAKE' if r['is_fake'] else 'REAL'} (error: {r['error']:.6f})")
|