File size: 1,607 Bytes
402b016 d30298c 402b016 d30298c 402b016 d30298c 402b016 d30298c 402b016 d30298c 402b016 d30298c 402b016 d30298c |
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 |
import torch
import torch.nn.functional as F
from PIL import Image
from facenet_pytorch import MTCNN
from transformers import Pipeline
class DeepFakePipeline(Pipeline):
def __init__(self, **kwargs):
Pipeline.__init__(self, **kwargs)
def _sanitize_parameters(self, **kwargs):
return {}, {}, {}
def preprocess(self, inputs):
return inputs
def _forward(self, input):
return input
def postprocess(self, confidences):
out = {"confidences": confidences}
return out
def predict(self, input_image: Image.Image):
DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
mtcnn = MTCNN(
select_largest=False,
post_process=False,
device=DEVICE)
mtcnn.to(DEVICE)
model = self.model.model
model.to(DEVICE)
face = mtcnn(input_image)
if face is None:
raise Exception('No face detected')
face = face.unsqueeze(0) # add the batch dimension
face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
face = face.to(DEVICE)
face = face.to(torch.float32)
face = face / 255.0
with torch.no_grad():
output = torch.sigmoid(model(face).squeeze(0))
real_prediction = 1 - output.item()
fake_prediction = output.item()
confidences = {
'real': real_prediction,
'fake': fake_prediction
}
return self.postprocess(confidences)
|