diff --git a/__init__.py b/__init__.py index e69de29..37f5bd7 100644 --- a/__init__.py +++ b/__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. diff --git a/deepfake.py b/deepfake.py index 0637ce8..a846080 100644 --- a/deepfake.py +++ b/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) diff --git a/kernel_utils.py b/kernel_utils.py index 23d05ca..67486e9 100644 --- a/kernel_utils.py +++ b/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))