logo
Browse Source

update the deepfake.

Signed-off-by: jinlingxu06 <jinling.xu@zilliz.com>
main
jinlingxu06 2 years ago
parent
commit
757cca3ce5
  1. 16
      README.md
  2. 25
      deepfake.py

16
README.md

@ -19,19 +19,21 @@ and use deepfake operator to predict the probabilities of fake videos.
```python ```python
import towhee
(
towhee.glob['path']('/home/test_video')
.deepfake['path', 'scores']()
.select['path', 'scores']()
.show()
from towhee.dc2 import pipe, ops, DataCollection
p = (
pipe.input('path')
.map('path', 'scores', ops.towhee.deepfake)
.output('scores')
) )
DataCollection(p('./deepfake_video/test/aagfhgtpmv.mp4').get_dict()).show()
``` ```
<img src="./deepfake.png" height="100px"/> <img src="./deepfake.png" height="100px"/>
```shell ```shell
[0.9893, 0.9097]
[0.99]
``` ```
<br /> <br />

25
deepfake.py

@ -30,7 +30,6 @@ class Deepfake(NNOperator):
self.model_paths = [os.path.join(weights_dir,model) for model in os.listdir(weights_dir)] 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") self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def __call__(self, filepath: string) -> list:
from kernel_utils import VideoReader, FaceExtractor, confident_strategy, predict_on_video from kernel_utils import VideoReader, FaceExtractor, confident_strategy, predict_on_video
from classifiers import DeepFakeClassifier from classifiers import DeepFakeClassifier
models = [] models = []
@ -42,17 +41,21 @@ class Deepfake(NNOperator):
model.load_state_dict({re.sub("^module.", "", k): v for k, v in state_dict.items()}, strict=False) model.load_state_dict({re.sub("^module.", "", k): v for k, v in state_dict.items()}, strict=False)
model.eval() model.eval()
del checkpoint del checkpoint
models.append(model.half())
frames_per_video = 32
models.append(model.float())
self.frames_per_video = 32
video_reader = VideoReader() video_reader = VideoReader()
video_read_fn = lambda x: video_reader.read_frames(x, num_frames=frames_per_video)
face_extractor = FaceExtractor(video_read_fn)
input_size = 384
strategy = confident_strategy
#stime = time.time()
prediction = predict_on_video(False, face_extractor=face_extractor, video_path=filepath,
input_size=input_size, batch_size=frames_per_video, models=models,
strategy=strategy, apply_compression=False)
video_read_fn = lambda x: video_reader.read_frames(x, num_frames=self.frames_per_video)
self.face_extractor = FaceExtractor(video_read_fn)
self.input_size = 384
self.strategy = confident_strategy
self.models = models
self.predict_on_video = predict_on_video
def __call__(self, filepath: string) -> list:
prediction = self.predict_on_video(False, face_extractor=self.face_extractor, video_path=filepath,
input_size=self.input_size, batch_size=self.frames_per_video, models=self.models,
strategy=self.strategy, apply_compression=False)
''' '''
test_videos = sorted([x for x in os.listdir(filepath) if x[-4:] == ".mp4"]) test_videos = sorted([x for x in os.listdir(filepath) if x[-4:] == ".mp4"])
print("Predicting {} videos".format(len(test_videos))) print("Predicting {} videos".format(len(test_videos)))

Loading…
Cancel
Save