|
from transformers.pipelines import PIPELINE_REGISTRY
|
|
from transformers import Pipeline, AutoModelForImageClassification
|
|
import torch
|
|
from PIL import Image
|
|
import cv2
|
|
from pytorch_grad_cam import GradCAM
|
|
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
|
|
from pytorch_grad_cam.utils.image import show_cam_on_image
|
|
from facenet_pytorch import MTCNN
|
|
import torch.nn.functional as F
|
|
|
|
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,face_with_mask):
|
|
out = {"confidences":confidences,
|
|
"face_with_mask": face_with_mask}
|
|
return out
|
|
|
|
def predict(self,input_image:str):
|
|
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)
|
|
|
|
input_image = Image.open(input_image)
|
|
face = mtcnn(input_image)
|
|
if face is None:
|
|
raise Exception('No face detected')
|
|
|
|
face = face.unsqueeze(0)
|
|
face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
|
|
|
|
|
|
prev_face = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
|
|
prev_face = prev_face.astype('uint8')
|
|
|
|
face = face.to(DEVICE)
|
|
face = face.to(torch.float32)
|
|
face = face / 255.0
|
|
face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
|
|
|
|
target_layers=[model.block8.branch1[-1]]
|
|
cam = GradCAM(model=model, target_layers=target_layers)
|
|
targets = [ClassifierOutputTarget(0)]
|
|
grayscale_cam = cam(input_tensor=face, targets=targets,eigen_smooth=True)
|
|
grayscale_cam = grayscale_cam[0, :]
|
|
visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True)
|
|
face_with_mask = cv2.addWeighted(prev_face, 1, visualization, 0.5, 0)
|
|
|
|
with torch.no_grad():
|
|
output = torch.sigmoid(model(face).squeeze(0))
|
|
prediction = "real" if output.item() < 0.5 else "fake"
|
|
|
|
real_prediction = 1 - output.item()
|
|
fake_prediction = output.item()
|
|
|
|
confidences = {
|
|
'real': real_prediction,
|
|
'fake': fake_prediction
|
|
}
|
|
return self.postprocess(confidences, face_with_mask) |
