fork-detect-fake / pipeline.py
wyyadd's picture
feat: update pipeline
d30298c
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)