From ea994bbe34941478644af6a70cc7ba0e9d5e45fa Mon Sep 17 00:00:00 2001 From: ko1N Date: Sun, 17 Jan 2021 15:24:57 +0100 Subject: [PATCH 1/5] Added Docker setup --- README.md | 13 +++++++++++++ benchmark/MiddleBury_Other.py | 2 +- benchmark/Vimeo90K.py | 2 +- docker/Dockerfile | 21 +++++++++++++++++++++ docker/rife.sh | 2 ++ inference_img.py | 2 +- inference_video.py | 15 ++++++++++----- 7 files changed, 49 insertions(+), 8 deletions(-) create mode 100644 docker/Dockerfile create mode 100644 docker/rife.sh diff --git a/README.md b/README.md index 7a3b820..26be3c6 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,19 @@ You can also use pngs to generate gif: ffmpeg -r 10 -f image2 -i output/img%d.png -s 448x256 -vf "split[s0][s1];[s0]palettegen=stats_mode=single[p];[s1][p]paletteuse=new=1" output/slomo.gif ``` +### Run in docker +Place the pre-trained models in the `./docker/pretrained_models directory` + +Building the container: +``` +docker build -t rife -f docker/Dockerfile . +``` + +Running the container: +``` +docker run --rm -it -v $PWD:/host rife:latest --exp=1 --video=untitled.mp4 --output=untitled_rife.mp4 +``` + ## Evaluation Download [RIFE model](https://drive.google.com/file/d/1c1R7iF-ypN6USo-D2YH_ORtaH3tukSlo/view?usp=sharing) or [RIFE2F1.5C model](https://drive.google.com/file/d/1ve9w-cRWotdvvbU1KcgtsSm12l-JUkeT/view?usp=sharing) reported by our paper. diff --git a/benchmark/MiddleBury_Other.py b/benchmark/MiddleBury_Other.py index 18b4e16..7c0f77f 100644 --- a/benchmark/MiddleBury_Other.py +++ b/benchmark/MiddleBury_Other.py @@ -12,7 +12,7 @@ from model.RIFE import Model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Model() -model.load_model('./train_log') +model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log')) model.eval() model.device() diff --git a/benchmark/Vimeo90K.py b/benchmark/Vimeo90K.py index c984b4c..cd70e48 100644 --- a/benchmark/Vimeo90K.py +++ b/benchmark/Vimeo90K.py @@ -13,7 +13,7 @@ from model.RIFE import Model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Model() -model.load_model('./train_log') +model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log')) model.eval() model.device() diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..84ad60c --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,21 @@ +FROM python:3.8-slim + +# install deps +RUN apt-get update && apt-get -y install \ + bash ffmpeg + +# setup RIFE +WORKDIR /rife +COPY . . +RUN pip3 install -r requirements.txt + +ADD docker/rife.sh /usr/local/bin/rife +RUN chmod +x /usr/local/bin/rife + +# add pre-trained models +COPY docker/pretrained_models /rife/train_log + +WORKDIR /host +ENTRYPOINT ["rife"] + +ENV NVIDIA_DRIVER_CAPABILITIES all \ No newline at end of file diff --git a/docker/rife.sh b/docker/rife.sh new file mode 100644 index 0000000..d718c5c --- /dev/null +++ b/docker/rife.sh @@ -0,0 +1,2 @@ +#!/bin/sh +python3 /rife/inference_video.py $@ diff --git a/inference_img.py b/inference_img.py index 24a9ab0..633bd55 100644 --- a/inference_img.py +++ b/inference_img.py @@ -19,7 +19,7 @@ parser.add_argument('--exp', default=4, type=int) args = parser.parse_args() model = Model() -model.load_model('./train_log', -1) +model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), -1) model.eval() model.device() diff --git a/inference_video.py b/inference_video.py index 4976851..28fd756 100644 --- a/inference_video.py +++ b/inference_video.py @@ -59,6 +59,7 @@ if torch.cuda.is_available(): parser = argparse.ArgumentParser(description='Interpolation for a pair of images') parser.add_argument('--video', dest='video', type=str, default=None) +parser.add_argument('--output', dest='output', type=str, default=None) parser.add_argument('--img', dest='img', type=str, default=None) parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video') parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video') @@ -74,7 +75,7 @@ if not args.img is None: from model.RIFE_HD import Model model = Model() -model.load_model('./train_log', -1) +model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), -1) model.eval() model.device() @@ -107,12 +108,17 @@ else: lastframe = cv2.imread(os.path.join(args.img, videogen[0]))[:, :, ::-1].copy() videogen = videogen[1:] h, w, _ = lastframe.shape +vid_out_name = None vid_out = None if args.png: if not os.path.exists('vid_out'): os.mkdir('vid_out') else: - vid_out = cv2.VideoWriter('{}_{}X_{}fps.{}'.format(video_path_wo_ext, (2 ** args.exp), int(np.round(args.fps)), args.ext), fourcc, args.fps, (w, h)) + if args.output is not None: + vid_out_name = args.output + else: + vid_out_name = '{}_{}X_{}fps.{}'.format(video_path_wo_ext, (2 ** args.exp), int(np.round(args.fps)), args.ext) + vid_out = cv2.VideoWriter(vid_out_name, fourcc, args.fps, (w, h)) def clear_write_buffer(user_args, write_buffer): cnt = 0 @@ -211,9 +217,8 @@ if not vid_out is None: # move audio to new video file if appropriate if args.png == False and fpsNotAssigned == True and not args.skip and not args.video is None: - outputVideoFileName = '{}_{}X_{}fps.{}'.format(video_path_wo_ext, 2 ** args.exp, int(np.round(args.fps)), args.ext) try: - transferAudio(args.video, outputVideoFileName) + transferAudio(args.video, vid_out_name) except: print("Audio transfer failed. Interpolated video will have no audio") - os.rename("noAudio_"+outputVideoFileName, outputVideoFileName) + os.rename("noAudio_"+vid_out_name, vid_out_name) From 0f980597a104cbd4907d8cda2ce5aa42c6af87ac Mon Sep 17 00:00:00 2001 From: ko1N Date: Sun, 17 Jan 2021 16:05:48 +0100 Subject: [PATCH 2/5] Added ability to run both inference scripts from docker --- README.md | 5 ++++- docker/Dockerfile | 8 +++++--- docker/inference_img | 2 ++ docker/{rife.sh => inference_video} | 0 4 files changed, 11 insertions(+), 4 deletions(-) create mode 100644 docker/inference_img rename docker/{rife.sh => inference_video} (100%) diff --git a/README.md b/README.md index 26be3c6..6e377f1 100644 --- a/README.md +++ b/README.md @@ -90,7 +90,10 @@ docker build -t rife -f docker/Dockerfile . Running the container: ``` -docker run --rm -it -v $PWD:/host rife:latest --exp=1 --video=untitled.mp4 --output=untitled_rife.mp4 +docker run --rm -it -v $PWD:/host rife:latest inference_video --exp=1 --video=untitled.mp4 --output=untitled_rife.mp4 +``` +``` +docker run --rm -it -v $PWD:/host rife:latest inference_img --img img0.png img1.png --exp=4 ``` ## Evaluation diff --git a/docker/Dockerfile b/docker/Dockerfile index 84ad60c..eca0589 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -9,13 +9,15 @@ WORKDIR /rife COPY . . RUN pip3 install -r requirements.txt -ADD docker/rife.sh /usr/local/bin/rife -RUN chmod +x /usr/local/bin/rife +ADD docker/inference_img /usr/local/bin/inference_img +RUN chmod +x /usr/local/bin/inference_img +ADD docker/inference_video /usr/local/bin/inference_video +RUN chmod +x /usr/local/bin/inference_video # add pre-trained models COPY docker/pretrained_models /rife/train_log WORKDIR /host -ENTRYPOINT ["rife"] +ENTRYPOINT ["/bin/bash"] ENV NVIDIA_DRIVER_CAPABILITIES all \ No newline at end of file diff --git a/docker/inference_img b/docker/inference_img new file mode 100644 index 0000000..5557be4 --- /dev/null +++ b/docker/inference_img @@ -0,0 +1,2 @@ +#!/bin/sh +python3 /rife/inference_img.py $@ diff --git a/docker/rife.sh b/docker/inference_video similarity index 100% rename from docker/rife.sh rename to docker/inference_video From f455a0573656fcf95da54a538341f1672c556126 Mon Sep 17 00:00:00 2001 From: ko1N Date: Sun, 17 Jan 2021 16:07:54 +0100 Subject: [PATCH 3/5] Dockerfile now uses the same train_log folder as the scripts --- README.md | 2 +- docker/Dockerfile | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 6e377f1..65e008e 100644 --- a/README.md +++ b/README.md @@ -81,7 +81,7 @@ ffmpeg -r 10 -f image2 -i output/img%d.png -s 448x256 -vf "split[s0][s1];[s0]pal ``` ### Run in docker -Place the pre-trained models in the `./docker/pretrained_models directory` +Place the pre-trained models in `train_log/\*.pkl` (as above) Building the container: ``` diff --git a/docker/Dockerfile b/docker/Dockerfile index eca0589..801dbb7 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -15,7 +15,7 @@ ADD docker/inference_video /usr/local/bin/inference_video RUN chmod +x /usr/local/bin/inference_video # add pre-trained models -COPY docker/pretrained_models /rife/train_log +COPY train_log /rife/train_log WORKDIR /host ENTRYPOINT ["/bin/bash"] From dd21a4bc9f1c14d1d5a676042fd7184dae2f0a0b Mon Sep 17 00:00:00 2001 From: ko1N Date: Sun, 17 Jan 2021 16:11:24 +0100 Subject: [PATCH 4/5] Updated docker gpu instructions --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index 65e008e..aaf2847 100644 --- a/README.md +++ b/README.md @@ -96,6 +96,11 @@ docker run --rm -it -v $PWD:/host rife:latest inference_video --exp=1 --video=un docker run --rm -it -v $PWD:/host rife:latest inference_img --img img0.png img1.png --exp=4 ``` +Using gpu acceleration (requires proper gpu drivers for docker): +``` +docker run --rm -it --gpus all -v /dev/dri:/dev/dri -v $PWD:/host rife:latest inference_video --exp=1 --video=untitled.mp4 --output=untitled_rife.mp4 +``` + ## Evaluation Download [RIFE model](https://drive.google.com/file/d/1c1R7iF-ypN6USo-D2YH_ORtaH3tukSlo/view?usp=sharing) or [RIFE2F1.5C model](https://drive.google.com/file/d/1ve9w-cRWotdvvbU1KcgtsSm12l-JUkeT/view?usp=sharing) reported by our paper. From 166d10e9a723584a15f27b6b88a5f2f24f1a50d5 Mon Sep 17 00:00:00 2001 From: ko1N Date: Sun, 17 Jan 2021 17:34:48 +0100 Subject: [PATCH 5/5] Fixed audio merge when using abs paths --- inference_video.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/inference_video.py b/inference_video.py index 28fd756..7b5d56c 100644 --- a/inference_video.py +++ b/inference_video.py @@ -27,25 +27,26 @@ def transferAudio(sourceVideo, targetVideo): os.makedirs("temp") # extract audio from video os.system("ffmpeg -y -i " + sourceVideo + " -c:a copy -vn " + tempAudioFileName) - - os.rename(targetVideo, "noAudio_"+targetVideo) + + targetNoAudio = os.path.splitext(targetVideo)[0] + "_noaudio" + os.path.splitext(targetVideo)[1] + os.rename(targetVideo, targetNoAudio) # combine audio file and new video file - os.system("ffmpeg -y -i " + "noAudio_"+targetVideo + " -i " + tempAudioFileName + " -c copy " + targetVideo) + os.system("ffmpeg -y -i " + targetNoAudio + " -i " + tempAudioFileName + " -c copy " + targetVideo) if os.path.getsize(targetVideo) == 0: # if ffmpeg failed to merge the video and audio together try converting the audio to aac tempAudioFileName = "./temp/audio.m4a" os.system("ffmpeg -y -i " + sourceVideo + " -c:a aac -b:a 160k -vn " + tempAudioFileName) - os.system("ffmpeg -y -i " + "noAudio_"+targetVideo + " -i " + tempAudioFileName + " -c copy " + targetVideo) + os.system("ffmpeg -y -i " + targetNoAudio + " -i " + tempAudioFileName + " -c copy " + targetVideo) if (os.path.getsize(targetVideo) == 0): # if aac is not supported by selected format - os.rename("noAudio_"+targetVideo, targetVideo) + os.rename(targetNoAudio, targetVideo) print("Audio transfer failed. Interpolated video will have no audio") else: print("Lossless audio transfer failed. Audio was transcoded to AAC (M4A) instead.") # remove audio-less video - os.remove("noAudio_"+targetVideo) + os.remove(targetNoAudio) else: - os.remove("noAudio_"+targetVideo) + os.remove(targetNoAudio) # remove temp directory shutil.rmtree("temp") @@ -221,4 +222,5 @@ if args.png == False and fpsNotAssigned == True and not args.skip and not args.v transferAudio(args.video, vid_out_name) except: print("Audio transfer failed. Interpolated video will have no audio") - os.rename("noAudio_"+vid_out_name, vid_out_name) + targetNoAudio = os.path.splitext(vid_out_name)[0] + "_noaudio" + os.path.splitext(vid_out_name)[1] + os.rename(targetNoAudio, vid_out_name)