File size: 3,855 Bytes
284919d | 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 | import torch
from torch import nn
from torchvision import models, transforms
import cv2
import face_recognition
import numpy as np
from PIL import Image
import os
# --- 1. MODEL ARCHITECTURE (From views.py) ---
class DeepfakeModel(nn.Module):
def __init__(self, num_classes=2, latent_dim=2048, lstm_layers=1, hidden_dim=2048, bidirectional=False):
super(DeepfakeModel, self).__init__()
resnext = models.resnext50_32x4d(pretrained=True)
self.model = nn.Sequential(*list(resnext.children())[:-2])
self.lstm = nn.LSTM(latent_dim, hidden_dim, lstm_layers, bidirectional)
self.dp = nn.Dropout(0.4)
self.linear1 = nn.Linear(2048, num_classes)
self.avgpool = nn.AdaptiveAvgPool2d(1)
def forward(self, x):
batch_size, seq_length, c, h, w = x.shape
x = x.view(batch_size * seq_length, c, h, w)
fmap = self.model(x)
x = self.avgpool(fmap)
x = x.view(batch_size, seq_length, 2048)
x_lstm, _ = self.lstm(x, None)
return fmap, self.dp(self.linear1(x_lstm[:, -1, :]))
# --- 2. CUSTOM HANDLER FOR HUGGING FACE ---
class EndpointHandler():
def __init__(self, path=""):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = DeepfakeModel().to(self.device)
# Look for the .pt file in the repository
model_weight_path = os.path.join(path, "pytorch_model.bin")
if not os.path.exists(model_weight_path):
# Fallback to any .pt file if pytorch_model.bin isn't found
for f in os.listdir(path):
if f.endswith(".pt"):
model_weight_path = os.path.join(path, f)
break
self.model.load_state_dict(torch.load(model_weight_path, map_location=self.device))
self.model.eval()
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((112, 112)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def __call__(self, data):
# Data normally contains "inputs" which is the raw video file (bytes)
inputs = data.pop("inputs", data)
# Save bytes to temporary file to process with OpenCV
temp_video = "temp_video.mp4"
with open(temp_video, "wb") as f:
f.write(inputs)
# 1. Extract Frames & Crop Faces
cap = cv2.VideoCapture(temp_video)
frames = []
while cap.isOpened() and len(frames) < 20: # Process first 20 frames
ret, frame = cap.read()
if not ret: break
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
face_locations = face_recognition.face_locations(rgb_frame)
if len(face_locations) > 0:
top, right, bottom, left = face_locations[0]
face_crop = rgb_frame[top:bottom, left:right]
frames.append(self.transform(face_crop))
cap.release()
os.remove(temp_video)
if len(frames) < 10:
return {"error": "Not enough faces detected in video."}
# 2. Run Inference
input_tensor = torch.stack(frames).unsqueeze(0).to(self.device) # [1, Seq, C, H, W]
with torch.no_grad():
_, outputs = self.model(input_tensor)
probabilities = torch.softmax(outputs, dim=1)
confidence, prediction = torch.max(probabilities, 1)
result = "REAL" if prediction.item() == 1 else "FAKE"
return {
"label": result,
"confidence": round(float(confidence.item()) * 100, 2)
}
|