feat: update pipeline

deepfakeconfig.py

@@ -1,7 +1,10 @@
 from transformers import PretrainedConfig
-import torch 
+import torch
+
+
 class DeepFakeConfig(PretrainedConfig):
     model_type = "ResNet"
-    def __init__(self,**kwargs):
+
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
-        self.DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        self.DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'

deepfakemodel.py

@@ -1,9 +1,11 @@
-from transformers import  PreTrainedModel
-from facenet_pytorch import MTCNN, InceptionResnetV1
+from facenet_pytorch import InceptionResnetV1
+from transformers import PreTrainedModel
 from .deepfakeconfig import DeepFakeConfig
 
+
 class DeepFakeModel(PreTrainedModel):
     config_class = DeepFakeConfig
+
     def __init__(self, config):
         super().__init__(config)
         self.model = InceptionResnetV1(
@@ -15,4 +17,4 @@ class DeepFakeModel(PreTrainedModel):
 
 
 DeepFakeConfig.register_for_auto_class()
-DeepFakeModel.register_for_auto_class("AutoModelForImageClassification")
+DeepFakeModel.register_for_auto_class("AutoModelForImageClassification")

pipeline.py

@@ -1,72 +1,53 @@
-from transformers.pipelines import PIPELINE_REGISTRY
-from transformers import Pipeline, AutoModelForImageClassification
 import torch
+import torch.nn.functional as F
 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
+from transformers import Pipeline
+
 
 class DeepFakePipeline(Pipeline):
-    def __init__(self,**kwargs):
-        Pipeline.__init__(self,**kwargs)
+    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)
+    def _forward(self, input):
+        return input
 
-      input_image = Image.open(input_image)
-      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)
-
-      # convert the face into a numpy array to be able to plot it
-      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)
+    def postprocess(self, confidences):
+        out = {"confidences": confidences}
+        return out
 
-      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)
+    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)

requirements.txt

@@ -1,6 +1,4 @@
 Pillow
 facenet-pytorch==2.5.2
-torch==1.11.0
-opencv-python
-grad-cam
+torch
 transformers