update

__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2021 Zilliz. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

deepfake.py

@@ -3,11 +3,11 @@ import os
 import re
 import string
 import time
-
+import sys
+from pathlib import Path
 import torch
 import pandas as pd
-from .kernel_utils import VideoReader, FaceExtractor, confident_strategy, predict_on_video_set
-from .classifiers import DeepFakeClassifier
+
 import towhee
 from towhee.operator.base import NNOperator, OperatorFlag
 from towhee import register
@@ -21,16 +21,21 @@ log = logging.getLogger()
 
 class Deepfake(NNOperator):
     '''
-    ViT Distillation
+    Deepfake
     '''
     def __init__(self):
-        pass
+        super().__init__()
+        sys.path.append(str(Path(__file__).parent))
+        weights_dir = os.path.join(str(Path(__file__).parent),"weights/")
+        self.model_paths = [os.path.join(weights_dir,model) for model in os.listdir(weights_dir)]
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
     def __call__(self, filepath: string) -> list:
-        weights_dir = "weights/"
+        from kernel_utils import VideoReader, FaceExtractor, confident_strategy, predict_on_video_set
+        from classifiers import DeepFakeClassifier
         models = []
-        model_paths = [os.path.join(weights_dir,model) for model in os.listdir(weights_dir)]
-        for path in model_paths:
-            model = DeepFakeClassifier(encoder="tf_efficientnet_b7_ns").to("cuda")
+        for path in self.model_paths:
+            model = DeepFakeClassifier(encoder="tf_efficientnet_b7_ns").to(self.device)
             print("loading state dict {}".format(path))
             checkpoint = torch.load(path, map_location="cpu")
             state_dict = checkpoint.get("state_dict", checkpoint)

kernel_utils.py

@@ -15,7 +15,7 @@ log = logging.getLogger()
 mean = [0.485, 0.456, 0.406]
 std = [0.229, 0.224, 0.225]
 normalize_transform = Normalize(mean, std)
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class VideoReader:
     """Helper class for reading one or more frames from a video file."""
@@ -201,7 +201,7 @@ class VideoReader:
 class FaceExtractor:
     def __init__(self, video_read_fn):
         self.video_read_fn = video_read_fn
-        self.detector = MTCNN(margin=0, thresholds=[0.7, 0.8, 0.8], device="cuda")
+        self.detector = MTCNN(margin=0, thresholds=[0.7, 0.8, 0.8], device=device)
 
     def process_videos(self, input_dir, filenames, video_idxs):
         videos_read = []
@@ -306,7 +306,7 @@ def isotropically_resize_image(img, size, interpolation_down=cv2.INTER_AREA, int
 def dist(p1, p2):
     return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
 
-detector = MTCNN(margin=0, thresholds=(0.7, 0.8, 0.8), device="cuda")
+detector = MTCNN(margin=0, thresholds=(0.7, 0.8, 0.8), device=device)
 def predict_on_video(distill, face_extractor, video_path, batch_size, input_size, models, strategy=np.mean,
                      apply_compression=False):
     batch_size *= 4
@@ -344,7 +344,7 @@ def predict_on_video(distill, face_extractor, video_path, batch_size, input_size
                     else:
                         pass
             if n > 0:
-                x = torch.tensor(x, device="cuda").float()
+                x = torch.tensor(x, device=device).float()
                 #e = torch.tensor(e, device="cuda").float() #eye
                 # Preprocess the images.
                 x = x.permute((0, 3, 1, 2))