Upload folder using huggingface_hub

.gitattributes

@@ -126,3 +126,4 @@ assets/assets/models/lsk-shell.glb filter=lfs diff=lfs merge=lfs -text
 assets/assets/models/lsk.glb filter=lfs diff=lfs merge=lfs -text
 assets/assets/models/outside.glb filter=lfs diff=lfs merge=lfs -text
 assets/t2v/couple.gif filter=lfs diff=lfs merge=lfs -text
+ECCV2022-RIFE-main/demo/I0_slomo_clipped.gif filter=lfs diff=lfs merge=lfs -text

ECCV2022-RIFE-main/.gitignore

@@ -0,0 +1,14 @@
+*.pyc
+*.py~
+*.py#
+
+*.pkl
+output/*
+train_log/*
+*.mp4
+
+test/
+.idea/
+*.npz
+
+*.zip

ECCV2022-RIFE-main/Colab_demo.ipynb

@@ -0,0 +1,125 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/hzwer/arXiv2020-RIFE/blob/main/Colab_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "FypCcZkNNt2p"
+      },
+      "outputs": [],
+      "source": [
+        "!git clone https://github.com/hzwer/arXiv2020-RIFE"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1wysVHxoN54f"
+      },
+      "outputs": [],
+      "source": [
+        "!mkdir /content/arXiv2020-RIFE/train_log\n",
+        "%cd /content/arXiv2020-RIFE/train_log\n",
+        "!gdown --id 1APIzVeI-4ZZCEuIRE1m6WYfSCaOsi_7_\n",
+        "!7z e RIFE_trained_model_v3.6.zip"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "AhbHfRBJRAUt"
+      },
+      "outputs": [],
+      "source": [
+        "%cd /content/arXiv2020-RIFE/\n",
+        "!gdown --id 1i3xlKb7ax7Y70khcTcuePi6E7crO_dFc\n",
+        "!pip install scikit-video"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "rirngW5uRMdg"
+      },
+      "source": [
+        "Please upload your video to content/arXiv2020-RIFE/video.mp4, or use our demo video."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dnLn4aHHPzN3"
+      },
+      "outputs": [],
+      "source": [
+        "!nvidia-smi\n",
+        "!python3 inference_video.py --exp=2 --video=demo.mp4 --montage"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77KK6lxHgJhf"
+      },
+      "source": [
+        "Our demo.mp4 is 25FPS. You can adjust the parameters for your own perference.\n",
+        "For example: \n",
+        "--fps=60 --exp=1 --video=mydemo.avi --png"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "code",
+        "id": "0zIBbVE3UfUD"
+      },
+      "outputs": [],
+      "source": [
+        "from IPython.display import display, Image\n",
+        "import moviepy.editor as mpy\n",
+        "display(mpy.ipython_display('demo_4X_100fps.mp4', height=256, max_duration=100.))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "tWkJCNgP3zXA"
+      },
+      "outputs": [],
+      "source": [
+        "!python3 inference_img.py --img demo/I0_0.png demo/I0_1.png\n",
+        "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\n",
+        "# Image interpolation"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "include_colab_link": true,
+      "name": "Untitled0.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

ECCV2022-RIFE-main/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright  (c)  Megvii  Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ECCV2022-RIFE-main/README.md

@@ -0,0 +1,195 @@
+# Real-Time Intermediate Flow Estimation for Video Frame Interpolation
+## Introduction
+This project is the implement of [Real-Time Intermediate Flow Estimation for Video Frame Interpolation](https://arxiv.org/abs/2011.06294). Currently, our model can run 30+FPS for 2X 720p interpolation on a 2080Ti GPU. It supports arbitrary-timestep interpolation between a pair of images.
+
+**2023.11 - We recently release new [v4.7-4.10](https://github.com/hzwer/Practical-RIFE/tree/main#model-list) optimized for anime scenes!** 🎉 We draw from [SAFA](https://github.com/megvii-research/WACV2024-SAFA/tree/main)’s research.
+
+2022.7.4 - Our paper is accepted by ECCV2022. Thanks to all relevant authors, contributors and users!
+
+From 2020 to 2022, we submitted RIFE for five submissions（rejected by CVPR21 ICCV21 AAAI22 CVPR22). Thanks to all anonymous reviewers, your suggestions have helped to significantly improve the paper! -> [author website](https://github.com/hzwer)
+
+[ECCV Poster](https://drive.google.com/file/d/1xCXuLUCSwhN61kvIF8jxDvQiUGtLK0kN/view?usp=sharing) | [ECCV 5-min presentation](https://youtu.be/qdp-NYqWQpA) | [论文中文介绍](https://zhuanlan.zhihu.com/p/568553080) | [rebuttal (2WA1WR->3WA)](https://drive.google.com/file/d/16IVjwRpwbTuJbYyTn4PizKX8I257QxY-/view?usp=sharing)
+
+## [YouTube](https://www.youtube.com/results?search_query=rife+interpolation&sp=CAM%253D) | [BiliBili](https://search.bilibili.com/all?keyword=SVFI&order=stow&duration=0&tids_1=0) | [Colab](https://colab.research.google.com/github/hzwer/ECCV2022-RIFE/blob/main/Colab_demo.ipynb) | [Tutorial](https://www.youtube.com/watch?v=gf_on-dbwyU&feature=emb_title) | [V2EX](https://www.v2ex.com/t/984548)
+
+**Pinned Software: [RIFE-App](https://grisk.itch.io/rife-app) | [FlowFrames](https://nmkd.itch.io/flowframes) | [SVFI (中文)](https://github.com/YiWeiHuang-stack/Squirrel-Video-Frame-Interpolation)**
+
+16X interpolation results from two input images: 
+
+![Demo](./demo/I2_slomo_clipped.gif)
+![Demo](./demo/D2_slomo_clipped.gif)
+
+## Software
+[Flowframes](https://nmkd.itch.io/flowframes) | [SVFI(中文)](https://github.com/YiWeiHuang-stack/Squirrel-Video-Frame-Interpolation) | [Waifu2x-Extension-GUI](https://github.com/AaronFeng753/Waifu2x-Extension-GUI) | [Autodesk Flame](https://vimeo.com/505942142) | [SVP](https://www.svp-team.com/wiki/RIFE_AI_interpolation) | [MPV_lazy](https://github.com/hooke007/MPV_lazy) | [enhancr](https://github.com/mafiosnik777/enhancr)
+
+[RIFE-App(Paid)](https://grisk.itch.io/rife-app) | [Steam-VFI(Paid)](https://store.steampowered.com/app/1692080/SVFI/) 
+
+We are not responsible for and participating in the development of above software. According to the open source license, we respect the commercial behavior of other developers.
+
+[VapourSynth-RIFE](https://github.com/HolyWu/vs-rife) | [RIFE-ncnn-vulkan](https://github.com/nihui/rife-ncnn-vulkan) | [VapourSynth-RIFE-ncnn-Vulkan](https://github.com/HomeOfVapourSynthEvolution/VapourSynth-RIFE-ncnn-Vulkan) 
+
+<img src="https://api.star-history.com/svg?repos=megvii-research/ECCV2022-RIFE,Justin62628/Squirrel-RIFE,n00mkrad/flowframes,nihui/rife-ncnn-vulkan,hzwer/Practical-RIFE&type=Date" height="320" width="480" />
+
+If you are a developer, welcome to follow [Practical-RIFE](https://github.com/hzwer/Practical-RIFE), which aims to make RIFE more practical for users by adding various features and design new models with faster speed.
+
+You may check [this pull request](https://github.com/megvii-research/ECCV2022-RIFE/pull/300) for supporting macOS.
+## CLI Usage
+
+### Installation
+
+```
+git clone git@github.com:megvii-research/ECCV2022-RIFE.git
+cd ECCV2022-RIFE
+pip3 install -r requirements.txt
+```
+
+* Download the pretrained **HD** models from [here](https://drive.google.com/file/d/1APIzVeI-4ZZCEuIRE1m6WYfSCaOsi_7_/view?usp=sharing). (百度网盘链接:https://pan.baidu.com/share/init?surl=u6Q7-i4Hu4Vx9_5BJibPPA 密码:hfk3，把压缩包解开后放在 train_log/\*)
+
+* Unzip and move the pretrained parameters to train_log/\*
+
+* This model is not reported by our paper, for our paper model please refer to [evaluation](https://github.com/hzwer/ECCV2022-RIFE#evaluation).
+
+### Run
+
+**Video Frame Interpolation**
+
+You can use our [demo video](https://drive.google.com/file/d/1i3xlKb7ax7Y70khcTcuePi6E7crO_dFc/view?usp=sharing) or your own video. 
+```
+python3 inference_video.py --exp=1 --video=video.mp4 
+```
+(generate video_2X_xxfps.mp4)
+```
+python3 inference_video.py --exp=2 --video=video.mp4
+```
+(for 4X interpolation)
+```
+python3 inference_video.py --exp=1 --video=video.mp4 --scale=0.5
+```
+(If your video has very high resolution such as 4K, we recommend set --scale=0.5 (default 1.0). If you generate disordered pattern on your videos, try set --scale=2.0. This parameter control the process resolution for optical flow model.)
+```
+python3 inference_video.py --exp=2 --img=input/
+```
+(to read video from pngs, like input/0.png ... input/612.png, ensure that the png names are numbers)
+```
+python3 inference_video.py --exp=2 --video=video.mp4 --fps=60
+```
+(add slomo effect, the audio will be removed)
+```
+python3 inference_video.py --video=video.mp4 --montage --png
+```
+(if you want to montage the origin video and save the png format output)
+
+**Optical Flow Estimation**
+
+You may refer to [#278](https://github.com/megvii-research/ECCV2022-RIFE/issues/278#event-7199085190).
+
+**Image Interpolation**
+
+```
+python3 inference_img.py --img img0.png img1.png --exp=4
+```
+(2^4=16X interpolation results)
+After that, you can use pngs to generate mp4:
+```
+ffmpeg -r 10 -f image2 -i output/img%d.png -s 448x256 -c:v libx264 -pix_fmt yuv420p output/slomo.mp4 -q:v 0 -q:a 0
+```
+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 `train_log/\*.pkl` (as above)
+
+Building the container:
+```
+docker build -t rife -f docker/Dockerfile .
+```
+
+Running the container:
+```
+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
+```
+
+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/1h42aGYPNJn2q8j_GVkS_yDu__G_UZ2GX/view?usp=sharing) or [RIFE_m model](https://drive.google.com/file/d/147XVsDXBfJPlyct2jfo9kpbL944mNeZr/view?usp=sharing) reported by our paper.
+
+**UCF101**: Download [UCF101 dataset](https://liuziwei7.github.io/projects/VoxelFlow) at ./UCF101/ucf101_interp_ours/
+
+**Vimeo90K**: Download [Vimeo90K dataset](http://toflow.csail.mit.edu/) at ./vimeo_interp_test
+
+**MiddleBury**: Download [MiddleBury OTHER dataset](https://vision.middlebury.edu/flow/data/) at ./other-data and ./other-gt-interp
+
+**HD**: Download [HD dataset](https://github.com/baowenbo/MEMC-Net) at ./HD_dataset. We also provide a [google drive download link](https://drive.google.com/file/d/1iHaLoR2g1-FLgr9MEv51NH_KQYMYz-FA/view?usp=sharing).
+```
+# RIFE
+python3 benchmark/UCF101.py
+# "PSNR: 35.282 SSIM: 0.9688"
+python3 benchmark/Vimeo90K.py
+# "PSNR: 35.615 SSIM: 0.9779"
+python3 benchmark/MiddleBury_Other.py
+# "IE: 1.956"
+python3 benchmark/HD.py
+# "PSNR: 32.14"
+
+# RIFE_m
+python3 benchmark/HD_multi_4X.py
+# "PSNR: 22.96(544*1280), 31.87(720p), 34.25(1080p)"
+```
+
+## Training and Reproduction
+Download [Vimeo90K dataset](http://toflow.csail.mit.edu/).
+
+We use 16 CPUs, 4 GPUs and 20G memory for training: 
+```
+python3 -m torch.distributed.launch --nproc_per_node=4 train.py --world_size=4
+```
+
+## Revision History
+
+2021.3.18 [arXiv](https://arxiv.org/pdf/2011.06294v5.pdf): Modify the main experimental data, especially the runtime related issues.
+
+2021.8.12 [arXiv](https://arxiv.org/pdf/2011.06294v6.pdf): Remove pre-trained model dependency and propose privileged distillation scheme for frame interpolation. Remove [census loss](https://github.com/hzwer/arXiv2021-RIFE/blob/0e241367847a0895748e64c6e1604c94db54d395/model/loss.py#L20) supervision.
+
+2021.11.17 [arXiv](https://arxiv.org/pdf/2011.06294v11.pdf): Support arbitrary-time frame interpolation, aka RIFEm and add more experiments.
+
+## Recommend
+We sincerely recommend some related papers:
+
+CVPR22 - [Optimizing Video Prediction via Video Frame Interpolation](https://openaccess.thecvf.com/content/CVPR2022/html/Wu_Optimizing_Video_Prediction_via_Video_Frame_Interpolation_CVPR_2022_paper.html)
+
+CVPR22 - [Video Frame Interpolation with Transformer](https://openaccess.thecvf.com/content/CVPR2022/html/Lu_Video_Frame_Interpolation_With_Transformer_CVPR_2022_paper.html)
+
+CVPR22 - [IFRNet: Intermediate Feature Refine Network for Efficient Frame Interpolation](https://openaccess.thecvf.com/content/CVPR2022/html/Kong_IFRNet_Intermediate_Feature_Refine_Network_for_Efficient_Frame_Interpolation_CVPR_2022_paper.html)
+
+CVPR23 - [A Dynamic Multi-Scale Voxel Flow Network for Video Prediction](https://huxiaotaostasy.github.io/DMVFN/)
+
+CVPR23 - [Extracting Motion and Appearance via Inter-Frame Attention for Efficient Video Frame Interpolation](https://arxiv.org/abs/2303.00440)
+
+## Citation
+If you think this project is helpful, please feel free to leave a star or cite our paper:
+
+```
+@inproceedings{huang2022rife,
+  title={Real-Time Intermediate Flow Estimation for Video Frame Interpolation},
+  author={Huang, Zhewei and Zhang, Tianyuan and Heng, Wen and Shi, Boxin and Zhou, Shuchang},
+  booktitle={Proceedings of the European Conference on Computer Vision (ECCV)},
+  year={2022}
+}
+```
+
+## Reference
+
+Optical Flow:
+[ARFlow](https://github.com/lliuz/ARFlow)  [pytorch-liteflownet](https://github.com/sniklaus/pytorch-liteflownet)  [RAFT](https://github.com/princeton-vl/RAFT)  [pytorch-PWCNet](https://github.com/sniklaus/pytorch-pwc)
+
+Video Interpolation: 
+[DVF](https://github.com/lxx1991/pytorch-voxel-flow)  [TOflow](https://github.com/Coldog2333/pytoflow)  [SepConv](https://github.com/sniklaus/sepconv-slomo)  [DAIN](https://github.com/baowenbo/DAIN)  [CAIN](https://github.com/myungsub/CAIN)  [MEMC-Net](https://github.com/baowenbo/MEMC-Net)   [SoftSplat](https://github.com/sniklaus/softmax-splatting)  [BMBC](https://github.com/JunHeum/BMBC)  [EDSC](https://github.com/Xianhang/EDSC-pytorch)  [EQVI](https://github.com/lyh-18/EQVI)

ECCV2022-RIFE-main/benchmark/ATD12K.py

@@ -0,0 +1,42 @@
+import os
+import sys
+sys.path.append('.')
+import cv2
+import math
+import torch
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from model.pytorch_msssim import ssim_matlab
+from model.RIFE import Model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model()
+model.load_model('train_log')
+model.eval()
+model.device()
+
+path = 'datasets/test_2k_540p/'
+dirs = os.listdir(path)
+psnr_list = []
+ssim_list = []
+print(len(dirs))
+for d in dirs:
+    img0 = (path + d + '/frame1.png')
+    img1 = (path + d + '/frame3.png')
+    gt = (path + d + '/frame2.png')
+    img0 = (torch.tensor(cv2.imread(img0).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0)
+    img1 = (torch.tensor(cv2.imread(img1).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0)
+    gt = (torch.tensor(cv2.imread(gt).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0)
+    pader = torch.nn.ReplicationPad2d([0, 0, 2, 2])
+    img0 = pader(img0)
+    img1 = pader(img1)
+    pred = model.inference(img0, img1)[0][:, 2:-2]
+    ssim = ssim_matlab(gt, torch.round(pred * 255).unsqueeze(0) / 255.).detach().cpu().numpy()
+    out = pred.detach().cpu().numpy().transpose(1, 2, 0)
+    out = np.round(out * 255) / 255.
+    gt = gt[0].cpu().numpy().transpose(1, 2, 0)
+    psnr = -10 * math.log10(((gt - out) * (gt - out)).mean())
+    psnr_list.append(psnr)
+    ssim_list.append(ssim)
+    print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list)))

ECCV2022-RIFE-main/benchmark/HD.py

@@ -0,0 +1,89 @@
+import os
+import sys
+sys.path.append('.')
+import cv2
+import math
+import torch
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from model.pytorch_msssim import ssim_matlab
+from model.RIFE import Model
+from skimage.color import rgb2yuv, yuv2rgb
+from yuv_frame_io import YUV_Read,YUV_Write
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model()
+model.load_model('train_log')
+model.eval()
+model.device()
+
+name_list = [
+    ('HD_dataset/HD720p_GT/parkrun_1280x720_50.yuv', 720, 1280),
+    ('HD_dataset/HD720p_GT/shields_1280x720_60.yuv', 720, 1280),
+    ('HD_dataset/HD720p_GT/stockholm_1280x720_60.yuv', 720, 1280),
+    ('HD_dataset/HD1080p_GT/BlueSky.yuv', 1080, 1920),
+    ('HD_dataset/HD1080p_GT/Kimono1_1920x1080_24.yuv', 1080, 1920),
+    ('HD_dataset/HD1080p_GT/ParkScene_1920x1080_24.yuv', 1080, 1920),
+    ('HD_dataset/HD1080p_GT/sunflower_1080p25.yuv', 1080, 1920),
+    ('HD_dataset/HD544p_GT/Sintel_Alley2_1280x544.yuv', 544, 1280),
+    ('HD_dataset/HD544p_GT/Sintel_Market5_1280x544.yuv', 544, 1280),
+    ('HD_dataset/HD544p_GT/Sintel_Temple1_1280x544.yuv', 544, 1280),
+    ('HD_dataset/HD544p_GT/Sintel_Temple2_1280x544.yuv', 544, 1280),
+]
+tot = 0.
+for data in name_list:
+    psnr_list = []
+    name = data[0]
+    h = data[1]
+    w = data[2]
+    if 'yuv' in name:
+        Reader = YUV_Read(name, h, w, toRGB=True)
+    else:
+        Reader = cv2.VideoCapture(name)
+    _, lastframe = Reader.read()
+    # fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
+    # video = cv2.VideoWriter(name + '.mp4', fourcc, 30, (w, h))
+    for index in range(0, 100, 2):
+        if 'yuv' in name:
+            IMAGE1, success1 = Reader.read(index)
+            gt, _ = Reader.read(index + 1)
+            IMAGE2, success2 = Reader.read(index + 2)
+            if not success2:
+                break
+        else:
+            success1, gt = Reader.read()
+            success2, frame = Reader.read()
+            IMAGE1 = lastframe
+            IMAGE2 = frame
+            lastframe = frame
+            if not success2:
+                break
+        I0 = torch.from_numpy(np.transpose(IMAGE1, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0)
+        I1 = torch.from_numpy(np.transpose(IMAGE2, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0)
+        
+        if h == 720:
+            pad = 24
+        elif h == 1080:
+            pad = 4
+        else:
+            pad = 16
+        pader = torch.nn.ReplicationPad2d([0, 0, pad, pad])
+        I0 = pader(I0)
+        I1 = pader(I1)
+        with torch.no_grad():
+            pred = model.inference(I0, I1)
+            pred = pred[:, :, pad: -pad]
+        out = (np.round(pred[0].detach().cpu().numpy().transpose(1, 2, 0) * 255)).astype('uint8')
+        # video.write(out)
+        if 'yuv' in name:
+            diff_rgb = 128.0 + rgb2yuv(gt / 255.)[:, :, 0] * 255 - rgb2yuv(out / 255.)[:, :, 0] * 255
+            mse = np.mean((diff_rgb - 128.0) ** 2)
+            PIXEL_MAX = 255.0
+            psnr = 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
+        else:
+            psnr = skim.compare_psnr(gt, out)
+        psnr_list.append(psnr)
+    print(np.mean(psnr_list))
+    tot += np.mean(psnr_list)
+print('avg psnr', tot / len(name_list))

ECCV2022-RIFE-main/benchmark/HD_multi_4X.py

@@ -0,0 +1,105 @@
+import os
+import sys
+sys.path.append('.')
+import cv2
+import math
+import torch
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from model.pytorch_msssim import ssim_matlab
+from model.RIFE import Model
+from skimage.color import rgb2yuv, yuv2rgb
+from yuv_frame_io import YUV_Read,YUV_Write
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model(arbitrary=True)
+model.load_model('RIFE_m_train_log')
+model.eval()
+model.device()
+
+name_list = [
+    ('HD_dataset/HD720p_GT/parkrun_1280x720_50.yuv', 720, 1280),
+    ('HD_dataset/HD720p_GT/shields_1280x720_60.yuv', 720, 1280),
+    ('HD_dataset/HD720p_GT/stockholm_1280x720_60.yuv', 720, 1280),
+    ('HD_dataset/HD1080p_GT/BlueSky.yuv', 1080, 1920),
+    ('HD_dataset/HD1080p_GT/Kimono1_1920x1080_24.yuv', 1080, 1920),
+    ('HD_dataset/HD1080p_GT/ParkScene_1920x1080_24.yuv', 1080, 1920),
+    ('HD_dataset/HD1080p_GT/sunflower_1080p25.yuv', 1080, 1920),
+    ('HD_dataset/HD544p_GT/Sintel_Alley2_1280x544.yuv', 544, 1280),
+    ('HD_dataset/HD544p_GT/Sintel_Market5_1280x544.yuv', 544, 1280),
+    ('HD_dataset/HD544p_GT/Sintel_Temple1_1280x544.yuv', 544, 1280),
+    ('HD_dataset/HD544p_GT/Sintel_Temple2_1280x544.yuv', 544, 1280),
+]
+def inference(I0, I1, pad, multi=2, arbitrary=True):
+    img = [I0, I1]
+    if not arbitrary:
+        for i in range(multi):
+            res = [I0]
+            for j in range(len(img) - 1):
+                res.append(model.inference(img[j], img[j + 1]))
+                res.append(img[j + 1])
+            img = res
+    else:
+        img = [I0]
+        p = 2**multi
+        for i in range(p-1):
+            img.append(model.inference(I0, I1, timestep=(i+1)*(1./p)))
+        img.append(I1)
+    for i in range(len(img)):
+        img[i] = img[i][0][:, pad: -pad]
+    return img[1: -1]
+        
+tot = []
+for data in name_list:
+    psnr_list = []
+    name = data[0]
+    h = data[1]
+    w = data[2]
+    if 'yuv' in name:
+        Reader = YUV_Read(name, h, w, toRGB=True)
+    else:
+        Reader = cv2.VideoCapture(name)
+    _, lastframe = Reader.read()
+    # fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
+    # video = cv2.VideoWriter(name + '.mp4', fourcc, 30, (w, h))    
+    for index in range(0, 100, 4):
+        gt = []
+        if 'yuv' in name:
+            IMAGE1, success1 = Reader.read(index)
+            IMAGE2, success2 = Reader.read(index + 4)
+            if not success2:
+                break
+            for i in range(1, 4):
+                tmp, _ = Reader.read(index + i)
+                gt.append(tmp)
+        else:
+            print('Not Implement')
+        I0 = torch.from_numpy(np.transpose(IMAGE1, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0)
+        I1 = torch.from_numpy(np.transpose(IMAGE2, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0)
+        
+        if h == 720:
+            pad = 24
+        elif h == 1080:
+            pad = 4
+        else:
+            pad = 16
+        pader = torch.nn.ReplicationPad2d([0, 0, pad, pad])
+        I0 = pader(I0)
+        I1 = pader(I1)
+        with torch.no_grad():
+            pred = inference(I0, I1, pad)
+        for i in range(4 - 1):
+            out = (np.round(pred[i].detach().cpu().numpy().transpose(1, 2, 0) * 255)).astype('uint8')
+            if 'yuv' in name:
+                diff_rgb = 128.0 + rgb2yuv(gt[i] / 255.)[:, :, 0] * 255 - rgb2yuv(out / 255.)[:, :, 0] * 255
+                mse = np.mean((diff_rgb - 128.0) ** 2)
+                PIXEL_MAX = 255.0
+                psnr = 20 * math.log10(PIXEL_MAX / math.sqrt(mse))
+            else:
+                print('Not Implement')
+            psnr_list.append(psnr)
+    print(np.mean(psnr_list))
+    tot.append(np.mean(psnr_list))
+
+print('PSNR: {}(544*1280), {}(720p), {}(1080p)'.format(np.mean(tot[7:11]), np.mean(tot[:3]), np.mean(tot[3:7])))

ECCV2022-RIFE-main/benchmark/MiddleBury_Other.py

@@ -0,0 +1,37 @@
+import os
+import sys
+sys.path.append('.')
+import cv2
+import math
+import torch
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from model.pytorch_msssim import ssim_matlab
+from model.RIFE import Model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model()
+model.load_model('train_log')
+model.eval()
+model.device()
+
+name = ['Beanbags', 'Dimetrodon', 'DogDance', 'Grove2', 'Grove3', 'Hydrangea', 'MiniCooper', 'RubberWhale', 'Urban2', 'Urban3', 'Venus', 'Walking']
+IE_list = []
+for i in name:
+    i0 = cv2.imread('other-data/{}/frame10.png'.format(i)).transpose(2, 0, 1) / 255.
+    i1 = cv2.imread('other-data/{}/frame11.png'.format(i)).transpose(2, 0, 1) / 255.
+    gt = cv2.imread('other-gt-interp/{}/frame10i11.png'.format(i)) 
+    h, w = i0.shape[1], i0.shape[2]
+    imgs = torch.zeros([1, 6, 480, 640]).to(device)
+    ph = (480 - h) // 2
+    pw = (640 - w) // 2
+    imgs[:, :3, :h, :w] = torch.from_numpy(i0).unsqueeze(0).float().to(device)
+    imgs[:, 3:, :h, :w] = torch.from_numpy(i1).unsqueeze(0).float().to(device)
+    I0 = imgs[:, :3]
+    I2 = imgs[:, 3:]
+    pred = model.inference(I0, I2)
+    out = pred[0].detach().cpu().numpy().transpose(1, 2, 0)
+    out = np.round(out[:h, :w] * 255)
+    IE_list.append(np.abs((out - gt * 1.0)).mean())
+    print(np.mean(IE_list))

ECCV2022-RIFE-main/benchmark/UCF101.py

@@ -0,0 +1,39 @@
+import os
+import sys
+sys.path.append('.')
+import cv2
+import math
+import torch
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from model.pytorch_msssim import ssim_matlab
+from model.RIFE import Model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model()
+model.load_model('train_log')
+model.eval()
+model.device()
+
+path = 'UCF101/ucf101_interp_ours/'
+dirs = os.listdir(path)
+psnr_list = []
+ssim_list = []
+print(len(dirs))
+for d in dirs:
+    img0 = (path + d + '/frame_00.png')
+    img1 = (path + d + '/frame_02.png')
+    gt = (path + d + '/frame_01_gt.png')
+    img0 = (torch.tensor(cv2.imread(img0).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0)
+    img1 = (torch.tensor(cv2.imread(img1).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0)
+    gt = (torch.tensor(cv2.imread(gt).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0)
+    pred = model.inference(img0, img1)[0]
+    ssim = ssim_matlab(gt, torch.round(pred * 255).unsqueeze(0) / 255.).detach().cpu().numpy()
+    out = pred.detach().cpu().numpy().transpose(1, 2, 0)
+    out = np.round(out * 255) / 255.
+    gt = gt[0].cpu().numpy().transpose(1, 2, 0)
+    psnr = -10 * math.log10(((gt - out) * (gt - out)).mean())
+    psnr_list.append(psnr)
+    ssim_list.append(ssim)
+    print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list)))

ECCV2022-RIFE-main/benchmark/Vimeo90K.py

@@ -0,0 +1,40 @@
+import os
+import sys
+sys.path.append('.')
+import cv2
+import math
+import torch
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from model.pytorch_msssim import ssim_matlab
+from model.RIFE import Model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = Model()
+model.load_model('train_log')
+model.eval()
+model.device()
+
+path = 'vimeo_interp_test/'
+f = open(path + 'tri_testlist.txt', 'r')
+psnr_list = []
+ssim_list = []
+for i in f:
+    name = str(i).strip()
+    if(len(name) <= 1):
+        continue
+    print(path + 'target/' + name + '/im1.png')
+    I0 = cv2.imread(path + 'target/' + name + '/im1.png')
+    I1 = cv2.imread(path + 'target/' + name + '/im2.png')
+    I2 = cv2.imread(path + 'target/' + name + '/im3.png')
+    I0 = (torch.tensor(I0.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0)
+    I2 = (torch.tensor(I2.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0)
+    mid = model.inference(I0, I2)[0]
+    ssim = ssim_matlab(torch.tensor(I1.transpose(2, 0, 1)).to(device).unsqueeze(0) / 255., torch.round(mid * 255).unsqueeze(0) / 255.).detach().cpu().numpy()
+    mid = np.round((mid * 255).detach().cpu().numpy()).astype('uint8').transpose(1, 2, 0) / 255.    
+    I1 = I1 / 255.
+    psnr = -10 * math.log10(((I1 - mid) * (I1 - mid)).mean())
+    psnr_list.append(psnr)
+    ssim_list.append(ssim)
+    print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list)))

ECCV2022-RIFE-main/benchmark/testtime.py

@@ -0,0 +1,29 @@
+import cv2
+import sys
+sys.path.append('.')
+import time
+import torch
+import torch.nn as nn
+from model.RIFE import Model
+
+model = Model()
+model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_grad_enabled(False)
+if torch.cuda.is_available():
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.benchmark = True
+       
+I0 = torch.rand(1, 3, 480, 640).to(device)
+I1 = torch.rand(1, 3, 480, 640).to(device)
+with torch.no_grad():
+    for i in range(100):
+        pred = model.inference(I0, I1)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    time_stamp = time.time()
+    for i in range(100):
+        pred = model.inference(I0, I1)
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+    print((time.time() - time_stamp) / 100)

ECCV2022-RIFE-main/benchmark/yuv_frame_io.py

@@ -0,0 +1,124 @@
+import sys
+import getopt
+import math
+import numpy
+import random
+import logging
+import numpy as np
+from skimage.color import rgb2yuv, yuv2rgb
+from PIL import Image
+import os
+from shutil import copyfile
+
+class YUV_Read():
+    def __init__(self, filepath, h, w, format='yuv420', toRGB=True):
+
+        self.h = h
+        self.w = w
+
+        self.fp = open(filepath, 'rb')
+
+        if format == 'yuv420':
+            self.frame_length = int(1.5 * h * w)
+            self.Y_length = h * w
+            self.Uv_length = int(0.25 * h * w)
+        else:
+            pass
+        self.toRGB = toRGB
+
+    def read(self, offset_frame=None):
+        if not offset_frame == None:
+            self.fp.seek(offset_frame * self.frame_length, 0)
+
+        Y = np.fromfile(self.fp, np.uint8, count=self.Y_length)
+        U = np.fromfile(self.fp, np.uint8, count=self.Uv_length)
+        V = np.fromfile(self.fp, np.uint8, count=self.Uv_length)
+        if Y.size < self.Y_length or \
+                        U.size < self.Uv_length or \
+                        V.size < self.Uv_length:
+            return None, False
+
+        Y = np.reshape(Y, [self.w, self.h], order='F')
+        Y = np.transpose(Y)
+
+        U = np.reshape(U, [int(self.w / 2), int(self.h / 2)], order='F')
+        U = np.transpose(U)
+
+        V = np.reshape(V, [int(self.w / 2), int(self.h / 2)], order='F')
+        V = np.transpose(V)
+
+        U = np.array(Image.fromarray(U).resize([self.w, self.h]))
+        V = np.array(Image.fromarray(V).resize([self.w, self.h]))
+
+        if self.toRGB:
+            Y = Y / 255.0
+            U = U / 255.0 - 0.5
+            V = V / 255.0 - 0.5
+
+            self.YUV = np.stack((Y, U, V), axis=-1)
+            self.RGB = (255.0 * np.clip(yuv2rgb(self.YUV), 0.0, 1.0)).astype('uint8')
+
+            self.YUV = None
+            return self.RGB, True
+        else:
+            self.YUV = np.stack((Y, U, V), axis=-1)
+            return self.YUV, True
+
+    def close(self):
+        self.fp.close()
+
+
+class YUV_Write():
+    def __init__(self, filepath, fromRGB=True):
+        if os.path.exists(filepath):
+            print(filepath)
+  
+        self.fp = open(filepath, 'wb')
+        self.fromRGB = fromRGB
+
+    def write(self, Frame):
+
+        self.h = Frame.shape[0]
+        self.w = Frame.shape[1]
+        c = Frame.shape[2]
+
+        assert c == 3
+        if format == 'yuv420':
+            self.frame_length = int(1.5 * self.h * self.w)
+            self.Y_length = self.h * self.w
+            self.Uv_length = int(0.25 * self.h * self.w)
+        else:
+            pass
+        if self.fromRGB:
+            Frame = Frame / 255.0
+            YUV = rgb2yuv(Frame)
+            Y, U, V = np.dsplit(YUV, 3)
+            Y = Y[:, :, 0]
+            U = U[:, :, 0]
+            V = V[:, :, 0]
+            U = np.clip(U + 0.5, 0.0, 1.0)
+            V = np.clip(V + 0.5, 0.0, 1.0)
+
+            U = U[::2, ::2]  # imresize(U,[int(self.h/2),int(self.w/2)],interp = 'nearest')
+            V = V[::2, ::2]  # imresize(V ,[int(self.h/2),int(self.w/2)],interp = 'nearest')
+            Y = (255.0 * Y).astype('uint8')
+            U = (255.0 * U).astype('uint8')
+            V = (255.0 * V).astype('uint8')
+        else:
+            YUV = Frame
+            Y = YUV[:, :, 0]
+            U = YUV[::2, ::2, 1]
+            V = YUV[::2, ::2, 2]
+
+        Y = Y.flatten()  # the first order is 0-dimension so don't need to transpose before flatten
+        U = U.flatten()
+        V = V.flatten()
+
+        Y.tofile(self.fp)
+        U.tofile(self.fp)
+        V.tofile(self.fp)
+
+        return True
+
+    def close(self):
+        self.fp.close()

ECCV2022-RIFE-main/dataset.py

@@ -0,0 +1,109 @@
+import os
+import cv2
+import ast
+import torch
+import numpy as np
+import random
+from torch.utils.data import DataLoader, Dataset
+
+cv2.setNumThreads(1)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+class VimeoDataset(Dataset):
+    def __init__(self, dataset_name, batch_size=32):
+        self.batch_size = batch_size
+        self.dataset_name = dataset_name        
+        self.h = 256
+        self.w = 448
+        self.data_root = 'vimeo_triplet'
+        self.image_root = os.path.join(self.data_root, 'sequences')
+        train_fn = os.path.join(self.data_root, 'tri_trainlist.txt')
+        test_fn = os.path.join(self.data_root, 'tri_testlist.txt')
+        with open(train_fn, 'r') as f:
+            self.trainlist = f.read().splitlines()
+        with open(test_fn, 'r') as f:
+            self.testlist = f.read().splitlines()   
+        self.load_data()
+
+    def __len__(self):
+        return len(self.meta_data)
+
+    def load_data(self):
+        cnt = int(len(self.trainlist) * 0.95)
+        if self.dataset_name == 'train':
+            self.meta_data = self.trainlist[:cnt]
+        elif self.dataset_name == 'test':
+            self.meta_data = self.testlist
+        else:
+            self.meta_data = self.trainlist[cnt:]
+           
+    def crop(self, img0, gt, img1, h, w):
+        ih, iw, _ = img0.shape
+        x = np.random.randint(0, ih - h + 1)
+        y = np.random.randint(0, iw - w + 1)
+        img0 = img0[x:x+h, y:y+w, :]
+        img1 = img1[x:x+h, y:y+w, :]
+        gt = gt[x:x+h, y:y+w, :]
+        return img0, gt, img1
+
+    def getimg(self, index):
+        imgpath = os.path.join(self.image_root, self.meta_data[index])
+        imgpaths = [imgpath + '/im1.png', imgpath + '/im2.png', imgpath + '/im3.png']
+
+        # Load images
+        img0 = cv2.imread(imgpaths[0])
+        gt = cv2.imread(imgpaths[1])
+        img1 = cv2.imread(imgpaths[2])
+        timestep = 0.5
+        return img0, gt, img1, timestep
+    
+        # RIFEm with Vimeo-Septuplet
+        # imgpaths = [imgpath + '/im1.png', imgpath + '/im2.png', imgpath + '/im3.png', imgpath + '/im4.png', imgpath + '/im5.png', imgpath + '/im6.png', imgpath + '/im7.png']
+        # ind = [0, 1, 2, 3, 4, 5, 6]
+        # random.shuffle(ind)
+        # ind = ind[:3]
+        # ind.sort()
+        # img0 = cv2.imread(imgpaths[ind[0]])
+        # gt = cv2.imread(imgpaths[ind[1]])
+        # img1 = cv2.imread(imgpaths[ind[2]])        
+        # timestep = (ind[1] - ind[0]) * 1.0 / (ind[2] - ind[0] + 1e-6)
+            
+    def __getitem__(self, index):        
+        img0, gt, img1, timestep = self.getimg(index)
+        if self.dataset_name == 'train':
+            img0, gt, img1 = self.crop(img0, gt, img1, 224, 224)
+            if random.uniform(0, 1) < 0.5:
+                img0 = img0[:, :, ::-1]
+                img1 = img1[:, :, ::-1]
+                gt = gt[:, :, ::-1]
+            if random.uniform(0, 1) < 0.5:
+                img0 = img0[::-1]
+                img1 = img1[::-1]
+                gt = gt[::-1]
+            if random.uniform(0, 1) < 0.5:
+                img0 = img0[:, ::-1]
+                img1 = img1[:, ::-1]
+                gt = gt[:, ::-1]
+            if random.uniform(0, 1) < 0.5:
+                tmp = img1
+                img1 = img0
+                img0 = tmp
+                timestep = 1 - timestep
+            # random rotation
+            p = random.uniform(0, 1)
+            if p < 0.25:
+                img0 = cv2.rotate(img0, cv2.ROTATE_90_CLOCKWISE)
+                gt = cv2.rotate(gt, cv2.ROTATE_90_CLOCKWISE)
+                img1 = cv2.rotate(img1, cv2.ROTATE_90_CLOCKWISE)
+            elif p < 0.5:
+                img0 = cv2.rotate(img0, cv2.ROTATE_180)
+                gt = cv2.rotate(gt, cv2.ROTATE_180)
+                img1 = cv2.rotate(img1, cv2.ROTATE_180)
+            elif p < 0.75:
+                img0 = cv2.rotate(img0, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                gt = cv2.rotate(gt, cv2.ROTATE_90_COUNTERCLOCKWISE)
+                img1 = cv2.rotate(img1, cv2.ROTATE_90_COUNTERCLOCKWISE)
+        img0 = torch.from_numpy(img0.copy()).permute(2, 0, 1)
+        img1 = torch.from_numpy(img1.copy()).permute(2, 0, 1)
+        gt = torch.from_numpy(gt.copy()).permute(2, 0, 1)
+        timestep = torch.tensor(timestep).reshape(1, 1, 1)
+        return torch.cat((img0, img1, gt), 0), timestep

ECCV2022-RIFE-main/docker/Dockerfile

@@ -0,0 +1,23 @@
+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/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 train_log /rife/train_log
+
+WORKDIR /host
+ENTRYPOINT ["/bin/bash"]
+
+ENV NVIDIA_DRIVER_CAPABILITIES all

ECCV2022-RIFE-main/docker/inference_img

@@ -0,0 +1,2 @@
+#!/bin/sh
+python3 /rife/inference_img.py $@

ECCV2022-RIFE-main/docker/inference_video

@@ -0,0 +1,2 @@
+#!/bin/sh
+python3 /rife/inference_video.py $@

ECCV2022-RIFE-main/inference_img.py

@@ -0,0 +1,111 @@
+import os
+import cv2
+import torch
+import argparse
+from torch.nn import functional as F
+import warnings
+warnings.filterwarnings("ignore")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_grad_enabled(False)
+if torch.cuda.is_available():
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.benchmark = True
+
+parser = argparse.ArgumentParser(description='Interpolation for a pair of images')
+parser.add_argument('--img', dest='img', nargs=2, required=True)
+parser.add_argument('--exp', default=4, type=int)
+parser.add_argument('--ratio', default=0, type=float, help='inference ratio between two images with 0 - 1 range')
+parser.add_argument('--rthreshold', default=0.02, type=float, help='returns image when actual ratio falls in given range threshold')
+parser.add_argument('--rmaxcycles', default=8, type=int, help='limit max number of bisectional cycles')
+parser.add_argument('--model', dest='modelDir', type=str, default='train_log', help='directory with trained model files')
+
+args = parser.parse_args()
+
+try:
+    try:
+        try:
+            from model.RIFE_HDv2 import Model
+            model = Model()
+            model.load_model(args.modelDir, -1)
+            print("Loaded v2.x HD model.")
+        except:
+            from train_log.RIFE_HDv3 import Model
+            model = Model()
+            model.load_model(args.modelDir, -1)
+            print("Loaded v3.x HD model.")
+    except:
+        from model.RIFE_HD import Model
+        model = Model()
+        model.load_model(args.modelDir, -1)
+        print("Loaded v1.x HD model")
+except:
+    from model.RIFE import Model
+    model = Model()
+    model.load_model(args.modelDir, -1)
+    print("Loaded ArXiv-RIFE model")
+model.eval()
+model.device()
+
+if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
+    img0 = cv2.imread(args.img[0], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
+    img1 = cv2.imread(args.img[1], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
+    img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device)).unsqueeze(0)
+    img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device)).unsqueeze(0)
+
+else:
+    img0 = cv2.imread(args.img[0], cv2.IMREAD_UNCHANGED)
+    img1 = cv2.imread(args.img[1], cv2.IMREAD_UNCHANGED)
+    img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0)
+    img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0)
+
+n, c, h, w = img0.shape
+ph = ((h - 1) // 32 + 1) * 32
+pw = ((w - 1) // 32 + 1) * 32
+padding = (0, pw - w, 0, ph - h)
+img0 = F.pad(img0, padding)
+img1 = F.pad(img1, padding)
+
+
+if args.ratio:
+    img_list = [img0]
+    img0_ratio = 0.0
+    img1_ratio = 1.0
+    if args.ratio <= img0_ratio + args.rthreshold / 2:
+        middle = img0
+    elif args.ratio >= img1_ratio - args.rthreshold / 2:
+        middle = img1
+    else:
+        tmp_img0 = img0
+        tmp_img1 = img1
+        for inference_cycle in range(args.rmaxcycles):
+            middle = model.inference(tmp_img0, tmp_img1)
+            middle_ratio = ( img0_ratio + img1_ratio ) / 2
+            if args.ratio - (args.rthreshold / 2) <= middle_ratio <= args.ratio + (args.rthreshold / 2):
+                break
+            if args.ratio > middle_ratio:
+                tmp_img0 = middle
+                img0_ratio = middle_ratio
+            else:
+                tmp_img1 = middle
+                img1_ratio = middle_ratio
+    img_list.append(middle)
+    img_list.append(img1)
+else:
+    img_list = [img0, img1]
+    for i in range(args.exp):
+        tmp = []
+        for j in range(len(img_list) - 1):
+            mid = model.inference(img_list[j], img_list[j + 1])
+            tmp.append(img_list[j])
+            tmp.append(mid)
+        tmp.append(img1)
+        img_list = tmp
+
+if not os.path.exists('output'):
+    os.mkdir('output')
+for i in range(len(img_list)):
+    if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
+        cv2.imwrite('output/img{}.exr'.format(i), (img_list[i][0]).cpu().numpy().transpose(1, 2, 0)[:h, :w], [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_HALF])
+    else:
+        cv2.imwrite('output/img{}.png'.format(i), (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w])

ECCV2022-RIFE-main/inference_video.py

@@ -0,0 +1,294 @@
+import os
+import cv2
+import torch
+import argparse
+import numpy as np
+from tqdm import tqdm
+from torch.nn import functional as F
+import warnings
+import _thread
+import skvideo.io
+from queue import Queue, Empty
+from model.pytorch_msssim import ssim_matlab
+
+warnings.filterwarnings("ignore")
+
+def transferAudio(sourceVideo, targetVideo):
+    import shutil
+    import moviepy.editor
+    tempAudioFileName = "./temp/audio.mkv"
+
+    # split audio from original video file and store in "temp" directory
+    if True:
+
+        # clear old "temp" directory if it exits
+        if os.path.isdir("temp"):
+            # remove temp directory
+            shutil.rmtree("temp")
+        # create new "temp" directory
+        os.makedirs("temp")
+        # extract audio from video
+        os.system('ffmpeg -y -i "{}" -c:a copy -vn {}'.format(sourceVideo, tempAudioFileName))
+
+    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 "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, 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 "{}" -c:a aac -b:a 160k -vn {}'.format(sourceVideo, tempAudioFileName))
+        os.system('ffmpeg -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo))
+        if (os.path.getsize(targetVideo) == 0): # if aac is not supported by selected format
+            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(targetNoAudio)
+    else:
+        os.remove(targetNoAudio)
+
+    # remove temp directory
+    shutil.rmtree("temp")
+
+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('--model', dest='modelDir', type=str, default='train_log', help='directory with trained model files')
+parser.add_argument('--fp16', dest='fp16', action='store_true', help='fp16 mode for faster and more lightweight inference on cards with Tensor Cores')
+parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video')
+parser.add_argument('--scale', dest='scale', type=float, default=1.0, help='Try scale=0.5 for 4k video')
+parser.add_argument('--skip', dest='skip', action='store_true', help='whether to remove static frames before processing')
+parser.add_argument('--fps', dest='fps', type=int, default=None)
+parser.add_argument('--png', dest='png', action='store_true', help='whether to vid_out png format vid_outs')
+parser.add_argument('--ext', dest='ext', type=str, default='mp4', help='vid_out video extension')
+parser.add_argument('--exp', dest='exp', type=int, default=1)
+args = parser.parse_args()
+assert (not args.video is None or not args.img is None)
+if args.skip:
+    print("skip flag is abandoned, please refer to issue #207.")
+if args.UHD and args.scale==1.0:
+    args.scale = 0.5
+assert args.scale in [0.25, 0.5, 1.0, 2.0, 4.0]
+if not args.img is None:
+    args.png = True
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_grad_enabled(False)
+if torch.cuda.is_available():
+    torch.backends.cudnn.enabled = True
+    torch.backends.cudnn.benchmark = True
+    if(args.fp16):
+        torch.set_default_tensor_type(torch.cuda.HalfTensor)
+
+try:
+    try:
+        try:
+            from model.RIFE_HDv2 import Model
+            model = Model()
+            model.load_model(args.modelDir, -1)
+            print("Loaded v2.x HD model.")
+        except:
+            from train_log.RIFE_HDv3 import Model
+            model = Model()
+            model.load_model(args.modelDir, -1)
+            print("Loaded v3.x HD model.")
+    except:
+        from model.RIFE_HD import Model
+        model = Model()
+        model.load_model(args.modelDir, -1)
+        print("Loaded v1.x HD model")
+except:
+    from model.RIFE import Model
+    model = Model()
+    model.load_model(args.modelDir, -1)
+    print("Loaded ArXiv-RIFE model")
+model.eval()
+model.device()
+
+if not args.video is None:
+    videoCapture = cv2.VideoCapture(args.video)
+    fps = videoCapture.get(cv2.CAP_PROP_FPS)
+    tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
+    videoCapture.release()
+    if args.fps is None:
+        fpsNotAssigned = True
+        args.fps = fps * (2 ** args.exp)
+    else:
+        fpsNotAssigned = False
+    videogen = skvideo.io.vreader(args.video)
+    lastframe = next(videogen)
+    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
+    video_path_wo_ext, ext = os.path.splitext(args.video)
+    print('{}.{}, {} frames in total, {}FPS to {}FPS'.format(video_path_wo_ext, args.ext, tot_frame, fps, args.fps))
+    if args.png == False and fpsNotAssigned == True:
+        print("The audio will be merged after interpolation process")
+    else:
+        print("Will not merge audio because using png or fps flag!")
+else:
+    videogen = []
+    for f in os.listdir(args.img):
+        if 'png' in f:
+            videogen.append(f)
+    tot_frame = len(videogen)
+    videogen.sort(key= lambda x:int(x[:-4]))
+    lastframe = cv2.imread(os.path.join(args.img, videogen[0]), cv2.IMREAD_UNCHANGED)[:, :, ::-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:
+    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
+    while True:
+        item = write_buffer.get()
+        if item is None:
+            break
+        if user_args.png:
+            cv2.imwrite('vid_out/{:0>7d}.png'.format(cnt), item[:, :, ::-1])
+            cnt += 1
+        else:
+            vid_out.write(item[:, :, ::-1])
+
+def build_read_buffer(user_args, read_buffer, videogen):
+    try:
+        for frame in videogen:
+             if not user_args.img is None:
+                  frame = cv2.imread(os.path.join(user_args.img, frame), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy()
+             if user_args.montage:
+                  frame = frame[:, left: left + w]
+             read_buffer.put(frame)
+    except:
+        pass
+    read_buffer.put(None)
+
+def make_inference(I0, I1, n):
+    global model
+    middle = model.inference(I0, I1, args.scale)
+    if n == 1:
+        return [middle]
+    first_half = make_inference(I0, middle, n=n//2)
+    second_half = make_inference(middle, I1, n=n//2)
+    if n%2:
+        return [*first_half, middle, *second_half]
+    else:
+        return [*first_half, *second_half]
+
+def pad_image(img):
+    if(args.fp16):
+        return F.pad(img, padding).half()
+    else:
+        return F.pad(img, padding)
+
+if args.montage:
+    left = w // 4
+    w = w // 2
+tmp = max(32, int(32 / args.scale))
+ph = ((h - 1) // tmp + 1) * tmp
+pw = ((w - 1) // tmp + 1) * tmp
+padding = (0, pw - w, 0, ph - h)
+pbar = tqdm(total=tot_frame)
+if args.montage:
+    lastframe = lastframe[:, left: left + w]
+write_buffer = Queue(maxsize=500)
+read_buffer = Queue(maxsize=500)
+_thread.start_new_thread(build_read_buffer, (args, read_buffer, videogen))
+_thread.start_new_thread(clear_write_buffer, (args, write_buffer))
+
+I1 = torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
+I1 = pad_image(I1)
+temp = None # save lastframe when processing static frame
+
+while True:
+    if temp is not None:
+        frame = temp
+        temp = None
+    else:
+        frame = read_buffer.get()
+    if frame is None:
+        break
+    I0 = I1
+    I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
+    I1 = pad_image(I1)
+    I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False)
+    I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
+    ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
+
+    break_flag = False
+    if ssim > 0.996:        
+        frame = read_buffer.get() # read a new frame
+        if frame is None:
+            break_flag = True
+            frame = lastframe
+        else:
+            temp = frame
+        I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.
+        I1 = pad_image(I1)
+        I1 = model.inference(I0, I1, args.scale)
+        I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
+        ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
+        frame = (I1[0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w]
+    
+    if ssim < 0.2:
+        output = []
+        for i in range((2 ** args.exp) - 1):
+            output.append(I0)
+        '''
+        output = []
+        step = 1 / (2 ** args.exp)
+        alpha = 0
+        for i in range((2 ** args.exp) - 1):
+            alpha += step
+            beta = 1-alpha
+            output.append(torch.from_numpy(np.transpose((cv2.addWeighted(frame[:, :, ::-1], alpha, lastframe[:, :, ::-1], beta, 0)[:, :, ::-1].copy()), (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255.)
+        '''
+    else:
+        output = make_inference(I0, I1, 2**args.exp-1) if args.exp else []
+
+    if args.montage:
+        write_buffer.put(np.concatenate((lastframe, lastframe), 1))
+        for mid in output:
+            mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
+            write_buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1))
+    else:
+        write_buffer.put(lastframe)
+        for mid in output:
+            mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0)))
+            write_buffer.put(mid[:h, :w])
+    pbar.update(1)
+    lastframe = frame
+    if break_flag:
+        break
+
+if args.montage:
+    write_buffer.put(np.concatenate((lastframe, lastframe), 1))
+else:
+    write_buffer.put(lastframe)
+import time
+while(not write_buffer.empty()):
+    time.sleep(0.1)
+pbar.close()
+if not vid_out is None:
+    vid_out.release()
+
+# move audio to new video file if appropriate
+if args.png == False and fpsNotAssigned == True and not args.video is None:
+    try:
+        transferAudio(args.video, vid_out_name)
+    except:
+        print("Audio transfer failed. Interpolated video will have no audio")
+        targetNoAudio = os.path.splitext(vid_out_name)[0] + "_noaudio" + os.path.splitext(vid_out_name)[1]
+        os.rename(targetNoAudio, vid_out_name)

ECCV2022-RIFE-main/model/IFNet.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.warplayer import warp
+from model.refine import *
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),
+        nn.PReLU(out_planes)
+    )
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c=64):
+        super(IFBlock, self).__init__()
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c//2, 3, 2, 1),
+            conv(c//2, c, 3, 2, 1),
+            )
+        self.convblock = nn.Sequential(
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+        )
+        self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
+
+    def forward(self, x, flow, scale):
+        if scale != 1:
+            x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False)
+        if flow != None:
+            flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scale
+            x = torch.cat((x, flow), 1)
+        x = self.conv0(x)
+        x = self.convblock(x) + x
+        tmp = self.lastconv(x)
+        tmp = F.interpolate(tmp, scale_factor = scale * 2, mode="bilinear", align_corners=False)
+        flow = tmp[:, :4] * scale * 2
+        mask = tmp[:, 4:5]
+        return flow, mask
+    
+class IFNet(nn.Module):
+    def __init__(self):
+        super(IFNet, self).__init__()
+        self.block0 = IFBlock(6, c=240)
+        self.block1 = IFBlock(13+4, c=150)
+        self.block2 = IFBlock(13+4, c=90)
+        self.block_tea = IFBlock(16+4, c=90)
+        self.contextnet = Contextnet()
+        self.unet = Unet()
+
+    def forward(self, x, scale=[4,2,1], timestep=0.5):
+        img0 = x[:, :3]
+        img1 = x[:, 3:6]
+        gt = x[:, 6:] # In inference time, gt is None
+        flow_list = []
+        merged = []
+        mask_list = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = None 
+        loss_distill = 0
+        stu = [self.block0, self.block1, self.block2]
+        for i in range(3):
+            if flow != None:
+                flow_d, mask_d = stu[i](torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i])
+                flow = flow + flow_d
+                mask = mask + mask_d
+            else:
+                flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i])
+            mask_list.append(torch.sigmoid(mask))
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged_student = (warped_img0, warped_img1)
+            merged.append(merged_student)
+        if gt.shape[1] == 3:
+            flow_d, mask_d = self.block_tea(torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1)
+            flow_teacher = flow + flow_d
+            warped_img0_teacher = warp(img0, flow_teacher[:, :2])
+            warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])
+            mask_teacher = torch.sigmoid(mask + mask_d)
+            merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)
+        else:
+            flow_teacher = None
+            merged_teacher = None
+        for i in range(3):
+            merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
+            if gt.shape[1] == 3:
+                loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach()
+                loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()
+        c0 = self.contextnet(img0, flow[:, :2])
+        c1 = self.contextnet(img1, flow[:, 2:4])
+        tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
+        res = tmp[:, :3] * 2 - 1
+        merged[2] = torch.clamp(merged[2] + res, 0, 1)
+        return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill

ECCV2022-RIFE-main/model/IFNet_2R.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.warplayer import warp
+from model.refine_2R import *
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),
+        nn.PReLU(out_planes)
+    )
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c=64):
+        super(IFBlock, self).__init__()
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c//2, 3, 1, 1),
+            conv(c//2, c, 3, 2, 1),
+            )
+        self.convblock = nn.Sequential(
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+        )
+        self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
+
+    def forward(self, x, flow, scale):
+        if scale != 1:
+            x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False)
+        if flow != None:
+            flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scale
+            x = torch.cat((x, flow), 1)
+        x = self.conv0(x)
+        x = self.convblock(x) + x
+        tmp = self.lastconv(x)
+        tmp = F.interpolate(tmp, scale_factor = scale, mode="bilinear", align_corners=False)
+        flow = tmp[:, :4] * scale
+        mask = tmp[:, 4:5]
+        return flow, mask
+    
+class IFNet(nn.Module):
+    def __init__(self):
+        super(IFNet, self).__init__()
+        self.block0 = IFBlock(6, c=240)
+        self.block1 = IFBlock(13+4, c=150)
+        self.block2 = IFBlock(13+4, c=90)
+        self.block_tea = IFBlock(16+4, c=90)
+        self.contextnet = Contextnet()
+        self.unet = Unet()
+
+    def forward(self, x, scale=[4,2,1], timestep=0.5):
+        img0 = x[:, :3]
+        img1 = x[:, 3:6]
+        gt = x[:, 6:] # In inference time, gt is None
+        flow_list = []
+        merged = []
+        mask_list = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = None 
+        loss_distill = 0
+        stu = [self.block0, self.block1, self.block2]
+        for i in range(3):
+            if flow != None:
+                flow_d, mask_d = stu[i](torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i])
+                flow = flow + flow_d
+                mask = mask + mask_d
+            else:
+                flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i])
+            mask_list.append(torch.sigmoid(mask))
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged_student = (warped_img0, warped_img1)
+            merged.append(merged_student)
+        if gt.shape[1] == 3:
+            flow_d, mask_d = self.block_tea(torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1)
+            flow_teacher = flow + flow_d
+            warped_img0_teacher = warp(img0, flow_teacher[:, :2])
+            warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])
+            mask_teacher = torch.sigmoid(mask + mask_d)
+            merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)
+        else:
+            flow_teacher = None
+            merged_teacher = None
+        for i in range(3):
+            merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
+            if gt.shape[1] == 3:
+                loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach()
+                loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()
+        c0 = self.contextnet(img0, flow[:, :2])
+        c1 = self.contextnet(img1, flow[:, 2:4])
+        tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
+        res = tmp[:, :3] * 2 - 1
+        merged[2] = torch.clamp(merged[2] + res, 0, 1)
+        return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill

ECCV2022-RIFE-main/model/IFNet_m.py

@@ -0,0 +1,112 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.warplayer import warp
+from model.refine import *
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1),
+        nn.PReLU(out_planes)
+    )
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, c=64):
+        super(IFBlock, self).__init__()
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c//2, 3, 2, 1),
+            conv(c//2, c, 3, 2, 1),
+            )
+        self.convblock = nn.Sequential(
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+            conv(c, c),
+        )
+        self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1)
+
+    def forward(self, x, flow, scale):
+        if scale != 1:
+            x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False)
+        if flow != None:
+            flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scale
+            x = torch.cat((x, flow), 1)
+        x = self.conv0(x)
+        x = self.convblock(x) + x
+        tmp = self.lastconv(x)
+        tmp = F.interpolate(tmp, scale_factor = scale * 2, mode="bilinear", align_corners=False)
+        flow = tmp[:, :4] * scale * 2
+        mask = tmp[:, 4:5]
+        return flow, mask
+    
+class IFNet_m(nn.Module):
+    def __init__(self):
+        super(IFNet_m, self).__init__()
+        self.block0 = IFBlock(6+1, c=240)
+        self.block1 = IFBlock(13+4+1, c=150)
+        self.block2 = IFBlock(13+4+1, c=90)
+        self.block_tea = IFBlock(16+4+1, c=90)
+        self.contextnet = Contextnet()
+        self.unet = Unet()
+
+    def forward(self, x, scale=[4,2,1], timestep=0.5, returnflow=False):
+        timestep = (x[:, :1].clone() * 0 + 1) * timestep
+        img0 = x[:, :3]
+        img1 = x[:, 3:6]
+        gt = x[:, 6:] # In inference time, gt is None
+        flow_list = []
+        merged = []
+        mask_list = []
+        warped_img0 = img0
+        warped_img1 = img1
+        flow = None 
+        loss_distill = 0
+        stu = [self.block0, self.block1, self.block2]
+        for i in range(3):
+            if flow != None:
+                flow_d, mask_d = stu[i](torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask), 1), flow, scale=scale[i])
+                flow = flow + flow_d
+                mask = mask + mask_d
+            else:
+                flow, mask = stu[i](torch.cat((img0, img1, timestep), 1), None, scale=scale[i])
+            mask_list.append(torch.sigmoid(mask))
+            flow_list.append(flow)
+            warped_img0 = warp(img0, flow[:, :2])
+            warped_img1 = warp(img1, flow[:, 2:4])
+            merged_student = (warped_img0, warped_img1)
+            merged.append(merged_student)
+        if gt.shape[1] == 3:
+            flow_d, mask_d = self.block_tea(torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask, gt), 1), flow, scale=1)
+            flow_teacher = flow + flow_d
+            warped_img0_teacher = warp(img0, flow_teacher[:, :2])
+            warped_img1_teacher = warp(img1, flow_teacher[:, 2:4])
+            mask_teacher = torch.sigmoid(mask + mask_d)
+            merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher)
+        else:
+            flow_teacher = None
+            merged_teacher = None
+        for i in range(3):
+            merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
+            if gt.shape[1] == 3:
+                loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach()
+                loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean()
+        if returnflow:
+            return flow
+        else:
+            c0 = self.contextnet(img0, flow[:, :2])
+            c1 = self.contextnet(img1, flow[:, 2:4])
+            tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
+            res = tmp[:, :3] * 2 - 1
+            merged[2] = torch.clamp(merged[2] + res, 0, 1)
+        return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill

ECCV2022-RIFE-main/model/RIFE.py

@@ -0,0 +1,95 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim import AdamW
+import torch.optim as optim
+import itertools
+from model.warplayer import warp
+from torch.nn.parallel import DistributedDataParallel as DDP
+from model.IFNet import *
+from model.IFNet_m import *
+import torch.nn.functional as F
+from model.loss import *
+from model.laplacian import *
+from model.refine import *
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+class Model:
+    def __init__(self, local_rank=-1, arbitrary=False):
+        if arbitrary == True:
+            self.flownet = IFNet_m()
+        else:
+            self.flownet = IFNet()
+        self.device()
+        self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-3) # use large weight decay may avoid NaN loss
+        self.epe = EPE()
+        self.lap = LapLoss()
+        self.sobel = SOBEL()
+        if local_rank != -1:
+            self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank)
+
+    def train(self):
+        self.flownet.train()
+
+    def eval(self):
+        self.flownet.eval()
+
+    def device(self):
+        self.flownet.to(device)
+
+    def load_model(self, path, rank=0):
+        def convert(param):
+            return {
+            k.replace("module.", ""): v
+                for k, v in param.items()
+                if "module." in k
+            }
+            
+        if rank <= 0:
+            self.flownet.load_state_dict(convert(torch.load('{}/flownet.pkl'.format(path))))
+        
+    def save_model(self, path, rank=0):
+        if rank == 0:
+            torch.save(self.flownet.state_dict(),'{}/flownet.pkl'.format(path))
+
+    def inference(self, img0, img1, scale=1, scale_list=[4, 2, 1], TTA=False, timestep=0.5):
+        for i in range(3):
+            scale_list[i] = scale_list[i] * 1.0 / scale
+        imgs = torch.cat((img0, img1), 1)
+        flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(imgs, scale_list, timestep=timestep)
+        if TTA == False:
+            return merged[2]
+        else:
+            flow2, mask2, merged2, flow_teacher2, merged_teacher2, loss_distill2 = self.flownet(imgs.flip(2).flip(3), scale_list, timestep=timestep)
+            return (merged[2] + merged2[2].flip(2).flip(3)) / 2
+    
+    def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
+        for param_group in self.optimG.param_groups:
+            param_group['lr'] = learning_rate
+        img0 = imgs[:, :3]
+        img1 = imgs[:, 3:]
+        if training:
+            self.train()
+        else:
+            self.eval()
+        flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(torch.cat((imgs, gt), 1), scale=[4, 2, 1])
+        loss_l1 = (self.lap(merged[2], gt)).mean()
+        loss_tea = (self.lap(merged_teacher, gt)).mean()
+        if training:
+            self.optimG.zero_grad()
+            loss_G = loss_l1 + loss_tea + loss_distill * 0.01 # when training RIFEm, the weight of loss_distill should be 0.005 or 0.002
+            loss_G.backward()
+            self.optimG.step()
+        else:
+            flow_teacher = flow[2]
+        return merged[2], {
+            'merged_tea': merged_teacher,
+            'mask': mask,
+            'mask_tea': mask,
+            'flow': flow[2][:, :2],
+            'flow_tea': flow_teacher,
+            'loss_l1': loss_l1,
+            'loss_tea': loss_tea,
+            'loss_distill': loss_distill,
+            }

ECCV2022-RIFE-main/model/laplacian.py

@@ -0,0 +1,59 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+import torch
+
+def gauss_kernel(size=5, channels=3):
+    kernel = torch.tensor([[1., 4., 6., 4., 1],
+                           [4., 16., 24., 16., 4.],
+                           [6., 24., 36., 24., 6.],
+                           [4., 16., 24., 16., 4.],
+                           [1., 4., 6., 4., 1.]])
+    kernel /= 256.
+    kernel = kernel.repeat(channels, 1, 1, 1)
+    kernel = kernel.to(device)
+    return kernel
+
+def downsample(x):
+    return x[:, :, ::2, ::2]
+
+def upsample(x):
+    cc = torch.cat([x, torch.zeros(x.shape[0], x.shape[1], x.shape[2], x.shape[3]).to(device)], dim=3)
+    cc = cc.view(x.shape[0], x.shape[1], x.shape[2]*2, x.shape[3])
+    cc = cc.permute(0,1,3,2)
+    cc = torch.cat([cc, torch.zeros(x.shape[0], x.shape[1], x.shape[3], x.shape[2]*2).to(device)], dim=3)
+    cc = cc.view(x.shape[0], x.shape[1], x.shape[3]*2, x.shape[2]*2)
+    x_up = cc.permute(0,1,3,2)
+    return conv_gauss(x_up, 4*gauss_kernel(channels=x.shape[1]))
+
+def conv_gauss(img, kernel):
+    img = torch.nn.functional.pad(img, (2, 2, 2, 2), mode='reflect')
+    out = torch.nn.functional.conv2d(img, kernel, groups=img.shape[1])
+    return out
+
+def laplacian_pyramid(img, kernel, max_levels=3):
+    current = img
+    pyr = []
+    for level in range(max_levels):
+        filtered = conv_gauss(current, kernel)
+        down = downsample(filtered)
+        up = upsample(down)
+        diff = current-up
+        pyr.append(diff)
+        current = down
+    return pyr
+
+class LapLoss(torch.nn.Module):
+    def __init__(self, max_levels=5, channels=3):
+        super(LapLoss, self).__init__()
+        self.max_levels = max_levels
+        self.gauss_kernel = gauss_kernel(channels=channels)
+        
+    def forward(self, input, target):
+        pyr_input  = laplacian_pyramid(img=input, kernel=self.gauss_kernel, max_levels=self.max_levels)
+        pyr_target = laplacian_pyramid(img=target, kernel=self.gauss_kernel, max_levels=self.max_levels)
+        return sum(torch.nn.functional.l1_loss(a, b) for a, b in zip(pyr_input, pyr_target))

ECCV2022-RIFE-main/model/loss.py

@@ -0,0 +1,128 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.models as models
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class EPE(nn.Module):
+    def __init__(self):
+        super(EPE, self).__init__()
+
+    def forward(self, flow, gt, loss_mask):
+        loss_map = (flow - gt.detach()) ** 2
+        loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5
+        return (loss_map * loss_mask)
+
+
+class Ternary(nn.Module):
+    def __init__(self):
+        super(Ternary, self).__init__()
+        patch_size = 7
+        out_channels = patch_size * patch_size
+        self.w = np.eye(out_channels).reshape(
+            (patch_size, patch_size, 1, out_channels))
+        self.w = np.transpose(self.w, (3, 2, 0, 1))
+        self.w = torch.tensor(self.w).float().to(device)
+
+    def transform(self, img):
+        patches = F.conv2d(img, self.w, padding=3, bias=None)
+        transf = patches - img
+        transf_norm = transf / torch.sqrt(0.81 + transf**2)
+        return transf_norm
+
+    def rgb2gray(self, rgb):
+        r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :]
+        gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
+        return gray
+
+    def hamming(self, t1, t2):
+        dist = (t1 - t2) ** 2
+        dist_norm = torch.mean(dist / (0.1 + dist), 1, True)
+        return dist_norm
+
+    def valid_mask(self, t, padding):
+        n, _, h, w = t.size()
+        inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t)
+        mask = F.pad(inner, [padding] * 4)
+        return mask
+
+    def forward(self, img0, img1):
+        img0 = self.transform(self.rgb2gray(img0))
+        img1 = self.transform(self.rgb2gray(img1))
+        return self.hamming(img0, img1) * self.valid_mask(img0, 1)
+
+
+class SOBEL(nn.Module):
+    def __init__(self):
+        super(SOBEL, self).__init__()
+        self.kernelX = torch.tensor([
+            [1, 0, -1],
+            [2, 0, -2],
+            [1, 0, -1],
+        ]).float()
+        self.kernelY = self.kernelX.clone().T
+        self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device)
+        self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device)
+
+    def forward(self, pred, gt):
+        N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3]
+        img_stack = torch.cat(
+            [pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0)
+        sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1)
+        sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1)
+        pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:]
+        pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:]
+
+        L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y)
+        loss = (L1X+L1Y)
+        return loss
+
+class MeanShift(nn.Conv2d):
+    def __init__(self, data_mean, data_std, data_range=1, norm=True):
+        c = len(data_mean)
+        super(MeanShift, self).__init__(c, c, kernel_size=1)
+        std = torch.Tensor(data_std)
+        self.weight.data = torch.eye(c).view(c, c, 1, 1)
+        if norm:
+            self.weight.data.div_(std.view(c, 1, 1, 1))
+            self.bias.data = -1 * data_range * torch.Tensor(data_mean)
+            self.bias.data.div_(std)
+        else:
+            self.weight.data.mul_(std.view(c, 1, 1, 1))
+            self.bias.data = data_range * torch.Tensor(data_mean)
+        self.requires_grad = False
+            
+class VGGPerceptualLoss(torch.nn.Module):
+    def __init__(self, rank=0):
+        super(VGGPerceptualLoss, self).__init__()
+        blocks = []
+        pretrained = True
+        self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features
+        self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda()
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, X, Y, indices=None):
+        X = self.normalize(X)
+        Y = self.normalize(Y)
+        indices = [2, 7, 12, 21, 30]
+        weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5]
+        k = 0
+        loss = 0
+        for i in range(indices[-1]):
+            X = self.vgg_pretrained_features[i](X)
+            Y = self.vgg_pretrained_features[i](Y)
+            if (i+1) in indices:
+                loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1
+                k += 1
+        return loss
+
+if __name__ == '__main__':
+    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
+    img1 = torch.tensor(np.random.normal(
+        0, 1, (3, 3, 256, 256))).float().to(device)
+    ternary_loss = Ternary()
+    print(ternary_loss(img0, img1).shape)

ECCV2022-RIFE-main/model/oldmodel/IFNet_HD.py

@@ -0,0 +1,122 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from model.warplayer import warp
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def conv_wo_act(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=False),
+        nn.BatchNorm2d(out_planes),
+    )
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=False),
+        nn.BatchNorm2d(out_planes),
+        nn.PReLU(out_planes)
+    )
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_planes, out_planes, stride=1):
+        super(ResBlock, self).__init__()
+        if in_planes == out_planes and stride == 1:
+            self.conv0 = nn.Identity()
+        else:
+            self.conv0 = nn.Conv2d(in_planes, out_planes,
+                                   3, stride, 1, bias=False)
+        self.conv1 = conv(in_planes, out_planes, 5, stride, 2)
+        self.conv2 = conv_wo_act(out_planes, out_planes, 3, 1, 1)
+        self.relu1 = nn.PReLU(1)
+        self.relu2 = nn.PReLU(out_planes)
+        self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False)
+        self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        y = self.conv0(x)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        w = x.mean(3, True).mean(2, True)
+        w = self.relu1(self.fc1(w))
+        w = torch.sigmoid(self.fc2(w))
+        x = self.relu2(x * w + y)
+        return x
+
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, scale=1, c=64):
+        super(IFBlock, self).__init__()
+        self.scale = scale
+        self.conv0 = conv(in_planes, c, 5, 2, 2)
+        self.res0 = ResBlock(c, c)
+        self.res1 = ResBlock(c, c)
+        self.res2 = ResBlock(c, c)
+        self.res3 = ResBlock(c, c)
+        self.res4 = ResBlock(c, c)
+        self.res5 = ResBlock(c, c)
+        self.conv1 = nn.Conv2d(c, 8, 3, 1, 1)
+        self.up = nn.PixelShuffle(2)
+
+    def forward(self, x):
+        if self.scale != 1:
+            x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear",
+                              align_corners=False)
+        x = self.conv0(x)
+        x = self.res0(x)
+        x = self.res1(x)
+        x = self.res2(x)
+        x = self.res3(x)
+        x = self.res4(x)
+        x = self.res5(x)
+        x = self.conv1(x)
+        flow = self.up(x)
+        if self.scale != 1:
+            flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear",
+                                 align_corners=False)
+        return flow
+
+
+class IFNet(nn.Module):
+    def __init__(self):
+        super(IFNet, self).__init__()
+        self.block0 = IFBlock(6, scale=8, c=192)
+        self.block1 = IFBlock(8, scale=4, c=128)
+        self.block2 = IFBlock(8, scale=2, c=96)
+        self.block3 = IFBlock(8, scale=1, c=48)
+
+    def forward(self, x, scale=1.0):
+        x = F.interpolate(x, scale_factor=0.5 * scale, mode="bilinear",
+                          align_corners=False)
+        flow0 = self.block0(x)
+        F1 = flow0
+        warped_img0 = warp(x[:, :3], F1)
+        warped_img1 = warp(x[:, 3:], -F1)
+        flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1), 1))
+        F2 = (flow0 + flow1)
+        warped_img0 = warp(x[:, :3], F2)
+        warped_img1 = warp(x[:, 3:], -F2)
+        flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2), 1))
+        F3 = (flow0 + flow1 + flow2)
+        warped_img0 = warp(x[:, :3], F3)
+        warped_img1 = warp(x[:, 3:], -F3)
+        flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3), 1))
+        F4 = (flow0 + flow1 + flow2 + flow3)
+        F4 = F.interpolate(F4, scale_factor=1 / scale, mode="bilinear",
+                           align_corners=False) / scale
+        return F4, [F1, F2, F3, F4]
+
+if __name__ == '__main__':
+    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
+    img1 = torch.tensor(np.random.normal(
+        0, 1, (3, 3, 256, 256))).float().to(device)
+    imgs = torch.cat((img0, img1), 1)
+    flownet = IFNet()
+    flow, _ = flownet(imgs)
+    print(flow.shape)

ECCV2022-RIFE-main/model/oldmodel/IFNet_HDv2.py

@@ -0,0 +1,95 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from model.warplayer import warp
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def conv_wo_act(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+    )
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+class IFBlock(nn.Module):
+    def __init__(self, in_planes, scale=1, c=64):
+        super(IFBlock, self).__init__()
+        self.scale = scale
+        self.conv0 = nn.Sequential(
+            conv(in_planes, c, 3, 2, 1),
+            conv(c, 2*c, 3, 2, 1),
+            )
+        self.convblock = nn.Sequential(
+            conv(2*c, 2*c),
+            conv(2*c, 2*c),
+            conv(2*c, 2*c),
+            conv(2*c, 2*c),
+            conv(2*c, 2*c),
+            conv(2*c, 2*c),
+        )        
+        self.conv1 = nn.ConvTranspose2d(2*c, 4, 4, 2, 1)
+                    
+    def forward(self, x):
+        if self.scale != 1:
+            x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear",
+                              align_corners=False)
+        x = self.conv0(x)
+        x = self.convblock(x)
+        x = self.conv1(x)
+        flow = x
+        if self.scale != 1:
+            flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear",
+                                 align_corners=False)
+        return flow
+
+
+class IFNet(nn.Module):
+    def __init__(self):
+        super(IFNet, self).__init__()
+        self.block0 = IFBlock(6, scale=8, c=192)
+        self.block1 = IFBlock(10, scale=4, c=128)
+        self.block2 = IFBlock(10, scale=2, c=96)
+        self.block3 = IFBlock(10, scale=1, c=48)
+
+    def forward(self, x, scale=1.0):
+        if scale != 1.0:
+            x = F.interpolate(x, scale_factor=scale, mode="bilinear", align_corners=False)
+        flow0 = self.block0(x)
+        F1 = flow0
+        F1_large = F.interpolate(F1, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0
+        warped_img0 = warp(x[:, :3], F1_large[:, :2])
+        warped_img1 = warp(x[:, 3:], F1_large[:, 2:4])
+        flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1_large), 1))
+        F2 = (flow0 + flow1)
+        F2_large = F.interpolate(F2, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0
+        warped_img0 = warp(x[:, :3], F2_large[:, :2])
+        warped_img1 = warp(x[:, 3:], F2_large[:, 2:4])
+        flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2_large), 1))
+        F3 = (flow0 + flow1 + flow2)
+        F3_large = F.interpolate(F3, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0
+        warped_img0 = warp(x[:, :3], F3_large[:, :2])
+        warped_img1 = warp(x[:, 3:], F3_large[:, 2:4])
+        flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3_large), 1))
+        F4 = (flow0 + flow1 + flow2 + flow3)
+        if scale != 1.0:
+            F4 = F.interpolate(F4, scale_factor=1 / scale, mode="bilinear", align_corners=False) / scale
+        return F4, [F1, F2, F3, F4]
+
+if __name__ == '__main__':
+    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
+    img1 = torch.tensor(np.random.normal(
+        0, 1, (3, 3, 256, 256))).float().to(device)
+    imgs = torch.cat((img0, img1), 1)
+    flownet = IFNet()
+    flow, _ = flownet(imgs)
+    print(flow.shape)

ECCV2022-RIFE-main/model/oldmodel/RIFE_HD.py

@@ -0,0 +1,260 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim import AdamW
+import torch.optim as optim
+import itertools
+from model.warplayer import warp
+from torch.nn.parallel import DistributedDataParallel as DDP
+from model.oldmodel.IFNet_HD import *
+import torch.nn.functional as F
+from model.loss import *
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes,
+                                 kernel_size=4, stride=2, padding=1, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+    )
+
+class ResBlock(nn.Module):
+    def __init__(self, in_planes, out_planes, stride=2):
+        super(ResBlock, self).__init__()
+        if in_planes == out_planes and stride == 1:
+            self.conv0 = nn.Identity()
+        else:
+            self.conv0 = nn.Conv2d(in_planes, out_planes,
+                                   3, stride, 1, bias=False)
+        self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
+        self.conv2 = conv_woact(out_planes, out_planes, 3, 1, 1)
+        self.relu1 = nn.PReLU(1)
+        self.relu2 = nn.PReLU(out_planes)
+        self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False)
+        self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        y = self.conv0(x)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        w = x.mean(3, True).mean(2, True)
+        w = self.relu1(self.fc1(w))
+        w = torch.sigmoid(self.fc2(w))
+        x = self.relu2(x * w + y)
+        return x
+
+c = 32
+
+class ContextNet(nn.Module):
+    def __init__(self):
+        super(ContextNet, self).__init__()
+        self.conv0 = conv(3, c, 3, 2, 1)
+        self.conv1 = ResBlock(c, c)
+        self.conv2 = ResBlock(c, 2*c)
+        self.conv3 = ResBlock(2*c, 4*c)
+        self.conv4 = ResBlock(4*c, 8*c)
+
+    def forward(self, x, flow):
+        x = self.conv0(x)
+        x = self.conv1(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5
+        f1 = warp(x, flow)
+        x = self.conv2(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
+                             align_corners=False) * 0.5
+        f2 = warp(x, flow)
+        x = self.conv3(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
+                             align_corners=False) * 0.5
+        f3 = warp(x, flow)
+        x = self.conv4(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
+                             align_corners=False) * 0.5
+        f4 = warp(x, flow)
+        return [f1, f2, f3, f4]
+
+
+class FusionNet(nn.Module):
+    def __init__(self):
+        super(FusionNet, self).__init__()
+        self.conv0 = conv(8, c, 3, 2, 1)
+        self.down0 = ResBlock(c, 2*c)
+        self.down1 = ResBlock(4*c, 4*c)
+        self.down2 = ResBlock(8*c, 8*c)
+        self.down3 = ResBlock(16*c, 16*c)
+        self.up0 = deconv(32*c, 8*c)
+        self.up1 = deconv(16*c, 4*c)
+        self.up2 = deconv(8*c, 2*c)
+        self.up3 = deconv(4*c, c)
+        self.conv = nn.Conv2d(c, 16, 3, 1, 1)
+        self.up4 = nn.PixelShuffle(2)
+
+    def forward(self, img0, img1, flow, c0, c1, flow_gt):
+        warped_img0 = warp(img0, flow)
+        warped_img1 = warp(img1, -flow)
+        if flow_gt == None:
+            warped_img0_gt, warped_img1_gt = None, None
+        else:
+            warped_img0_gt = warp(img0, flow_gt[:, :2])
+            warped_img1_gt = warp(img1, flow_gt[:, 2:4])
+        x = self.conv0(torch.cat((warped_img0, warped_img1, flow), 1))
+        s0 = self.down0(x)
+        s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
+        s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
+        s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
+        x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
+        x = self.up1(torch.cat((x, s2), 1))
+        x = self.up2(torch.cat((x, s1), 1))
+        x = self.up3(torch.cat((x, s0), 1))
+        x = self.up4(self.conv(x))
+        return x, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
+
+
+class Model:
+    def __init__(self, local_rank=-1):
+        self.flownet = IFNet()
+        self.contextnet = ContextNet()
+        self.fusionnet = FusionNet()
+        self.device()
+        self.optimG = AdamW(itertools.chain(
+            self.flownet.parameters(),
+            self.contextnet.parameters(),
+            self.fusionnet.parameters()), lr=1e-6, weight_decay=1e-4)
+        self.schedulerG = optim.lr_scheduler.CyclicLR(
+            self.optimG, base_lr=1e-6, max_lr=1e-3, step_size_up=8000, cycle_momentum=False)
+        self.epe = EPE()
+        self.ter = Ternary()
+        self.sobel = SOBEL()
+        if local_rank != -1:
+            self.flownet = DDP(self.flownet, device_ids=[
+                               local_rank], output_device=local_rank)
+            self.contextnet = DDP(self.contextnet, device_ids=[
+                                  local_rank], output_device=local_rank)
+            self.fusionnet = DDP(self.fusionnet, device_ids=[
+                                 local_rank], output_device=local_rank)
+
+    def train(self):
+        self.flownet.train()
+        self.contextnet.train()
+        self.fusionnet.train()
+
+    def eval(self):
+        self.flownet.eval()
+        self.contextnet.eval()
+        self.fusionnet.eval()
+
+    def device(self):
+        self.flownet.to(device)
+        self.contextnet.to(device)
+        self.fusionnet.to(device)
+
+    def load_model(self, path, rank):
+        def convert(param):
+            if rank == -1:
+                return {
+                    k.replace("module.", ""): v
+                    for k, v in param.items()
+                    if "module." in k
+                }
+            else:
+                return param
+        if rank <= 0:
+            self.flownet.load_state_dict(
+                convert(torch.load('{}/flownet.pkl'.format(path), map_location=device)))
+            self.contextnet.load_state_dict(
+                convert(torch.load('{}/contextnet.pkl'.format(path), map_location=device)))
+            self.fusionnet.load_state_dict(
+                convert(torch.load('{}/unet.pkl'.format(path), map_location=device)))
+
+    def save_model(self, path, rank):
+        if rank == 0:
+            torch.save(self.flownet.state_dict(), '{}/flownet.pkl'.format(path))
+            torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path))
+            torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path))
+
+    def predict(self, imgs, flow, training=True, flow_gt=None):
+        img0 = imgs[:, :3]
+        img1 = imgs[:, 3:]
+        c0 = self.contextnet(img0, flow)
+        c1 = self.contextnet(img1, -flow)
+        flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear",
+                             align_corners=False) * 2.0
+        refine_output, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.fusionnet(
+            img0, img1, flow, c0, c1, flow_gt)
+        res = torch.sigmoid(refine_output[:, :3]) * 2 - 1
+        mask = torch.sigmoid(refine_output[:, 3:4])
+        merged_img = warped_img0 * mask + warped_img1 * (1 - mask)
+        pred = merged_img + res
+        pred = torch.clamp(pred, 0, 1)
+        if training:
+            return pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
+        else:
+            return pred
+
+    def inference(self, img0, img1, scale=1.0):
+        imgs = torch.cat((img0, img1), 1)
+        flow, _ = self.flownet(imgs, scale)
+        return self.predict(imgs, flow, training=False)
+
+    def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
+        for param_group in self.optimG.param_groups:
+            param_group['lr'] = learning_rate
+        if training:
+            self.train()
+        else:
+            self.eval()
+        flow, flow_list = self.flownet(imgs)
+        pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.predict(
+            imgs, flow, flow_gt=flow_gt)
+        loss_ter = self.ter(pred, gt).mean()
+        if training:
+            with torch.no_grad():
+                loss_flow = torch.abs(warped_img0_gt - gt).mean()
+                loss_mask = torch.abs(
+                    merged_img - gt).sum(1, True).float().detach()
+                loss_mask = F.interpolate(loss_mask, scale_factor=0.5, mode="bilinear",
+                                          align_corners=False).detach()
+                flow_gt = (F.interpolate(flow_gt, scale_factor=0.5, mode="bilinear",
+                                         align_corners=False) * 0.5).detach()
+            loss_cons = 0
+            for i in range(3):
+                loss_cons += self.epe(flow_list[i], flow_gt[:, :2], 1)
+                loss_cons += self.epe(-flow_list[i], flow_gt[:, 2:4], 1)
+            loss_cons = loss_cons.mean() * 0.01
+        else:
+            loss_cons = torch.tensor([0])
+            loss_flow = torch.abs(warped_img0 - gt).mean()
+            loss_mask = 1
+        loss_l1 = (((pred - gt) ** 2 + 1e-6) ** 0.5).mean()
+        if training:
+            self.optimG.zero_grad()
+            loss_G = loss_l1 + loss_cons + loss_ter
+            loss_G.backward()
+            self.optimG.step()
+        return pred, merged_img, flow, loss_l1, loss_flow, loss_cons, loss_ter, loss_mask
+
+
+if __name__ == '__main__':
+    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
+    img1 = torch.tensor(np.random.normal(
+        0, 1, (3, 3, 256, 256))).float().to(device)
+    imgs = torch.cat((img0, img1), 1)
+    model = Model()
+    model.eval()
+    print(model.inference(imgs).shape)

ECCV2022-RIFE-main/model/oldmodel/RIFE_HDv2.py

@@ -0,0 +1,245 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torch.optim import AdamW
+import torch.optim as optim
+import itertools
+from model.warplayer import warp
+from torch.nn.parallel import DistributedDataParallel as DDP
+from model.oldmodel.IFNet_HDv2 import *
+import torch.nn.functional as F
+from model.loss import *
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes,
+                                 kernel_size=4, stride=2, padding=1, bias=True),
+        nn.PReLU(out_planes)
+    )
+
+def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+    )
+
+class Conv2(nn.Module):
+    def __init__(self, in_planes, out_planes, stride=2):
+        super(Conv2, self).__init__()
+        self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
+        self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+c = 32
+
+class ContextNet(nn.Module):
+    def __init__(self):
+        super(ContextNet, self).__init__()
+        self.conv0 = Conv2(3, c)
+        self.conv1 = Conv2(c, c)
+        self.conv2 = Conv2(c, 2*c)
+        self.conv3 = Conv2(2*c, 4*c)
+        self.conv4 = Conv2(4*c, 8*c)
+
+    def forward(self, x, flow):
+        x = self.conv0(x)
+        x = self.conv1(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5
+        f1 = warp(x, flow)
+        x = self.conv2(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
+                             align_corners=False) * 0.5
+        f2 = warp(x, flow)
+        x = self.conv3(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
+                             align_corners=False) * 0.5
+        f3 = warp(x, flow)
+        x = self.conv4(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear",
+                             align_corners=False) * 0.5
+        f4 = warp(x, flow)
+        return [f1, f2, f3, f4]
+
+
+class FusionNet(nn.Module):
+    def __init__(self):
+        super(FusionNet, self).__init__()
+        self.conv0 = Conv2(10, c)
+        self.down0 = Conv2(c, 2*c)
+        self.down1 = Conv2(4*c, 4*c)
+        self.down2 = Conv2(8*c, 8*c)
+        self.down3 = Conv2(16*c, 16*c)
+        self.up0 = deconv(32*c, 8*c)
+        self.up1 = deconv(16*c, 4*c)
+        self.up2 = deconv(8*c, 2*c)
+        self.up3 = deconv(4*c, c)
+        self.conv = nn.ConvTranspose2d(c, 4, 4, 2, 1)
+
+    def forward(self, img0, img1, flow, c0, c1, flow_gt):
+        warped_img0 = warp(img0, flow[:, :2])
+        warped_img1 = warp(img1, flow[:, 2:4])
+        if flow_gt == None:
+            warped_img0_gt, warped_img1_gt = None, None
+        else:
+            warped_img0_gt = warp(img0, flow_gt[:, :2])
+            warped_img1_gt = warp(img1, flow_gt[:, 2:4])
+        x = self.conv0(torch.cat((warped_img0, warped_img1, flow), 1))
+        s0 = self.down0(x)
+        s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
+        s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
+        s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
+        x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
+        x = self.up1(torch.cat((x, s2), 1))
+        x = self.up2(torch.cat((x, s1), 1))
+        x = self.up3(torch.cat((x, s0), 1))
+        x = self.conv(x)
+        return x, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
+
+
+class Model:
+    def __init__(self, local_rank=-1):
+        self.flownet = IFNet()
+        self.contextnet = ContextNet()
+        self.fusionnet = FusionNet()
+        self.device()
+        self.optimG = AdamW(itertools.chain(
+            self.flownet.parameters(),
+            self.contextnet.parameters(),
+            self.fusionnet.parameters()), lr=1e-6, weight_decay=1e-4)
+        self.schedulerG = optim.lr_scheduler.CyclicLR(
+            self.optimG, base_lr=1e-6, max_lr=1e-3, step_size_up=8000, cycle_momentum=False)
+        self.epe = EPE()
+        self.ter = Ternary()
+        self.sobel = SOBEL()
+        if local_rank != -1:
+            self.flownet = DDP(self.flownet, device_ids=[
+                               local_rank], output_device=local_rank)
+            self.contextnet = DDP(self.contextnet, device_ids=[
+                                  local_rank], output_device=local_rank)
+            self.fusionnet = DDP(self.fusionnet, device_ids=[
+                                 local_rank], output_device=local_rank)
+
+    def train(self):
+        self.flownet.train()
+        self.contextnet.train()
+        self.fusionnet.train()
+
+    def eval(self):
+        self.flownet.eval()
+        self.contextnet.eval()
+        self.fusionnet.eval()
+
+    def device(self):
+        self.flownet.to(device)
+        self.contextnet.to(device)
+        self.fusionnet.to(device)
+
+    def load_model(self, path, rank):
+        def convert(param):
+            if rank == -1:
+                return {
+                    k.replace("module.", ""): v
+                    for k, v in param.items()
+                    if "module." in k
+                }
+            else:
+                return param
+        if rank <= 0:
+            self.flownet.load_state_dict(
+                convert(torch.load('{}/flownet.pkl'.format(path), map_location=device)))
+            self.contextnet.load_state_dict(
+                convert(torch.load('{}/contextnet.pkl'.format(path), map_location=device)))
+            self.fusionnet.load_state_dict(
+                convert(torch.load('{}/unet.pkl'.format(path), map_location=device)))
+
+    def save_model(self, path, rank):
+        if rank == 0:
+            torch.save(self.flownet.state_dict(), '{}/flownet.pkl'.format(path))
+            torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path))
+            torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path))
+
+    def predict(self, imgs, flow, training=True, flow_gt=None):
+        img0 = imgs[:, :3]
+        img1 = imgs[:, 3:]
+        c0 = self.contextnet(img0, flow[:, :2])
+        c1 = self.contextnet(img1, flow[:, 2:4])
+        flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear",
+                             align_corners=False) * 2.0
+        refine_output, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.fusionnet(
+            img0, img1, flow, c0, c1, flow_gt)
+        res = torch.sigmoid(refine_output[:, :3]) * 2 - 1
+        mask = torch.sigmoid(refine_output[:, 3:4])
+        merged_img = warped_img0 * mask + warped_img1 * (1 - mask)
+        pred = merged_img + res
+        pred = torch.clamp(pred, 0, 1)
+        if training:
+            return pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt
+        else:
+            return pred
+
+    def inference(self, img0, img1, scale=1.0):
+        imgs = torch.cat((img0, img1), 1)
+        flow, _ = self.flownet(imgs, scale)
+        return self.predict(imgs, flow, training=False)
+
+    def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None):
+        for param_group in self.optimG.param_groups:
+            param_group['lr'] = learning_rate
+        if training:
+            self.train()
+        else:
+            self.eval()
+        flow, flow_list = self.flownet(imgs)
+        pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.predict(
+            imgs, flow, flow_gt=flow_gt)
+        loss_ter = self.ter(pred, gt).mean()
+        if training:
+            with torch.no_grad():
+                loss_flow = torch.abs(warped_img0_gt - gt).mean()
+                loss_mask = torch.abs(
+                    merged_img - gt).sum(1, True).float().detach()
+                loss_mask = F.interpolate(loss_mask, scale_factor=0.5, mode="bilinear",
+                                          align_corners=False).detach()
+                flow_gt = (F.interpolate(flow_gt, scale_factor=0.5, mode="bilinear",
+                                         align_corners=False) * 0.5).detach()
+            loss_cons = 0
+            for i in range(4):
+                loss_cons += self.epe(flow_list[i][:, :2], flow_gt[:, :2], 1)
+                loss_cons += self.epe(flow_list[i][:, 2:4], flow_gt[:, 2:4], 1)
+            loss_cons = loss_cons.mean() * 0.01
+        else:
+            loss_cons = torch.tensor([0])
+            loss_flow = torch.abs(warped_img0 - gt).mean()
+            loss_mask = 1
+        loss_l1 = (((pred - gt) ** 2 + 1e-6) ** 0.5).mean()
+        if training:
+            self.optimG.zero_grad()
+            loss_G = loss_l1 + loss_cons + loss_ter
+            loss_G.backward()
+            self.optimG.step()
+        return pred, merged_img, flow, loss_l1, loss_flow, loss_cons, loss_ter, loss_mask
+
+
+if __name__ == '__main__':
+    img0 = torch.zeros(3, 3, 256, 256).float().to(device)
+    img1 = torch.tensor(np.random.normal(
+        0, 1, (3, 3, 256, 256))).float().to(device)
+    imgs = torch.cat((img0, img1), 1)
+    model = Model()
+    model.eval()
+    print(model.inference(imgs).shape)

ECCV2022-RIFE-main/model/pytorch_msssim/__init__.py

@@ -0,0 +1,200 @@
+import torch
+import torch.nn.functional as F
+from math import exp
+import numpy as np
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+    return gauss/gauss.sum()
+
+
+def create_window(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device)
+    window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
+    return window
+
+def create_window_3d(window_size, channel=1):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t())
+    _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t())
+    window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device)
+    return window
+
+
+def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
+    # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
+    if val_range is None:
+        if torch.max(img1) > 128:
+            max_val = 255
+        else:
+            max_val = 1
+
+        if torch.min(img1) < -0.5:
+            min_val = -1
+        else:
+            min_val = 0
+        L = max_val - min_val
+    else:
+        L = val_range
+
+    padd = 0
+    (_, channel, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window(real_size, channel=channel).to(img1.device)
+    
+    # mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
+    # mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
+    mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel)
+    mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_mu2
+
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+
+    if full:
+        return ret, cs
+    return ret
+
+
+def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):
+    # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).
+    if val_range is None:
+        if torch.max(img1) > 128:
+            max_val = 255
+        else:
+            max_val = 1
+
+        if torch.min(img1) < -0.5:
+            min_val = -1
+        else:
+            min_val = 0
+        L = max_val - min_val
+    else:
+        L = val_range
+
+    padd = 0
+    (_, _, height, width) = img1.size()
+    if window is None:
+        real_size = min(window_size, height, width)
+        window = create_window_3d(real_size, channel=1).to(img1.device)
+        # Channel is set to 1 since we consider color images as volumetric images
+
+    img1 = img1.unsqueeze(1)
+    img2 = img2.unsqueeze(1)
+
+    mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1)
+    mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_sq
+    sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu2_sq
+    sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_mu2
+
+    C1 = (0.01 * L) ** 2
+    C2 = (0.03 * L) ** 2
+
+    v1 = 2.0 * sigma12 + C2
+    v2 = sigma1_sq + sigma2_sq + C2
+    cs = torch.mean(v1 / v2)  # contrast sensitivity
+
+    ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)
+
+    if size_average:
+        ret = ssim_map.mean()
+    else:
+        ret = ssim_map.mean(1).mean(1).mean(1)
+
+    if full:
+        return ret, cs
+    return ret
+
+
+def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False):
+    device = img1.device
+    weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device)
+    levels = weights.size()[0]
+    mssim = []
+    mcs = []
+    for _ in range(levels):
+        sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)
+        mssim.append(sim)
+        mcs.append(cs)
+
+        img1 = F.avg_pool2d(img1, (2, 2))
+        img2 = F.avg_pool2d(img2, (2, 2))
+
+    mssim = torch.stack(mssim)
+    mcs = torch.stack(mcs)
+
+    # Normalize (to avoid NaNs during training unstable models, not compliant with original definition)
+    if normalize:
+        mssim = (mssim + 1) / 2
+        mcs = (mcs + 1) / 2
+
+    pow1 = mcs ** weights
+    pow2 = mssim ** weights
+    # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/
+    output = torch.prod(pow1[:-1] * pow2[-1])
+    return output
+
+
+# Classes to re-use window
+class SSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True, val_range=None):
+        super(SSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.val_range = val_range
+
+        # Assume 3 channel for SSIM
+        self.channel = 3
+        self.window = create_window(window_size, channel=self.channel)
+
+    def forward(self, img1, img2):
+        (_, channel, _, _) = img1.size()
+
+        if channel == self.channel and self.window.dtype == img1.dtype:
+            window = self.window
+        else:
+            window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)
+            self.window = window
+            self.channel = channel
+
+        _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)
+        dssim = (1 - _ssim) / 2
+        return dssim
+
+class MSSSIM(torch.nn.Module):
+    def __init__(self, window_size=11, size_average=True, channel=3):
+        super(MSSSIM, self).__init__()
+        self.window_size = window_size
+        self.size_average = size_average
+        self.channel = channel
+
+    def forward(self, img1, img2):
+        return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)

ECCV2022-RIFE-main/model/refine.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.optim as optim
+import itertools
+from model.warplayer import warp
+import torch.nn.functional as F
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+        )
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True),
+        nn.PReLU(out_planes)
+        )
+            
+class Conv2(nn.Module):
+    def __init__(self, in_planes, out_planes, stride=2):
+        super(Conv2, self).__init__()
+        self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
+        self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+    
+c = 16
+class Contextnet(nn.Module):
+    def __init__(self):
+        super(Contextnet, self).__init__()
+        self.conv1 = Conv2(3, c)
+        self.conv2 = Conv2(c, 2*c)
+        self.conv3 = Conv2(2*c, 4*c)
+        self.conv4 = Conv2(4*c, 8*c)
+    
+    def forward(self, x, flow):
+        x = self.conv1(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f1 = warp(x, flow)        
+        x = self.conv2(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f2 = warp(x, flow)
+        x = self.conv3(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f3 = warp(x, flow)
+        x = self.conv4(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f4 = warp(x, flow)
+        return [f1, f2, f3, f4]
+    
+class Unet(nn.Module):
+    def __init__(self):
+        super(Unet, self).__init__()
+        self.down0 = Conv2(17, 2*c)
+        self.down1 = Conv2(4*c, 4*c)
+        self.down2 = Conv2(8*c, 8*c)
+        self.down3 = Conv2(16*c, 16*c)
+        self.up0 = deconv(32*c, 8*c)
+        self.up1 = deconv(16*c, 4*c)
+        self.up2 = deconv(8*c, 2*c)
+        self.up3 = deconv(4*c, c)
+        self.conv = nn.Conv2d(c, 3, 3, 1, 1)
+
+    def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
+        s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1))
+        s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
+        s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
+        s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
+        x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
+        x = self.up1(torch.cat((x, s2), 1)) 
+        x = self.up2(torch.cat((x, s1), 1)) 
+        x = self.up3(torch.cat((x, s0), 1)) 
+        x = self.conv(x)
+        return torch.sigmoid(x)

ECCV2022-RIFE-main/model/refine_2R.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.optim as optim
+import itertools
+from model.warplayer import warp
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.nn.functional as F
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1):
+    return nn.Sequential(
+        nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride,
+                  padding=padding, dilation=dilation, bias=True),
+        nn.PReLU(out_planes)
+        )
+
+def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1):
+    return nn.Sequential(
+        torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True),
+        nn.PReLU(out_planes)
+        )
+            
+class Conv2(nn.Module):
+    def __init__(self, in_planes, out_planes, stride=2):
+        super(Conv2, self).__init__()
+        self.conv1 = conv(in_planes, out_planes, 3, stride, 1)
+        self.conv2 = conv(out_planes, out_planes, 3, 1, 1)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+    
+c = 16
+class Contextnet(nn.Module):
+    def __init__(self):
+        super(Contextnet, self).__init__()
+        self.conv1 = Conv2(3, c, 1)
+        self.conv2 = Conv2(c, 2*c)
+        self.conv3 = Conv2(2*c, 4*c)
+        self.conv4 = Conv2(4*c, 8*c)
+    
+    def forward(self, x, flow):
+        x = self.conv1(x)
+        # flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f1 = warp(x, flow)        
+        x = self.conv2(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f2 = warp(x, flow)
+        x = self.conv3(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f3 = warp(x, flow)
+        x = self.conv4(x)
+        flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5
+        f4 = warp(x, flow)
+        return [f1, f2, f3, f4]
+    
+class Unet(nn.Module):
+    def __init__(self):
+        super(Unet, self).__init__()
+        self.down0 = Conv2(17, 2*c, 1)
+        self.down1 = Conv2(4*c, 4*c)
+        self.down2 = Conv2(8*c, 8*c)
+        self.down3 = Conv2(16*c, 16*c)
+        self.up0 = deconv(32*c, 8*c)
+        self.up1 = deconv(16*c, 4*c)
+        self.up2 = deconv(8*c, 2*c)
+        self.up3 = deconv(4*c, c)
+        self.conv = nn.Conv2d(c, 3, 3, 2, 1)
+
+    def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1):
+        s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1))
+        s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1))
+        s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1))
+        s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1))
+        x = self.up0(torch.cat((s3, c0[3], c1[3]), 1))
+        x = self.up1(torch.cat((x, s2), 1)) 
+        x = self.up2(torch.cat((x, s1), 1)) 
+        x = self.up3(torch.cat((x, s0), 1)) 
+        x = self.conv(x)
+        return torch.sigmoid(x)

ECCV2022-RIFE-main/model/warplayer.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+backwarp_tenGrid = {}
+
+
+def warp(tenInput, tenFlow):
+    k = (str(tenFlow.device), str(tenFlow.size()))
+    if k not in backwarp_tenGrid:
+        tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device).view(
+            1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1)
+        tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device).view(
+            1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3])
+        backwarp_tenGrid[k] = torch.cat(
+            [tenHorizontal, tenVertical], 1).to(device)
+
+    tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0),
+                         tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1)
+
+    g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1)
+    return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True)

ECCV2022-RIFE-main/requirements.txt

@@ -0,0 +1,7 @@
+numpy>=1.16, <=1.23.5
+tqdm>=4.35.0 
+sk-video>=1.1.10 
+torch>=1.6.0 
+opencv-python>=4.1.2 
+moviepy>=1.0.3 
+torchvision>=0.7.0

ECCV2022-RIFE-main/train.py

@@ -0,0 +1,155 @@
+import os
+import cv2
+import math
+import time
+import torch
+import torch.distributed as dist
+import numpy as np
+import random
+import argparse
+
+from model.RIFE import Model
+from dataset import *
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.tensorboard import SummaryWriter
+from torch.utils.data.distributed import DistributedSampler
+
+device = torch.device("cuda")
+
+log_path = 'train_log'
+
+def get_learning_rate(step):
+    if step < 2000:
+        mul = step / 2000.
+        return 3e-4 * mul
+    else:
+        mul = np.cos((step - 2000) / (args.epoch * args.step_per_epoch - 2000.) * math.pi) * 0.5 + 0.5
+        return (3e-4 - 3e-6) * mul + 3e-6
+
+def flow2rgb(flow_map_np):
+    h, w, _ = flow_map_np.shape
+    rgb_map = np.ones((h, w, 3)).astype(np.float32)
+    normalized_flow_map = flow_map_np / (np.abs(flow_map_np).max())
+    
+    rgb_map[:, :, 0] += normalized_flow_map[:, :, 0]
+    rgb_map[:, :, 1] -= 0.5 * (normalized_flow_map[:, :, 0] + normalized_flow_map[:, :, 1])
+    rgb_map[:, :, 2] += normalized_flow_map[:, :, 1]
+    return rgb_map.clip(0, 1)
+
+def train(model, local_rank):
+    if local_rank == 0:
+        writer = SummaryWriter('train')
+        writer_val = SummaryWriter('validate')
+    else:
+        writer = None
+        writer_val = None
+    step = 0
+    nr_eval = 0
+    dataset = VimeoDataset('train')
+    sampler = DistributedSampler(dataset)
+    train_data = DataLoader(dataset, batch_size=args.batch_size, num_workers=8, pin_memory=True, drop_last=True, sampler=sampler)
+    args.step_per_epoch = train_data.__len__()
+    dataset_val = VimeoDataset('validation')
+    val_data = DataLoader(dataset_val, batch_size=16, pin_memory=True, num_workers=8)
+    print('training...')
+    time_stamp = time.time()
+    for epoch in range(args.epoch):
+        sampler.set_epoch(epoch)
+        for i, data in enumerate(train_data):
+            data_time_interval = time.time() - time_stamp
+            time_stamp = time.time()
+            data_gpu, timestep = data
+            data_gpu = data_gpu.to(device, non_blocking=True) / 255.
+            timestep = timestep.to(device, non_blocking=True)
+            imgs = data_gpu[:, :6]
+            gt = data_gpu[:, 6:9]
+            learning_rate = get_learning_rate(step) * args.world_size / 4
+            pred, info = model.update(imgs, gt, learning_rate, training=True) # pass timestep if you are training RIFEm
+            train_time_interval = time.time() - time_stamp
+            time_stamp = time.time()
+            if step % 200 == 1 and local_rank == 0:
+                writer.add_scalar('learning_rate', learning_rate, step)
+                writer.add_scalar('loss/l1', info['loss_l1'], step)
+                writer.add_scalar('loss/tea', info['loss_tea'], step)
+                writer.add_scalar('loss/distill', info['loss_distill'], step)
+            if step % 1000 == 1 and local_rank == 0:
+                gt = (gt.permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8')
+                mask = (torch.cat((info['mask'], info['mask_tea']), 3).permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8')
+                pred = (pred.permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8')
+                merged_img = (info['merged_tea'].permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8')
+                flow0 = info['flow'].permute(0, 2, 3, 1).detach().cpu().numpy()
+                flow1 = info['flow_tea'].permute(0, 2, 3, 1).detach().cpu().numpy()
+                for i in range(5):
+                    imgs = np.concatenate((merged_img[i], pred[i], gt[i]), 1)[:, :, ::-1]
+                    writer.add_image(str(i) + '/img', imgs, step, dataformats='HWC')
+                    writer.add_image(str(i) + '/flow', np.concatenate((flow2rgb(flow0[i]), flow2rgb(flow1[i])), 1), step, dataformats='HWC')
+                    writer.add_image(str(i) + '/mask', mask[i], step, dataformats='HWC')
+                writer.flush()
+            if local_rank == 0:
+                print('epoch:{} {}/{} time:{:.2f}+{:.2f} loss_l1:{:.4e}'.format(epoch, i, args.step_per_epoch, data_time_interval, train_time_interval, info['loss_l1']))
+            step += 1
+        nr_eval += 1
+        if nr_eval % 5 == 0:
+            evaluate(model, val_data, step, local_rank, writer_val)
+        model.save_model(log_path, local_rank)    
+        dist.barrier()
+
+def evaluate(model, val_data, nr_eval, local_rank, writer_val):
+    loss_l1_list = []
+    loss_distill_list = []
+    loss_tea_list = []
+    psnr_list = []
+    psnr_list_teacher = []
+    time_stamp = time.time()
+    for i, data in enumerate(val_data):
+        data_gpu, timestep = data
+        data_gpu = data_gpu.to(device, non_blocking=True) / 255.        
+        imgs = data_gpu[:, :6]
+        gt = data_gpu[:, 6:9]
+        with torch.no_grad():
+            pred, info = model.update(imgs, gt, training=False)
+            merged_img = info['merged_tea']
+        loss_l1_list.append(info['loss_l1'].cpu().numpy())
+        loss_tea_list.append(info['loss_tea'].cpu().numpy())
+        loss_distill_list.append(info['loss_distill'].cpu().numpy())
+        for j in range(gt.shape[0]):
+            psnr = -10 * math.log10(torch.mean((gt[j] - pred[j]) * (gt[j] - pred[j])).cpu().data)
+            psnr_list.append(psnr)
+            psnr = -10 * math.log10(torch.mean((merged_img[j] - gt[j]) * (merged_img[j] - gt[j])).cpu().data)
+            psnr_list_teacher.append(psnr)
+        gt = (gt.permute(0, 2, 3, 1).cpu().numpy() * 255).astype('uint8')
+        pred = (pred.permute(0, 2, 3, 1).cpu().numpy() * 255).astype('uint8')
+        merged_img = (merged_img.permute(0, 2, 3, 1).cpu().numpy() * 255).astype('uint8')
+        flow0 = info['flow'].permute(0, 2, 3, 1).cpu().numpy()
+        flow1 = info['flow_tea'].permute(0, 2, 3, 1).cpu().numpy()
+        if i == 0 and local_rank == 0:
+            for j in range(10):
+                imgs = np.concatenate((merged_img[j], pred[j], gt[j]), 1)[:, :, ::-1]
+                writer_val.add_image(str(j) + '/img', imgs.copy(), nr_eval, dataformats='HWC')
+                writer_val.add_image(str(j) + '/flow', flow2rgb(flow0[j][:, :, ::-1]), nr_eval, dataformats='HWC')
+    
+    eval_time_interval = time.time() - time_stamp
+
+    if local_rank != 0:
+        return
+    writer_val.add_scalar('psnr', np.array(psnr_list).mean(), nr_eval)
+    writer_val.add_scalar('psnr_teacher', np.array(psnr_list_teacher).mean(), nr_eval)
+        
+if __name__ == "__main__":    
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--epoch', default=300, type=int)
+    parser.add_argument('--batch_size', default=16, type=int, help='minibatch size')
+    parser.add_argument('--local_rank', default=0, type=int, help='local rank')
+    parser.add_argument('--world_size', default=4, type=int, help='world size')
+    args = parser.parse_args()
+    torch.distributed.init_process_group(backend="nccl", world_size=args.world_size)
+    torch.cuda.set_device(args.local_rank)
+    seed = 1234
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = True
+    model = Model(args.local_rank)
+    train(model, args.local_rank)
+