| import os.path |
| import re |
| import torch |
| import time |
| import tempfile |
|
|
| import streamlit as st |
| from training.zoo.classifiers import DeepFakeClassifier |
| from kernel_utils import VideoReader, FaceExtractor, confident_strategy, predict_on_video_set |
|
|
|
|
| def load_model(): |
| path = 'weights/best.pth' |
| model = DeepFakeClassifier(encoder="tf_efficientnet_b7_ns") |
| print("loading state dict {}".format(path)) |
| checkpoint = torch.load(path, map_location="cpu") |
| state_dict = checkpoint.get("state_dict", checkpoint) |
| model.load_state_dict( |
| {re.sub("^module.", "", k): v for k, v in state_dict.items()}, |
| strict=True) |
| model.eval() |
| del checkpoint |
| return model |
|
|
|
|
| def write_bytesio_to_file(filename, bytesio): |
| with open(filename, "wb") as outfile: |
| outfile.write(bytesio.getbuffer()) |
|
|
|
|
| def load_video(): |
| uploaded_file = st.file_uploader(label='Pick a video (mp4) file to test') |
| if uploaded_file is not None: |
| video_data = uploaded_file.getvalue() |
| tfile = tempfile.NamedTemporaryFile(delete=False) |
| tfile.write(video_data) |
| return tfile.name |
| else: |
| return None |
|
|
|
|
| def inference(model, test_video): |
| frames_per_video = 32 |
| 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 = 380 |
| strategy = confident_strategy |
|
|
| test_videos = [test_video] |
| print("Predicting {} videos".format(len(test_videos))) |
| models = [model] |
| predictions = predict_on_video_set(face_extractor=face_extractor, |
| input_size=input_size, models=models, |
| strategy=strategy, |
| frames_per_video=frames_per_video, |
| videos=test_videos, |
| num_workers=6, test_dir="test_video") |
| st.write("Prediction: ", predictions[0]) |
|
|
|
|
| def main(): |
| st.title('Deepfake video inference demo') |
| model = load_model() |
| video_data_path = load_video() |
|
|
| if video_data_path is not None and os.path.exists(video_data_path): |
| st.video(video_data_path) |
|
|
| result = st.button('Run on video') |
| if result: |
| st.write("Inference on video...") |
| stime = time.time() |
| inference(model, video_data_path) |
| st.write("Elapsed time: ", time.time() - stime, " seconds") |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
