| 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)
|
|
|
|
|
| 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
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
| 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
|
|
|
|
|
|
|
| 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
|
|
|
|
|
