File size: 2,951 Bytes
dbbbc4b | 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 100 101 102 103 104 105 106 107 | import torch
import torch.nn as nn
import json
import numpy as np
import os
from PIL import Image, ImageChops, ImageEnhance
IMAGE_SIZE = (128, 128)
# ---------------- ELA ----------------
def convert_to_ela_image(image_path, quality=95):
original = Image.open(image_path).convert("RGB")
temp_path = "temp_ela.jpg"
original.save(temp_path, "JPEG", quality=quality)
compressed = Image.open(temp_path)
diff = ImageChops.difference(original, compressed)
extrema = diff.getextrema()
max_diff = max([ex[1] for ex in extrema])
scale = 255.0 / max_diff if max_diff != 0 else 1
diff = ImageEnhance.Brightness(diff).enhance(scale)
diff = diff.resize(IMAGE_SIZE)
ela = np.array(diff, dtype=np.float32) / 255.0
ela = torch.tensor(ela).permute(2, 0, 1).unsqueeze(0)
return ela
# ---------------- MODEL ----------------
class DeepFakeCNN(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=5),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, kernel_size=3),
nn.ReLU(),
nn.MaxPool2d(2)
)
self._to_linear = None
self._get_flatten_size()
self.classifier = nn.Sequential(
nn.Linear(self._to_linear, 128),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(128, 2)
)
def _get_flatten_size(self):
with torch.no_grad():
x = torch.zeros(1, 3, 128, 128)
x = self.features(x)
self._to_linear = x.view(1, -1).shape[1]
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
return self.classifier(x)
# ---------------- LOAD MODEL ----------------
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def load_model(
model_path=os.path.join(BASE_DIR, "models", "deepfake_cnn_2.pth"),
config_path=os.path.join(BASE_DIR, "models", "config_cnn_2.json")
):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
with open(config_path, "r") as f:
config = json.load(f)
model = DeepFakeCNN().to(device)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval()
return model, config, device
# Prediction Function
def predict_image(image_path, model, config, device):
ela = convert_to_ela_image(image_path).to(device)
with torch.no_grad():
outputs = model(ela)
probs = torch.softmax(outputs, dim=1)
pred = torch.argmax(probs, dim=1).item()
confidence = probs[0, pred].item()
label = config["class_mapping"][str(pred)]
return label, confidence
|