├── docker ├── inference_img ├── inference_video └── Dockerfile ├── demo ├── I0_0.png ├── I0_1.png ├── I1_0.png ├── I1_1.png ├── I2_0.png ├── I2_1.png ├── intro.png ├── D0_slomo_clipped.gif ├── D2_slomo_clipped.gif ├── I0_slomo_clipped.gif └── I2_slomo_clipped.gif ├── .gitignore ├── requirements.txt ├── benchmark ├── testtime.py ├── MiddleBury_Other.py ├── UCF101.py ├── Vimeo90K.py ├── ATD12K.py ├── HD.py ├── yuv_frame_io.py └── HD_multi_4X.py ├── LICENSE ├── model ├── warplayer.py ├── laplacian.py ├── refine.py ├── refine_2R.py ├── RIFE.py ├── oldmodel │ ├── IFNet_HDv2.py │ ├── IFNet_HD.py │ ├── RIFE_HDv2.py │ └── RIFE_HD.py ├── IFNet.py ├── IFNet_2R.py ├── IFNet_m.py ├── loss.py └── pytorch_msssim │ └── __init__.py ├── Colab_demo.ipynb ├── inference_img.py ├── dataset.py ├── train.py ├── README.md └── inference_video.py /docker/inference_img: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | python3 /rife/inference_img.py $@ 3 | -------------------------------------------------------------------------------- /docker/inference_video: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | python3 /rife/inference_video.py $@ 3 | -------------------------------------------------------------------------------- /demo/I0_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I0_0.png -------------------------------------------------------------------------------- /demo/I0_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I0_1.png -------------------------------------------------------------------------------- /demo/I1_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I1_0.png -------------------------------------------------------------------------------- /demo/I1_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I1_1.png -------------------------------------------------------------------------------- /demo/I2_0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I2_0.png -------------------------------------------------------------------------------- /demo/I2_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I2_1.png -------------------------------------------------------------------------------- /demo/intro.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/intro.png -------------------------------------------------------------------------------- /demo/D0_slomo_clipped.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/D0_slomo_clipped.gif -------------------------------------------------------------------------------- /demo/D2_slomo_clipped.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/D2_slomo_clipped.gif -------------------------------------------------------------------------------- /demo/I0_slomo_clipped.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I0_slomo_clipped.gif -------------------------------------------------------------------------------- /demo/I2_slomo_clipped.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/hzwer/ECCV2022-RIFE/HEAD/demo/I2_slomo_clipped.gif -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.py~ 3 | *.py# 4 | 5 | *.pkl 6 | output/* 7 | train_log/* 8 | *.mp4 9 | 10 | test/ 11 | .idea/ 12 | *.npz 13 | 14 | *.zip 15 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy>=1.16, <=1.23.5 2 | tqdm>=4.35.0 3 | sk-video>=1.1.10 4 | torch>=1.6.0 5 | opencv-python>=4.1.2 6 | moviepy>=1.0.3 7 | torchvision>=0.7.0 8 | -------------------------------------------------------------------------------- /docker/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.8-slim 2 | 3 | # install deps 4 | RUN apt-get update && apt-get -y install \ 5 | bash ffmpeg 6 | 7 | # setup RIFE 8 | WORKDIR /rife 9 | COPY . . 10 | RUN pip3 install -r requirements.txt 11 | 12 | ADD docker/inference_img /usr/local/bin/inference_img 13 | RUN chmod +x /usr/local/bin/inference_img 14 | ADD docker/inference_video /usr/local/bin/inference_video 15 | RUN chmod +x /usr/local/bin/inference_video 16 | 17 | # add pre-trained models 18 | COPY train_log /rife/train_log 19 | 20 | WORKDIR /host 21 | ENTRYPOINT ["/bin/bash"] 22 | 23 | ENV NVIDIA_DRIVER_CAPABILITIES all -------------------------------------------------------------------------------- /benchmark/testtime.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import sys 3 | sys.path.append('.') 4 | import time 5 | import torch 6 | import torch.nn as nn 7 | from model.RIFE import Model 8 | 9 | model = Model() 10 | model.eval() 11 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 12 | torch.set_grad_enabled(False) 13 | if torch.cuda.is_available(): 14 | torch.backends.cudnn.enabled = True 15 | torch.backends.cudnn.benchmark = True 16 | 17 | I0 = torch.rand(1, 3, 480, 640).to(device) 18 | I1 = torch.rand(1, 3, 480, 640).to(device) 19 | with torch.no_grad(): 20 | for i in range(100): 21 | pred = model.inference(I0, I1) 22 | if torch.cuda.is_available(): 23 | torch.cuda.synchronize() 24 | time_stamp = time.time() 25 | for i in range(100): 26 | pred = model.inference(I0, I1) 27 | if torch.cuda.is_available(): 28 | torch.cuda.synchronize() 29 | print((time.time() - time_stamp) / 100) 30 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) Megvii Inc. 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /model/warplayer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 5 | backwarp_tenGrid = {} 6 | 7 | 8 | def warp(tenInput, tenFlow): 9 | k = (str(tenFlow.device), str(tenFlow.size())) 10 | if k not in backwarp_tenGrid: 11 | tenHorizontal = torch.linspace(-1.0, 1.0, tenFlow.shape[3], device=device).view( 12 | 1, 1, 1, tenFlow.shape[3]).expand(tenFlow.shape[0], -1, tenFlow.shape[2], -1) 13 | tenVertical = torch.linspace(-1.0, 1.0, tenFlow.shape[2], device=device).view( 14 | 1, 1, tenFlow.shape[2], 1).expand(tenFlow.shape[0], -1, -1, tenFlow.shape[3]) 15 | backwarp_tenGrid[k] = torch.cat( 16 | [tenHorizontal, tenVertical], 1).to(device) 17 | 18 | tenFlow = torch.cat([tenFlow[:, 0:1, :, :] / ((tenInput.shape[3] - 1.0) / 2.0), 19 | tenFlow[:, 1:2, :, :] / ((tenInput.shape[2] - 1.0) / 2.0)], 1) 20 | 21 | g = (backwarp_tenGrid[k] + tenFlow).permute(0, 2, 3, 1) 22 | return torch.nn.functional.grid_sample(input=tenInput, grid=g, mode='bilinear', padding_mode='border', align_corners=True) 23 | -------------------------------------------------------------------------------- /benchmark/MiddleBury_Other.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.append('.') 4 | import cv2 5 | import math 6 | import torch 7 | import argparse 8 | import numpy as np 9 | from torch.nn import functional as F 10 | from model.pytorch_msssim import ssim_matlab 11 | from model.RIFE import Model 12 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 13 | 14 | model = Model() 15 | model.load_model('train_log') 16 | model.eval() 17 | model.device() 18 | 19 | name = ['Beanbags', 'Dimetrodon', 'DogDance', 'Grove2', 'Grove3', 'Hydrangea', 'MiniCooper', 'RubberWhale', 'Urban2', 'Urban3', 'Venus', 'Walking'] 20 | IE_list = [] 21 | for i in name: 22 | i0 = cv2.imread('other-data/{}/frame10.png'.format(i)).transpose(2, 0, 1) / 255. 23 | i1 = cv2.imread('other-data/{}/frame11.png'.format(i)).transpose(2, 0, 1) / 255. 24 | gt = cv2.imread('other-gt-interp/{}/frame10i11.png'.format(i)) 25 | h, w = i0.shape[1], i0.shape[2] 26 | imgs = torch.zeros([1, 6, 480, 640]).to(device) 27 | ph = (480 - h) // 2 28 | pw = (640 - w) // 2 29 | imgs[:, :3, :h, :w] = torch.from_numpy(i0).unsqueeze(0).float().to(device) 30 | imgs[:, 3:, :h, :w] = torch.from_numpy(i1).unsqueeze(0).float().to(device) 31 | I0 = imgs[:, :3] 32 | I2 = imgs[:, 3:] 33 | pred = model.inference(I0, I2) 34 | out = pred[0].detach().cpu().numpy().transpose(1, 2, 0) 35 | out = np.round(out[:h, :w] * 255) 36 | IE_list.append(np.abs((out - gt * 1.0)).mean()) 37 | print(np.mean(IE_list)) 38 | -------------------------------------------------------------------------------- /benchmark/UCF101.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.append('.') 4 | import cv2 5 | import math 6 | import torch 7 | import argparse 8 | import numpy as np 9 | from torch.nn import functional as F 10 | from model.pytorch_msssim import ssim_matlab 11 | from model.RIFE import Model 12 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 13 | 14 | model = Model() 15 | model.load_model('train_log') 16 | model.eval() 17 | model.device() 18 | 19 | path = 'UCF101/ucf101_interp_ours/' 20 | dirs = os.listdir(path) 21 | psnr_list = [] 22 | ssim_list = [] 23 | print(len(dirs)) 24 | for d in dirs: 25 | img0 = (path + d + '/frame_00.png') 26 | img1 = (path + d + '/frame_02.png') 27 | gt = (path + d + '/frame_01_gt.png') 28 | img0 = (torch.tensor(cv2.imread(img0).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0) 29 | img1 = (torch.tensor(cv2.imread(img1).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0) 30 | gt = (torch.tensor(cv2.imread(gt).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0) 31 | pred = model.inference(img0, img1)[0] 32 | ssim = ssim_matlab(gt, torch.round(pred * 255).unsqueeze(0) / 255.).detach().cpu().numpy() 33 | out = pred.detach().cpu().numpy().transpose(1, 2, 0) 34 | out = np.round(out * 255) / 255. 35 | gt = gt[0].cpu().numpy().transpose(1, 2, 0) 36 | psnr = -10 * math.log10(((gt - out) * (gt - out)).mean()) 37 | psnr_list.append(psnr) 38 | ssim_list.append(ssim) 39 | print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list))) 40 | -------------------------------------------------------------------------------- /benchmark/Vimeo90K.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.append('.') 4 | import cv2 5 | import math 6 | import torch 7 | import argparse 8 | import numpy as np 9 | from torch.nn import functional as F 10 | from model.pytorch_msssim import ssim_matlab 11 | from model.RIFE import Model 12 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 13 | 14 | model = Model() 15 | model.load_model('train_log') 16 | model.eval() 17 | model.device() 18 | 19 | path = 'vimeo_interp_test/' 20 | f = open(path + 'tri_testlist.txt', 'r') 21 | psnr_list = [] 22 | ssim_list = [] 23 | for i in f: 24 | name = str(i).strip() 25 | if(len(name) <= 1): 26 | continue 27 | print(path + 'target/' + name + '/im1.png') 28 | I0 = cv2.imread(path + 'target/' + name + '/im1.png') 29 | I1 = cv2.imread(path + 'target/' + name + '/im2.png') 30 | I2 = cv2.imread(path + 'target/' + name + '/im3.png') 31 | I0 = (torch.tensor(I0.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0) 32 | I2 = (torch.tensor(I2.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0) 33 | mid = model.inference(I0, I2)[0] 34 | 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() 35 | mid = np.round((mid * 255).detach().cpu().numpy()).astype('uint8').transpose(1, 2, 0) / 255. 36 | I1 = I1 / 255. 37 | psnr = -10 * math.log10(((I1 - mid) * (I1 - mid)).mean()) 38 | psnr_list.append(psnr) 39 | ssim_list.append(ssim) 40 | print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list))) 41 | -------------------------------------------------------------------------------- /benchmark/ATD12K.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.append('.') 4 | import cv2 5 | import math 6 | import torch 7 | import argparse 8 | import numpy as np 9 | from torch.nn import functional as F 10 | from model.pytorch_msssim import ssim_matlab 11 | from model.RIFE import Model 12 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 13 | 14 | model = Model() 15 | model.load_model('train_log') 16 | model.eval() 17 | model.device() 18 | 19 | path = 'datasets/test_2k_540p/' 20 | dirs = os.listdir(path) 21 | psnr_list = [] 22 | ssim_list = [] 23 | print(len(dirs)) 24 | for d in dirs: 25 | img0 = (path + d + '/frame1.png') 26 | img1 = (path + d + '/frame3.png') 27 | gt = (path + d + '/frame2.png') 28 | img0 = (torch.tensor(cv2.imread(img0).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0) 29 | img1 = (torch.tensor(cv2.imread(img1).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0) 30 | gt = (torch.tensor(cv2.imread(gt).transpose(2, 0, 1) / 255.)).to(device).float().unsqueeze(0) 31 | pader = torch.nn.ReplicationPad2d([0, 0, 2, 2]) 32 | img0 = pader(img0) 33 | img1 = pader(img1) 34 | pred = model.inference(img0, img1)[0][:, 2:-2] 35 | ssim = ssim_matlab(gt, torch.round(pred * 255).unsqueeze(0) / 255.).detach().cpu().numpy() 36 | out = pred.detach().cpu().numpy().transpose(1, 2, 0) 37 | out = np.round(out * 255) / 255. 38 | gt = gt[0].cpu().numpy().transpose(1, 2, 0) 39 | psnr = -10 * math.log10(((gt - out) * (gt - out)).mean()) 40 | psnr_list.append(psnr) 41 | ssim_list.append(ssim) 42 | print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list))) 43 | -------------------------------------------------------------------------------- /model/laplacian.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 7 | 8 | import torch 9 | 10 | def gauss_kernel(size=5, channels=3): 11 | kernel = torch.tensor([[1., 4., 6., 4., 1], 12 | [4., 16., 24., 16., 4.], 13 | [6., 24., 36., 24., 6.], 14 | [4., 16., 24., 16., 4.], 15 | [1., 4., 6., 4., 1.]]) 16 | kernel /= 256. 17 | kernel = kernel.repeat(channels, 1, 1, 1) 18 | kernel = kernel.to(device) 19 | return kernel 20 | 21 | def downsample(x): 22 | return x[:, :, ::2, ::2] 23 | 24 | def upsample(x): 25 | cc = torch.cat([x, torch.zeros(x.shape[0], x.shape[1], x.shape[2], x.shape[3]).to(device)], dim=3) 26 | cc = cc.view(x.shape[0], x.shape[1], x.shape[2]*2, x.shape[3]) 27 | cc = cc.permute(0,1,3,2) 28 | cc = torch.cat([cc, torch.zeros(x.shape[0], x.shape[1], x.shape[3], x.shape[2]*2).to(device)], dim=3) 29 | cc = cc.view(x.shape[0], x.shape[1], x.shape[3]*2, x.shape[2]*2) 30 | x_up = cc.permute(0,1,3,2) 31 | return conv_gauss(x_up, 4*gauss_kernel(channels=x.shape[1])) 32 | 33 | def conv_gauss(img, kernel): 34 | img = torch.nn.functional.pad(img, (2, 2, 2, 2), mode='reflect') 35 | out = torch.nn.functional.conv2d(img, kernel, groups=img.shape[1]) 36 | return out 37 | 38 | def laplacian_pyramid(img, kernel, max_levels=3): 39 | current = img 40 | pyr = [] 41 | for level in range(max_levels): 42 | filtered = conv_gauss(current, kernel) 43 | down = downsample(filtered) 44 | up = upsample(down) 45 | diff = current-up 46 | pyr.append(diff) 47 | current = down 48 | return pyr 49 | 50 | class LapLoss(torch.nn.Module): 51 | def __init__(self, max_levels=5, channels=3): 52 | super(LapLoss, self).__init__() 53 | self.max_levels = max_levels 54 | self.gauss_kernel = gauss_kernel(channels=channels) 55 | 56 | def forward(self, input, target): 57 | pyr_input = laplacian_pyramid(img=input, kernel=self.gauss_kernel, max_levels=self.max_levels) 58 | pyr_target = laplacian_pyramid(img=target, kernel=self.gauss_kernel, max_levels=self.max_levels) 59 | return sum(torch.nn.functional.l1_loss(a, b) for a, b in zip(pyr_input, pyr_target)) 60 | -------------------------------------------------------------------------------- /model/refine.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | import torch.optim as optim 5 | import itertools 6 | from model.warplayer import warp 7 | import torch.nn.functional as F 8 | 9 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 10 | 11 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 12 | return nn.Sequential( 13 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 14 | padding=padding, dilation=dilation, bias=True), 15 | nn.PReLU(out_planes) 16 | ) 17 | 18 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 19 | return nn.Sequential( 20 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True), 21 | nn.PReLU(out_planes) 22 | ) 23 | 24 | class Conv2(nn.Module): 25 | def __init__(self, in_planes, out_planes, stride=2): 26 | super(Conv2, self).__init__() 27 | self.conv1 = conv(in_planes, out_planes, 3, stride, 1) 28 | self.conv2 = conv(out_planes, out_planes, 3, 1, 1) 29 | 30 | def forward(self, x): 31 | x = self.conv1(x) 32 | x = self.conv2(x) 33 | return x 34 | 35 | c = 16 36 | class Contextnet(nn.Module): 37 | def __init__(self): 38 | super(Contextnet, self).__init__() 39 | self.conv1 = Conv2(3, c) 40 | self.conv2 = Conv2(c, 2*c) 41 | self.conv3 = Conv2(2*c, 4*c) 42 | self.conv4 = Conv2(4*c, 8*c) 43 | 44 | def forward(self, x, flow): 45 | x = self.conv1(x) 46 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 47 | f1 = warp(x, flow) 48 | x = self.conv2(x) 49 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 50 | f2 = warp(x, flow) 51 | x = self.conv3(x) 52 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 53 | f3 = warp(x, flow) 54 | x = self.conv4(x) 55 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 56 | f4 = warp(x, flow) 57 | return [f1, f2, f3, f4] 58 | 59 | class Unet(nn.Module): 60 | def __init__(self): 61 | super(Unet, self).__init__() 62 | self.down0 = Conv2(17, 2*c) 63 | self.down1 = Conv2(4*c, 4*c) 64 | self.down2 = Conv2(8*c, 8*c) 65 | self.down3 = Conv2(16*c, 16*c) 66 | self.up0 = deconv(32*c, 8*c) 67 | self.up1 = deconv(16*c, 4*c) 68 | self.up2 = deconv(8*c, 2*c) 69 | self.up3 = deconv(4*c, c) 70 | self.conv = nn.Conv2d(c, 3, 3, 1, 1) 71 | 72 | def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): 73 | s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)) 74 | s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) 75 | s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) 76 | s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) 77 | x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) 78 | x = self.up1(torch.cat((x, s2), 1)) 79 | x = self.up2(torch.cat((x, s1), 1)) 80 | x = self.up3(torch.cat((x, s0), 1)) 81 | x = self.conv(x) 82 | return torch.sigmoid(x) 83 | -------------------------------------------------------------------------------- /benchmark/HD.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.append('.') 4 | import cv2 5 | import math 6 | import torch 7 | import argparse 8 | import numpy as np 9 | from torch.nn import functional as F 10 | from model.pytorch_msssim import ssim_matlab 11 | from model.RIFE import Model 12 | from skimage.color import rgb2yuv, yuv2rgb 13 | from yuv_frame_io import YUV_Read,YUV_Write 14 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 15 | 16 | model = Model() 17 | model.load_model('train_log') 18 | model.eval() 19 | model.device() 20 | 21 | name_list = [ 22 | ('HD_dataset/HD720p_GT/parkrun_1280x720_50.yuv', 720, 1280), 23 | ('HD_dataset/HD720p_GT/shields_1280x720_60.yuv', 720, 1280), 24 | ('HD_dataset/HD720p_GT/stockholm_1280x720_60.yuv', 720, 1280), 25 | ('HD_dataset/HD1080p_GT/BlueSky.yuv', 1080, 1920), 26 | ('HD_dataset/HD1080p_GT/Kimono1_1920x1080_24.yuv', 1080, 1920), 27 | ('HD_dataset/HD1080p_GT/ParkScene_1920x1080_24.yuv', 1080, 1920), 28 | ('HD_dataset/HD1080p_GT/sunflower_1080p25.yuv', 1080, 1920), 29 | ('HD_dataset/HD544p_GT/Sintel_Alley2_1280x544.yuv', 544, 1280), 30 | ('HD_dataset/HD544p_GT/Sintel_Market5_1280x544.yuv', 544, 1280), 31 | ('HD_dataset/HD544p_GT/Sintel_Temple1_1280x544.yuv', 544, 1280), 32 | ('HD_dataset/HD544p_GT/Sintel_Temple2_1280x544.yuv', 544, 1280), 33 | ] 34 | tot = 0. 35 | for data in name_list: 36 | psnr_list = [] 37 | name = data[0] 38 | h = data[1] 39 | w = data[2] 40 | if 'yuv' in name: 41 | Reader = YUV_Read(name, h, w, toRGB=True) 42 | else: 43 | Reader = cv2.VideoCapture(name) 44 | _, lastframe = Reader.read() 45 | # fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') 46 | # video = cv2.VideoWriter(name + '.mp4', fourcc, 30, (w, h)) 47 | for index in range(0, 100, 2): 48 | if 'yuv' in name: 49 | IMAGE1, success1 = Reader.read(index) 50 | gt, _ = Reader.read(index + 1) 51 | IMAGE2, success2 = Reader.read(index + 2) 52 | if not success2: 53 | break 54 | else: 55 | success1, gt = Reader.read() 56 | success2, frame = Reader.read() 57 | IMAGE1 = lastframe 58 | IMAGE2 = frame 59 | lastframe = frame 60 | if not success2: 61 | break 62 | I0 = torch.from_numpy(np.transpose(IMAGE1, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0) 63 | I1 = torch.from_numpy(np.transpose(IMAGE2, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0) 64 | 65 | if h == 720: 66 | pad = 24 67 | elif h == 1080: 68 | pad = 4 69 | else: 70 | pad = 16 71 | pader = torch.nn.ReplicationPad2d([0, 0, pad, pad]) 72 | I0 = pader(I0) 73 | I1 = pader(I1) 74 | with torch.no_grad(): 75 | pred = model.inference(I0, I1) 76 | pred = pred[:, :, pad: -pad] 77 | out = (np.round(pred[0].detach().cpu().numpy().transpose(1, 2, 0) * 255)).astype('uint8') 78 | # video.write(out) 79 | if 'yuv' in name: 80 | diff_rgb = 128.0 + rgb2yuv(gt / 255.)[:, :, 0] * 255 - rgb2yuv(out / 255.)[:, :, 0] * 255 81 | mse = np.mean((diff_rgb - 128.0) ** 2) 82 | PIXEL_MAX = 255.0 83 | psnr = 20 * math.log10(PIXEL_MAX / math.sqrt(mse)) 84 | else: 85 | psnr = skim.compare_psnr(gt, out) 86 | psnr_list.append(psnr) 87 | print(np.mean(psnr_list)) 88 | tot += np.mean(psnr_list) 89 | print('avg psnr', tot / len(name_list)) 90 | -------------------------------------------------------------------------------- /model/refine_2R.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | import torch.optim as optim 5 | import itertools 6 | from model.warplayer import warp 7 | from torch.nn.parallel import DistributedDataParallel as DDP 8 | import torch.nn.functional as F 9 | 10 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 11 | 12 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 13 | return nn.Sequential( 14 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 15 | padding=padding, dilation=dilation, bias=True), 16 | nn.PReLU(out_planes) 17 | ) 18 | 19 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 20 | return nn.Sequential( 21 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1, bias=True), 22 | nn.PReLU(out_planes) 23 | ) 24 | 25 | class Conv2(nn.Module): 26 | def __init__(self, in_planes, out_planes, stride=2): 27 | super(Conv2, self).__init__() 28 | self.conv1 = conv(in_planes, out_planes, 3, stride, 1) 29 | self.conv2 = conv(out_planes, out_planes, 3, 1, 1) 30 | 31 | def forward(self, x): 32 | x = self.conv1(x) 33 | x = self.conv2(x) 34 | return x 35 | 36 | c = 16 37 | class Contextnet(nn.Module): 38 | def __init__(self): 39 | super(Contextnet, self).__init__() 40 | self.conv1 = Conv2(3, c, 1) 41 | self.conv2 = Conv2(c, 2*c) 42 | self.conv3 = Conv2(2*c, 4*c) 43 | self.conv4 = Conv2(4*c, 8*c) 44 | 45 | def forward(self, x, flow): 46 | x = self.conv1(x) 47 | # flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 48 | f1 = warp(x, flow) 49 | x = self.conv2(x) 50 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 51 | f2 = warp(x, flow) 52 | x = self.conv3(x) 53 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 54 | f3 = warp(x, flow) 55 | x = self.conv4(x) 56 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False, recompute_scale_factor=False) * 0.5 57 | f4 = warp(x, flow) 58 | return [f1, f2, f3, f4] 59 | 60 | class Unet(nn.Module): 61 | def __init__(self): 62 | super(Unet, self).__init__() 63 | self.down0 = Conv2(17, 2*c, 1) 64 | self.down1 = Conv2(4*c, 4*c) 65 | self.down2 = Conv2(8*c, 8*c) 66 | self.down3 = Conv2(16*c, 16*c) 67 | self.up0 = deconv(32*c, 8*c) 68 | self.up1 = deconv(16*c, 4*c) 69 | self.up2 = deconv(8*c, 2*c) 70 | self.up3 = deconv(4*c, c) 71 | self.conv = nn.Conv2d(c, 3, 3, 2, 1) 72 | 73 | def forward(self, img0, img1, warped_img0, warped_img1, mask, flow, c0, c1): 74 | s0 = self.down0(torch.cat((img0, img1, warped_img0, warped_img1, mask, flow), 1)) 75 | s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) 76 | s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) 77 | s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) 78 | x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) 79 | x = self.up1(torch.cat((x, s2), 1)) 80 | x = self.up2(torch.cat((x, s1), 1)) 81 | x = self.up3(torch.cat((x, s0), 1)) 82 | x = self.conv(x) 83 | return torch.sigmoid(x) 84 | -------------------------------------------------------------------------------- /model/RIFE.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from torch.optim import AdamW 5 | import torch.optim as optim 6 | import itertools 7 | from model.warplayer import warp 8 | from torch.nn.parallel import DistributedDataParallel as DDP 9 | from model.IFNet import * 10 | from model.IFNet_m import * 11 | import torch.nn.functional as F 12 | from model.loss import * 13 | from model.laplacian import * 14 | from model.refine import * 15 | 16 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 17 | 18 | class Model: 19 | def __init__(self, local_rank=-1, arbitrary=False): 20 | if arbitrary == True: 21 | self.flownet = IFNet_m() 22 | else: 23 | self.flownet = IFNet() 24 | self.device() 25 | self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-3) # use large weight decay may avoid NaN loss 26 | self.epe = EPE() 27 | self.lap = LapLoss() 28 | self.sobel = SOBEL() 29 | if local_rank != -1: 30 | self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank) 31 | 32 | def train(self): 33 | self.flownet.train() 34 | 35 | def eval(self): 36 | self.flownet.eval() 37 | 38 | def device(self): 39 | self.flownet.to(device) 40 | 41 | def load_model(self, path, rank=0): 42 | def convert(param): 43 | return { 44 | k.replace("module.", ""): v 45 | for k, v in param.items() 46 | if "module." in k 47 | } 48 | 49 | if rank <= 0: 50 | self.flownet.load_state_dict(convert(torch.load('{}/flownet.pkl'.format(path)))) 51 | 52 | def save_model(self, path, rank=0): 53 | if rank == 0: 54 | torch.save(self.flownet.state_dict(),'{}/flownet.pkl'.format(path)) 55 | 56 | def inference(self, img0, img1, scale=1, scale_list=None, TTA=False, timestep=0.5): 57 | if scale_list is None: 58 | scale_list = [4, 2, 1] 59 | for i in range(3): 60 | scale_list[i] = scale_list[i] * 1.0 / scale 61 | imgs = torch.cat((img0, img1), 1) 62 | flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(imgs, scale_list, timestep=timestep) 63 | if TTA == False: 64 | return merged[2] 65 | else: 66 | flow2, mask2, merged2, flow_teacher2, merged_teacher2, loss_distill2 = self.flownet(imgs.flip(2).flip(3), scale_list, timestep=timestep) 67 | return (merged[2] + merged2[2].flip(2).flip(3)) / 2 68 | 69 | def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): 70 | for param_group in self.optimG.param_groups: 71 | param_group['lr'] = learning_rate 72 | img0 = imgs[:, :3] 73 | img1 = imgs[:, 3:] 74 | if training: 75 | self.train() 76 | else: 77 | self.eval() 78 | flow, mask, merged, flow_teacher, merged_teacher, loss_distill = self.flownet(torch.cat((imgs, gt), 1), scale=[4, 2, 1]) 79 | loss_l1 = (self.lap(merged[2], gt)).mean() 80 | loss_tea = (self.lap(merged_teacher, gt)).mean() 81 | if training: 82 | self.optimG.zero_grad() 83 | 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 84 | loss_G.backward() 85 | self.optimG.step() 86 | else: 87 | flow_teacher = flow[2] 88 | return merged[2], { 89 | 'merged_tea': merged_teacher, 90 | 'mask': mask, 91 | 'mask_tea': mask, 92 | 'flow': flow[2][:, :2], 93 | 'flow_tea': flow_teacher, 94 | 'loss_l1': loss_l1, 95 | 'loss_tea': loss_tea, 96 | 'loss_distill': loss_distill, 97 | } 98 | -------------------------------------------------------------------------------- /Colab_demo.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": { 6 | "colab_type": "text", 7 | "id": "view-in-github" 8 | }, 9 | "source": [ 10 | "\"Open" 11 | ] 12 | }, 13 | { 14 | "cell_type": "code", 15 | "execution_count": null, 16 | "metadata": { 17 | "id": "FypCcZkNNt2p" 18 | }, 19 | "outputs": [], 20 | "source": [ 21 | "!git clone https://github.com/hzwer/arXiv2020-RIFE" 22 | ] 23 | }, 24 | { 25 | "cell_type": "code", 26 | "execution_count": null, 27 | "metadata": { 28 | "id": "1wysVHxoN54f" 29 | }, 30 | "outputs": [], 31 | "source": [ 32 | "!mkdir /content/arXiv2020-RIFE/train_log\n", 33 | "%cd /content/arXiv2020-RIFE/train_log\n", 34 | "!gdown --id 1APIzVeI-4ZZCEuIRE1m6WYfSCaOsi_7_\n", 35 | "!7z e RIFE_trained_model_v3.6.zip" 36 | ] 37 | }, 38 | { 39 | "cell_type": "code", 40 | "execution_count": null, 41 | "metadata": { 42 | "id": "AhbHfRBJRAUt" 43 | }, 44 | "outputs": [], 45 | "source": [ 46 | "%cd /content/arXiv2020-RIFE/\n", 47 | "!gdown --id 1i3xlKb7ax7Y70khcTcuePi6E7crO_dFc\n", 48 | "!pip install git+https://github.com/rk-exxec/scikit-video.git@numpy_deprecation" 49 | ] 50 | }, 51 | { 52 | "cell_type": "markdown", 53 | "metadata": { 54 | "id": "rirngW5uRMdg" 55 | }, 56 | "source": [ 57 | "Please upload your video to content/arXiv2020-RIFE/video.mp4, or use our demo video." 58 | ] 59 | }, 60 | { 61 | "cell_type": "code", 62 | "execution_count": null, 63 | "metadata": { 64 | "id": "dnLn4aHHPzN3" 65 | }, 66 | "outputs": [], 67 | "source": [ 68 | "!nvidia-smi\n", 69 | "!python3 inference_video.py --exp=2 --video=demo.mp4 --montage" 70 | ] 71 | }, 72 | { 73 | "cell_type": "markdown", 74 | "metadata": { 75 | "id": "77KK6lxHgJhf" 76 | }, 77 | "source": [ 78 | "Our demo.mp4 is 25FPS. You can adjust the parameters for your own perference.\n", 79 | "For example: \n", 80 | "--fps=60 --exp=1 --video=mydemo.avi --png" 81 | ] 82 | }, 83 | { 84 | "cell_type": "code", 85 | "execution_count": null, 86 | "metadata": { 87 | "cellView": "code", 88 | "id": "0zIBbVE3UfUD" 89 | }, 90 | "outputs": [], 91 | "source": [ 92 | "from IPython.display import display, Image\n", 93 | "import moviepy.editor as mpy\n", 94 | "display(mpy.ipython_display('demo_4X_100fps.mp4', height=256, max_duration=100.))" 95 | ] 96 | }, 97 | { 98 | "cell_type": "code", 99 | "execution_count": null, 100 | "metadata": { 101 | "id": "tWkJCNgP3zXA" 102 | }, 103 | "outputs": [], 104 | "source": [ 105 | "!python3 inference_img.py --img demo/I0_0.png demo/I0_1.png\n", 106 | "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", 107 | "# Image interpolation" 108 | ] 109 | } 110 | ], 111 | "metadata": { 112 | "accelerator": "GPU", 113 | "colab": { 114 | "include_colab_link": true, 115 | "name": "Untitled0.ipynb", 116 | "provenance": [] 117 | }, 118 | "kernelspec": { 119 | "display_name": "Python 3", 120 | "name": "python3" 121 | } 122 | }, 123 | "nbformat": 4, 124 | "nbformat_minor": 0 125 | } 126 | -------------------------------------------------------------------------------- /model/oldmodel/IFNet_HDv2.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from model.warplayer import warp 6 | 7 | 8 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 9 | 10 | def conv_wo_act(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 11 | return nn.Sequential( 12 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 13 | padding=padding, dilation=dilation, bias=True), 14 | ) 15 | 16 | 17 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 18 | return nn.Sequential( 19 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 20 | padding=padding, dilation=dilation, bias=True), 21 | nn.PReLU(out_planes) 22 | ) 23 | 24 | class IFBlock(nn.Module): 25 | def __init__(self, in_planes, scale=1, c=64): 26 | super(IFBlock, self).__init__() 27 | self.scale = scale 28 | self.conv0 = nn.Sequential( 29 | conv(in_planes, c, 3, 2, 1), 30 | conv(c, 2*c, 3, 2, 1), 31 | ) 32 | self.convblock = nn.Sequential( 33 | conv(2*c, 2*c), 34 | conv(2*c, 2*c), 35 | conv(2*c, 2*c), 36 | conv(2*c, 2*c), 37 | conv(2*c, 2*c), 38 | conv(2*c, 2*c), 39 | ) 40 | self.conv1 = nn.ConvTranspose2d(2*c, 4, 4, 2, 1) 41 | 42 | def forward(self, x): 43 | if self.scale != 1: 44 | x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear", 45 | align_corners=False) 46 | x = self.conv0(x) 47 | x = self.convblock(x) 48 | x = self.conv1(x) 49 | flow = x 50 | if self.scale != 1: 51 | flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear", 52 | align_corners=False) 53 | return flow 54 | 55 | 56 | class IFNet(nn.Module): 57 | def __init__(self): 58 | super(IFNet, self).__init__() 59 | self.block0 = IFBlock(6, scale=8, c=192) 60 | self.block1 = IFBlock(10, scale=4, c=128) 61 | self.block2 = IFBlock(10, scale=2, c=96) 62 | self.block3 = IFBlock(10, scale=1, c=48) 63 | 64 | def forward(self, x, scale=1.0): 65 | if scale != 1.0: 66 | x = F.interpolate(x, scale_factor=scale, mode="bilinear", align_corners=False) 67 | flow0 = self.block0(x) 68 | F1 = flow0 69 | F1_large = F.interpolate(F1, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0 70 | warped_img0 = warp(x[:, :3], F1_large[:, :2]) 71 | warped_img1 = warp(x[:, 3:], F1_large[:, 2:4]) 72 | flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1_large), 1)) 73 | F2 = (flow0 + flow1) 74 | F2_large = F.interpolate(F2, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0 75 | warped_img0 = warp(x[:, :3], F2_large[:, :2]) 76 | warped_img1 = warp(x[:, 3:], F2_large[:, 2:4]) 77 | flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2_large), 1)) 78 | F3 = (flow0 + flow1 + flow2) 79 | F3_large = F.interpolate(F3, scale_factor=2.0, mode="bilinear", align_corners=False) * 2.0 80 | warped_img0 = warp(x[:, :3], F3_large[:, :2]) 81 | warped_img1 = warp(x[:, 3:], F3_large[:, 2:4]) 82 | flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3_large), 1)) 83 | F4 = (flow0 + flow1 + flow2 + flow3) 84 | if scale != 1.0: 85 | F4 = F.interpolate(F4, scale_factor=1 / scale, mode="bilinear", align_corners=False) / scale 86 | return F4, [F1, F2, F3, F4] 87 | 88 | if __name__ == '__main__': 89 | img0 = torch.zeros(3, 3, 256, 256).float().to(device) 90 | img1 = torch.tensor(np.random.normal( 91 | 0, 1, (3, 3, 256, 256))).float().to(device) 92 | imgs = torch.cat((img0, img1), 1) 93 | flownet = IFNet() 94 | flow, _ = flownet(imgs) 95 | print(flow.shape) 96 | -------------------------------------------------------------------------------- /benchmark/yuv_frame_io.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import getopt 3 | import math 4 | import numpy 5 | import random 6 | import logging 7 | import numpy as np 8 | from skimage.color import rgb2yuv, yuv2rgb 9 | from PIL import Image 10 | import os 11 | from shutil import copyfile 12 | 13 | class YUV_Read(): 14 | def __init__(self, filepath, h, w, format='yuv420', toRGB=True): 15 | 16 | self.h = h 17 | self.w = w 18 | 19 | self.fp = open(filepath, 'rb') 20 | 21 | if format == 'yuv420': 22 | self.frame_length = int(1.5 * h * w) 23 | self.Y_length = h * w 24 | self.Uv_length = int(0.25 * h * w) 25 | else: 26 | pass 27 | self.toRGB = toRGB 28 | 29 | def read(self, offset_frame=None): 30 | if not offset_frame == None: 31 | self.fp.seek(offset_frame * self.frame_length, 0) 32 | 33 | Y = np.fromfile(self.fp, np.uint8, count=self.Y_length) 34 | U = np.fromfile(self.fp, np.uint8, count=self.Uv_length) 35 | V = np.fromfile(self.fp, np.uint8, count=self.Uv_length) 36 | if Y.size < self.Y_length or \ 37 | U.size < self.Uv_length or \ 38 | V.size < self.Uv_length: 39 | return None, False 40 | 41 | Y = np.reshape(Y, [self.w, self.h], order='F') 42 | Y = np.transpose(Y) 43 | 44 | U = np.reshape(U, [int(self.w / 2), int(self.h / 2)], order='F') 45 | U = np.transpose(U) 46 | 47 | V = np.reshape(V, [int(self.w / 2), int(self.h / 2)], order='F') 48 | V = np.transpose(V) 49 | 50 | U = np.array(Image.fromarray(U).resize([self.w, self.h])) 51 | V = np.array(Image.fromarray(V).resize([self.w, self.h])) 52 | 53 | if self.toRGB: 54 | Y = Y / 255.0 55 | U = U / 255.0 - 0.5 56 | V = V / 255.0 - 0.5 57 | 58 | self.YUV = np.stack((Y, U, V), axis=-1) 59 | self.RGB = (255.0 * np.clip(yuv2rgb(self.YUV), 0.0, 1.0)).astype('uint8') 60 | 61 | self.YUV = None 62 | return self.RGB, True 63 | else: 64 | self.YUV = np.stack((Y, U, V), axis=-1) 65 | return self.YUV, True 66 | 67 | def close(self): 68 | self.fp.close() 69 | 70 | 71 | class YUV_Write(): 72 | def __init__(self, filepath, fromRGB=True): 73 | if os.path.exists(filepath): 74 | print(filepath) 75 | 76 | self.fp = open(filepath, 'wb') 77 | self.fromRGB = fromRGB 78 | 79 | def write(self, Frame): 80 | 81 | self.h = Frame.shape[0] 82 | self.w = Frame.shape[1] 83 | c = Frame.shape[2] 84 | 85 | assert c == 3 86 | if format == 'yuv420': 87 | self.frame_length = int(1.5 * self.h * self.w) 88 | self.Y_length = self.h * self.w 89 | self.Uv_length = int(0.25 * self.h * self.w) 90 | else: 91 | pass 92 | if self.fromRGB: 93 | Frame = Frame / 255.0 94 | YUV = rgb2yuv(Frame) 95 | Y, U, V = np.dsplit(YUV, 3) 96 | Y = Y[:, :, 0] 97 | U = U[:, :, 0] 98 | V = V[:, :, 0] 99 | U = np.clip(U + 0.5, 0.0, 1.0) 100 | V = np.clip(V + 0.5, 0.0, 1.0) 101 | 102 | U = U[::2, ::2] # imresize(U,[int(self.h/2),int(self.w/2)],interp = 'nearest') 103 | V = V[::2, ::2] # imresize(V ,[int(self.h/2),int(self.w/2)],interp = 'nearest') 104 | Y = (255.0 * Y).astype('uint8') 105 | U = (255.0 * U).astype('uint8') 106 | V = (255.0 * V).astype('uint8') 107 | else: 108 | YUV = Frame 109 | Y = YUV[:, :, 0] 110 | U = YUV[::2, ::2, 1] 111 | V = YUV[::2, ::2, 2] 112 | 113 | Y = Y.flatten() # the first order is 0-dimension so don't need to transpose before flatten 114 | U = U.flatten() 115 | V = V.flatten() 116 | 117 | Y.tofile(self.fp) 118 | U.tofile(self.fp) 119 | V.tofile(self.fp) 120 | 121 | return True 122 | 123 | def close(self): 124 | self.fp.close() 125 | -------------------------------------------------------------------------------- /benchmark/HD_multi_4X.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | sys.path.append('.') 4 | import cv2 5 | import math 6 | import torch 7 | import argparse 8 | import numpy as np 9 | from torch.nn import functional as F 10 | from model.pytorch_msssim import ssim_matlab 11 | from model.RIFE import Model 12 | from skimage.color import rgb2yuv, yuv2rgb 13 | from yuv_frame_io import YUV_Read,YUV_Write 14 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 15 | 16 | model = Model(arbitrary=True) 17 | model.load_model('RIFE_m_train_log') 18 | model.eval() 19 | model.device() 20 | 21 | name_list = [ 22 | ('HD_dataset/HD720p_GT/parkrun_1280x720_50.yuv', 720, 1280), 23 | ('HD_dataset/HD720p_GT/shields_1280x720_60.yuv', 720, 1280), 24 | ('HD_dataset/HD720p_GT/stockholm_1280x720_60.yuv', 720, 1280), 25 | ('HD_dataset/HD1080p_GT/BlueSky.yuv', 1080, 1920), 26 | ('HD_dataset/HD1080p_GT/Kimono1_1920x1080_24.yuv', 1080, 1920), 27 | ('HD_dataset/HD1080p_GT/ParkScene_1920x1080_24.yuv', 1080, 1920), 28 | ('HD_dataset/HD1080p_GT/sunflower_1080p25.yuv', 1080, 1920), 29 | ('HD_dataset/HD544p_GT/Sintel_Alley2_1280x544.yuv', 544, 1280), 30 | ('HD_dataset/HD544p_GT/Sintel_Market5_1280x544.yuv', 544, 1280), 31 | ('HD_dataset/HD544p_GT/Sintel_Temple1_1280x544.yuv', 544, 1280), 32 | ('HD_dataset/HD544p_GT/Sintel_Temple2_1280x544.yuv', 544, 1280), 33 | ] 34 | def inference(I0, I1, pad, multi=2, arbitrary=True): 35 | img = [I0, I1] 36 | if not arbitrary: 37 | for i in range(multi): 38 | res = [I0] 39 | for j in range(len(img) - 1): 40 | res.append(model.inference(img[j], img[j + 1])) 41 | res.append(img[j + 1]) 42 | img = res 43 | else: 44 | img = [I0] 45 | p = 2**multi 46 | for i in range(p-1): 47 | img.append(model.inference(I0, I1, timestep=(i+1)*(1./p))) 48 | img.append(I1) 49 | for i in range(len(img)): 50 | img[i] = img[i][0][:, pad: -pad] 51 | return img[1: -1] 52 | 53 | tot = [] 54 | for data in name_list: 55 | psnr_list = [] 56 | name = data[0] 57 | h = data[1] 58 | w = data[2] 59 | if 'yuv' in name: 60 | Reader = YUV_Read(name, h, w, toRGB=True) 61 | else: 62 | Reader = cv2.VideoCapture(name) 63 | _, lastframe = Reader.read() 64 | # fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') 65 | # video = cv2.VideoWriter(name + '.mp4', fourcc, 30, (w, h)) 66 | for index in range(0, 100, 4): 67 | gt = [] 68 | if 'yuv' in name: 69 | IMAGE1, success1 = Reader.read(index) 70 | IMAGE2, success2 = Reader.read(index + 4) 71 | if not success2: 72 | break 73 | for i in range(1, 4): 74 | tmp, _ = Reader.read(index + i) 75 | gt.append(tmp) 76 | else: 77 | print('Not Implement') 78 | I0 = torch.from_numpy(np.transpose(IMAGE1, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0) 79 | I1 = torch.from_numpy(np.transpose(IMAGE2, (2,0,1)).astype("float32") / 255.).cuda().unsqueeze(0) 80 | 81 | if h == 720: 82 | pad = 24 83 | elif h == 1080: 84 | pad = 4 85 | else: 86 | pad = 16 87 | pader = torch.nn.ReplicationPad2d([0, 0, pad, pad]) 88 | I0 = pader(I0) 89 | I1 = pader(I1) 90 | with torch.no_grad(): 91 | pred = inference(I0, I1, pad) 92 | for i in range(4 - 1): 93 | out = (np.round(pred[i].detach().cpu().numpy().transpose(1, 2, 0) * 255)).astype('uint8') 94 | if 'yuv' in name: 95 | diff_rgb = 128.0 + rgb2yuv(gt[i] / 255.)[:, :, 0] * 255 - rgb2yuv(out / 255.)[:, :, 0] * 255 96 | mse = np.mean((diff_rgb - 128.0) ** 2) 97 | PIXEL_MAX = 255.0 98 | psnr = 20 * math.log10(PIXEL_MAX / math.sqrt(mse)) 99 | else: 100 | print('Not Implement') 101 | psnr_list.append(psnr) 102 | print(np.mean(psnr_list)) 103 | tot.append(np.mean(psnr_list)) 104 | 105 | print('PSNR: {}(544*1280), {}(720p), {}(1080p)'.format(np.mean(tot[7:11]), np.mean(tot[:3]), np.mean(tot[3:7]))) 106 | -------------------------------------------------------------------------------- /inference_img.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import torch 4 | import argparse 5 | from torch.nn import functional as F 6 | import warnings 7 | warnings.filterwarnings("ignore") 8 | 9 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 10 | torch.set_grad_enabled(False) 11 | if torch.cuda.is_available(): 12 | torch.backends.cudnn.enabled = True 13 | torch.backends.cudnn.benchmark = True 14 | 15 | parser = argparse.ArgumentParser(description='Interpolation for a pair of images') 16 | parser.add_argument('--img', dest='img', nargs=2, required=True) 17 | parser.add_argument('--exp', default=4, type=int) 18 | parser.add_argument('--ratio', default=0, type=float, help='inference ratio between two images with 0 - 1 range') 19 | parser.add_argument('--rthreshold', default=0.02, type=float, help='returns image when actual ratio falls in given range threshold') 20 | parser.add_argument('--rmaxcycles', default=8, type=int, help='limit max number of bisectional cycles') 21 | parser.add_argument('--model', dest='modelDir', type=str, default='train_log', help='directory with trained model files') 22 | 23 | args = parser.parse_args() 24 | 25 | try: 26 | try: 27 | try: 28 | from model.RIFE_HDv2 import Model 29 | model = Model() 30 | model.load_model(args.modelDir, -1) 31 | print("Loaded v2.x HD model.") 32 | except: 33 | from train_log.RIFE_HDv3 import Model 34 | model = Model() 35 | model.load_model(args.modelDir, -1) 36 | print("Loaded v3.x HD model.") 37 | except: 38 | from model.RIFE_HD import Model 39 | model = Model() 40 | model.load_model(args.modelDir, -1) 41 | print("Loaded v1.x HD model") 42 | except: 43 | from model.RIFE import Model 44 | model = Model() 45 | model.load_model(args.modelDir, -1) 46 | print("Loaded ArXiv-RIFE model") 47 | model.eval() 48 | model.device() 49 | 50 | if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'): 51 | img0 = cv2.imread(args.img[0], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH) 52 | img1 = cv2.imread(args.img[1], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH) 53 | img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device)).unsqueeze(0) 54 | img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device)).unsqueeze(0) 55 | 56 | else: 57 | img0 = cv2.imread(args.img[0], cv2.IMREAD_UNCHANGED) 58 | img1 = cv2.imread(args.img[1], cv2.IMREAD_UNCHANGED) 59 | img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0) 60 | img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device) / 255.).unsqueeze(0) 61 | 62 | n, c, h, w = img0.shape 63 | ph = ((h - 1) // 32 + 1) * 32 64 | pw = ((w - 1) // 32 + 1) * 32 65 | padding = (0, pw - w, 0, ph - h) 66 | img0 = F.pad(img0, padding) 67 | img1 = F.pad(img1, padding) 68 | 69 | 70 | if args.ratio: 71 | img_list = [img0] 72 | img0_ratio = 0.0 73 | img1_ratio = 1.0 74 | if args.ratio <= img0_ratio + args.rthreshold / 2: 75 | middle = img0 76 | elif args.ratio >= img1_ratio - args.rthreshold / 2: 77 | middle = img1 78 | else: 79 | tmp_img0 = img0 80 | tmp_img1 = img1 81 | for inference_cycle in range(args.rmaxcycles): 82 | middle = model.inference(tmp_img0, tmp_img1) 83 | middle_ratio = ( img0_ratio + img1_ratio ) / 2 84 | if args.ratio - (args.rthreshold / 2) <= middle_ratio <= args.ratio + (args.rthreshold / 2): 85 | break 86 | if args.ratio > middle_ratio: 87 | tmp_img0 = middle 88 | img0_ratio = middle_ratio 89 | else: 90 | tmp_img1 = middle 91 | img1_ratio = middle_ratio 92 | img_list.append(middle) 93 | img_list.append(img1) 94 | else: 95 | img_list = [img0, img1] 96 | for i in range(args.exp): 97 | tmp = [] 98 | for j in range(len(img_list) - 1): 99 | mid = model.inference(img_list[j], img_list[j + 1]) 100 | tmp.append(img_list[j]) 101 | tmp.append(mid) 102 | tmp.append(img1) 103 | img_list = tmp 104 | 105 | if not os.path.exists('output'): 106 | os.mkdir('output') 107 | for i in range(len(img_list)): 108 | if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'): 109 | 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]) 110 | else: 111 | cv2.imwrite('output/img{}.png'.format(i), (img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w]) 112 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import ast 4 | import torch 5 | import numpy as np 6 | import random 7 | from torch.utils.data import DataLoader, Dataset 8 | 9 | cv2.setNumThreads(1) 10 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 11 | class VimeoDataset(Dataset): 12 | def __init__(self, dataset_name, batch_size=32): 13 | self.batch_size = batch_size 14 | self.dataset_name = dataset_name 15 | self.h = 256 16 | self.w = 448 17 | self.data_root = 'vimeo_triplet' 18 | self.image_root = os.path.join(self.data_root, 'sequences') 19 | train_fn = os.path.join(self.data_root, 'tri_trainlist.txt') 20 | test_fn = os.path.join(self.data_root, 'tri_testlist.txt') 21 | with open(train_fn, 'r') as f: 22 | self.trainlist = f.read().splitlines() 23 | with open(test_fn, 'r') as f: 24 | self.testlist = f.read().splitlines() 25 | self.load_data() 26 | 27 | def __len__(self): 28 | return len(self.meta_data) 29 | 30 | def load_data(self): 31 | cnt = int(len(self.trainlist) * 0.95) 32 | if self.dataset_name == 'train': 33 | self.meta_data = self.trainlist[:cnt] 34 | elif self.dataset_name == 'test': 35 | self.meta_data = self.testlist 36 | else: 37 | self.meta_data = self.trainlist[cnt:] 38 | 39 | def crop(self, img0, gt, img1, h, w): 40 | ih, iw, _ = img0.shape 41 | x = np.random.randint(0, ih - h + 1) 42 | y = np.random.randint(0, iw - w + 1) 43 | img0 = img0[x:x+h, y:y+w, :] 44 | img1 = img1[x:x+h, y:y+w, :] 45 | gt = gt[x:x+h, y:y+w, :] 46 | return img0, gt, img1 47 | 48 | def getimg(self, index): 49 | imgpath = os.path.join(self.image_root, self.meta_data[index]) 50 | imgpaths = [imgpath + '/im1.png', imgpath + '/im2.png', imgpath + '/im3.png'] 51 | 52 | # Load images 53 | img0 = cv2.imread(imgpaths[0]) 54 | gt = cv2.imread(imgpaths[1]) 55 | img1 = cv2.imread(imgpaths[2]) 56 | timestep = 0.5 57 | return img0, gt, img1, timestep 58 | 59 | # RIFEm with Vimeo-Septuplet 60 | # imgpaths = [imgpath + '/im1.png', imgpath + '/im2.png', imgpath + '/im3.png', imgpath + '/im4.png', imgpath + '/im5.png', imgpath + '/im6.png', imgpath + '/im7.png'] 61 | # ind = [0, 1, 2, 3, 4, 5, 6] 62 | # random.shuffle(ind) 63 | # ind = ind[:3] 64 | # ind.sort() 65 | # img0 = cv2.imread(imgpaths[ind[0]]) 66 | # gt = cv2.imread(imgpaths[ind[1]]) 67 | # img1 = cv2.imread(imgpaths[ind[2]]) 68 | # timestep = (ind[1] - ind[0]) * 1.0 / (ind[2] - ind[0] + 1e-6) 69 | 70 | def __getitem__(self, index): 71 | img0, gt, img1, timestep = self.getimg(index) 72 | if self.dataset_name == 'train': 73 | img0, gt, img1 = self.crop(img0, gt, img1, 224, 224) 74 | if random.uniform(0, 1) < 0.5: 75 | img0 = img0[:, :, ::-1] 76 | img1 = img1[:, :, ::-1] 77 | gt = gt[:, :, ::-1] 78 | if random.uniform(0, 1) < 0.5: 79 | img0 = img0[::-1] 80 | img1 = img1[::-1] 81 | gt = gt[::-1] 82 | if random.uniform(0, 1) < 0.5: 83 | img0 = img0[:, ::-1] 84 | img1 = img1[:, ::-1] 85 | gt = gt[:, ::-1] 86 | if random.uniform(0, 1) < 0.5: 87 | tmp = img1 88 | img1 = img0 89 | img0 = tmp 90 | timestep = 1 - timestep 91 | # random rotation 92 | p = random.uniform(0, 1) 93 | if p < 0.25: 94 | img0 = cv2.rotate(img0, cv2.ROTATE_90_CLOCKWISE) 95 | gt = cv2.rotate(gt, cv2.ROTATE_90_CLOCKWISE) 96 | img1 = cv2.rotate(img1, cv2.ROTATE_90_CLOCKWISE) 97 | elif p < 0.5: 98 | img0 = cv2.rotate(img0, cv2.ROTATE_180) 99 | gt = cv2.rotate(gt, cv2.ROTATE_180) 100 | img1 = cv2.rotate(img1, cv2.ROTATE_180) 101 | elif p < 0.75: 102 | img0 = cv2.rotate(img0, cv2.ROTATE_90_COUNTERCLOCKWISE) 103 | gt = cv2.rotate(gt, cv2.ROTATE_90_COUNTERCLOCKWISE) 104 | img1 = cv2.rotate(img1, cv2.ROTATE_90_COUNTERCLOCKWISE) 105 | img0 = torch.from_numpy(img0.copy()).permute(2, 0, 1) 106 | img1 = torch.from_numpy(img1.copy()).permute(2, 0, 1) 107 | gt = torch.from_numpy(gt.copy()).permute(2, 0, 1) 108 | timestep = torch.tensor(timestep).reshape(1, 1, 1) 109 | return torch.cat((img0, img1, gt), 0), timestep 110 | -------------------------------------------------------------------------------- /model/IFNet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from model.warplayer import warp 5 | from model.refine import * 6 | 7 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 8 | return nn.Sequential( 9 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1), 10 | nn.PReLU(out_planes) 11 | ) 12 | 13 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 14 | return nn.Sequential( 15 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 16 | padding=padding, dilation=dilation, bias=True), 17 | nn.PReLU(out_planes) 18 | ) 19 | 20 | class IFBlock(nn.Module): 21 | def __init__(self, in_planes, c=64): 22 | super(IFBlock, self).__init__() 23 | self.conv0 = nn.Sequential( 24 | conv(in_planes, c//2, 3, 2, 1), 25 | conv(c//2, c, 3, 2, 1), 26 | ) 27 | self.convblock = nn.Sequential( 28 | conv(c, c), 29 | conv(c, c), 30 | conv(c, c), 31 | conv(c, c), 32 | conv(c, c), 33 | conv(c, c), 34 | conv(c, c), 35 | conv(c, c), 36 | ) 37 | self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) 38 | 39 | def forward(self, x, flow, scale): 40 | if scale != 1: 41 | x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False) 42 | if flow != None: 43 | flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scale 44 | x = torch.cat((x, flow), 1) 45 | x = self.conv0(x) 46 | x = self.convblock(x) + x 47 | tmp = self.lastconv(x) 48 | tmp = F.interpolate(tmp, scale_factor = scale * 2, mode="bilinear", align_corners=False) 49 | flow = tmp[:, :4] * scale * 2 50 | mask = tmp[:, 4:5] 51 | return flow, mask 52 | 53 | class IFNet(nn.Module): 54 | def __init__(self): 55 | super(IFNet, self).__init__() 56 | self.block0 = IFBlock(6, c=240) 57 | self.block1 = IFBlock(13+4, c=150) 58 | self.block2 = IFBlock(13+4, c=90) 59 | self.block_tea = IFBlock(16+4, c=90) 60 | self.contextnet = Contextnet() 61 | self.unet = Unet() 62 | 63 | def forward(self, x, scale=[4,2,1], timestep=0.5): 64 | img0 = x[:, :3] 65 | img1 = x[:, 3:6] 66 | gt = x[:, 6:] # In inference time, gt is None 67 | flow_list = [] 68 | merged = [] 69 | mask_list = [] 70 | warped_img0 = img0 71 | warped_img1 = img1 72 | flow = None 73 | loss_distill = 0 74 | stu = [self.block0, self.block1, self.block2] 75 | for i in range(3): 76 | if flow != None: 77 | flow_d, mask_d = stu[i](torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i]) 78 | flow = flow + flow_d 79 | mask = mask + mask_d 80 | else: 81 | flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i]) 82 | mask_list.append(torch.sigmoid(mask)) 83 | flow_list.append(flow) 84 | warped_img0 = warp(img0, flow[:, :2]) 85 | warped_img1 = warp(img1, flow[:, 2:4]) 86 | merged_student = (warped_img0, warped_img1) 87 | merged.append(merged_student) 88 | if gt.shape[1] == 3: 89 | flow_d, mask_d = self.block_tea(torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1) 90 | flow_teacher = flow + flow_d 91 | warped_img0_teacher = warp(img0, flow_teacher[:, :2]) 92 | warped_img1_teacher = warp(img1, flow_teacher[:, 2:4]) 93 | mask_teacher = torch.sigmoid(mask + mask_d) 94 | merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher) 95 | else: 96 | flow_teacher = None 97 | merged_teacher = None 98 | for i in range(3): 99 | merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) 100 | if gt.shape[1] == 3: 101 | loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach() 102 | loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean() 103 | c0 = self.contextnet(img0, flow[:, :2]) 104 | c1 = self.contextnet(img1, flow[:, 2:4]) 105 | tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) 106 | res = tmp[:, :3] * 2 - 1 107 | merged[2] = torch.clamp(merged[2] + res, 0, 1) 108 | return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill 109 | -------------------------------------------------------------------------------- /model/IFNet_2R.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from model.warplayer import warp 5 | from model.refine_2R import * 6 | 7 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 8 | return nn.Sequential( 9 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1), 10 | nn.PReLU(out_planes) 11 | ) 12 | 13 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 14 | return nn.Sequential( 15 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 16 | padding=padding, dilation=dilation, bias=True), 17 | nn.PReLU(out_planes) 18 | ) 19 | 20 | class IFBlock(nn.Module): 21 | def __init__(self, in_planes, c=64): 22 | super(IFBlock, self).__init__() 23 | self.conv0 = nn.Sequential( 24 | conv(in_planes, c//2, 3, 1, 1), 25 | conv(c//2, c, 3, 2, 1), 26 | ) 27 | self.convblock = nn.Sequential( 28 | conv(c, c), 29 | conv(c, c), 30 | conv(c, c), 31 | conv(c, c), 32 | conv(c, c), 33 | conv(c, c), 34 | conv(c, c), 35 | conv(c, c), 36 | ) 37 | self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) 38 | 39 | def forward(self, x, flow, scale): 40 | if scale != 1: 41 | x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False) 42 | if flow != None: 43 | flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scale 44 | x = torch.cat((x, flow), 1) 45 | x = self.conv0(x) 46 | x = self.convblock(x) + x 47 | tmp = self.lastconv(x) 48 | tmp = F.interpolate(tmp, scale_factor = scale, mode="bilinear", align_corners=False) 49 | flow = tmp[:, :4] * scale 50 | mask = tmp[:, 4:5] 51 | return flow, mask 52 | 53 | class IFNet(nn.Module): 54 | def __init__(self): 55 | super(IFNet, self).__init__() 56 | self.block0 = IFBlock(6, c=240) 57 | self.block1 = IFBlock(13+4, c=150) 58 | self.block2 = IFBlock(13+4, c=90) 59 | self.block_tea = IFBlock(16+4, c=90) 60 | self.contextnet = Contextnet() 61 | self.unet = Unet() 62 | 63 | def forward(self, x, scale=[4,2,1], timestep=0.5): 64 | img0 = x[:, :3] 65 | img1 = x[:, 3:6] 66 | gt = x[:, 6:] # In inference time, gt is None 67 | flow_list = [] 68 | merged = [] 69 | mask_list = [] 70 | warped_img0 = img0 71 | warped_img1 = img1 72 | flow = None 73 | loss_distill = 0 74 | stu = [self.block0, self.block1, self.block2] 75 | for i in range(3): 76 | if flow != None: 77 | flow_d, mask_d = stu[i](torch.cat((img0, img1, warped_img0, warped_img1, mask), 1), flow, scale=scale[i]) 78 | flow = flow + flow_d 79 | mask = mask + mask_d 80 | else: 81 | flow, mask = stu[i](torch.cat((img0, img1), 1), None, scale=scale[i]) 82 | mask_list.append(torch.sigmoid(mask)) 83 | flow_list.append(flow) 84 | warped_img0 = warp(img0, flow[:, :2]) 85 | warped_img1 = warp(img1, flow[:, 2:4]) 86 | merged_student = (warped_img0, warped_img1) 87 | merged.append(merged_student) 88 | if gt.shape[1] == 3: 89 | flow_d, mask_d = self.block_tea(torch.cat((img0, img1, warped_img0, warped_img1, mask, gt), 1), flow, scale=1) 90 | flow_teacher = flow + flow_d 91 | warped_img0_teacher = warp(img0, flow_teacher[:, :2]) 92 | warped_img1_teacher = warp(img1, flow_teacher[:, 2:4]) 93 | mask_teacher = torch.sigmoid(mask + mask_d) 94 | merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher) 95 | else: 96 | flow_teacher = None 97 | merged_teacher = None 98 | for i in range(3): 99 | merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) 100 | if gt.shape[1] == 3: 101 | loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach() 102 | loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean() 103 | c0 = self.contextnet(img0, flow[:, :2]) 104 | c1 = self.contextnet(img1, flow[:, 2:4]) 105 | tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) 106 | res = tmp[:, :3] * 2 - 1 107 | merged[2] = torch.clamp(merged[2] + res, 0, 1) 108 | return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill 109 | -------------------------------------------------------------------------------- /model/oldmodel/IFNet_HD.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from model.warplayer import warp 6 | 7 | 8 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 9 | 10 | def conv_wo_act(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 11 | return nn.Sequential( 12 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 13 | padding=padding, dilation=dilation, bias=False), 14 | nn.BatchNorm2d(out_planes), 15 | ) 16 | 17 | 18 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 19 | return nn.Sequential( 20 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 21 | padding=padding, dilation=dilation, bias=False), 22 | nn.BatchNorm2d(out_planes), 23 | nn.PReLU(out_planes) 24 | ) 25 | 26 | 27 | class ResBlock(nn.Module): 28 | def __init__(self, in_planes, out_planes, stride=1): 29 | super(ResBlock, self).__init__() 30 | if in_planes == out_planes and stride == 1: 31 | self.conv0 = nn.Identity() 32 | else: 33 | self.conv0 = nn.Conv2d(in_planes, out_planes, 34 | 3, stride, 1, bias=False) 35 | self.conv1 = conv(in_planes, out_planes, 5, stride, 2) 36 | self.conv2 = conv_wo_act(out_planes, out_planes, 3, 1, 1) 37 | self.relu1 = nn.PReLU(1) 38 | self.relu2 = nn.PReLU(out_planes) 39 | self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False) 40 | self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False) 41 | 42 | def forward(self, x): 43 | y = self.conv0(x) 44 | x = self.conv1(x) 45 | x = self.conv2(x) 46 | w = x.mean(3, True).mean(2, True) 47 | w = self.relu1(self.fc1(w)) 48 | w = torch.sigmoid(self.fc2(w)) 49 | x = self.relu2(x * w + y) 50 | return x 51 | 52 | 53 | class IFBlock(nn.Module): 54 | def __init__(self, in_planes, scale=1, c=64): 55 | super(IFBlock, self).__init__() 56 | self.scale = scale 57 | self.conv0 = conv(in_planes, c, 5, 2, 2) 58 | self.res0 = ResBlock(c, c) 59 | self.res1 = ResBlock(c, c) 60 | self.res2 = ResBlock(c, c) 61 | self.res3 = ResBlock(c, c) 62 | self.res4 = ResBlock(c, c) 63 | self.res5 = ResBlock(c, c) 64 | self.conv1 = nn.Conv2d(c, 8, 3, 1, 1) 65 | self.up = nn.PixelShuffle(2) 66 | 67 | def forward(self, x): 68 | if self.scale != 1: 69 | x = F.interpolate(x, scale_factor=1. / self.scale, mode="bilinear", 70 | align_corners=False) 71 | x = self.conv0(x) 72 | x = self.res0(x) 73 | x = self.res1(x) 74 | x = self.res2(x) 75 | x = self.res3(x) 76 | x = self.res4(x) 77 | x = self.res5(x) 78 | x = self.conv1(x) 79 | flow = self.up(x) 80 | if self.scale != 1: 81 | flow = F.interpolate(flow, scale_factor=self.scale, mode="bilinear", 82 | align_corners=False) 83 | return flow 84 | 85 | 86 | class IFNet(nn.Module): 87 | def __init__(self): 88 | super(IFNet, self).__init__() 89 | self.block0 = IFBlock(6, scale=8, c=192) 90 | self.block1 = IFBlock(8, scale=4, c=128) 91 | self.block2 = IFBlock(8, scale=2, c=96) 92 | self.block3 = IFBlock(8, scale=1, c=48) 93 | 94 | def forward(self, x, scale=1.0): 95 | x = F.interpolate(x, scale_factor=0.5 * scale, mode="bilinear", 96 | align_corners=False) 97 | flow0 = self.block0(x) 98 | F1 = flow0 99 | warped_img0 = warp(x[:, :3], F1) 100 | warped_img1 = warp(x[:, 3:], -F1) 101 | flow1 = self.block1(torch.cat((warped_img0, warped_img1, F1), 1)) 102 | F2 = (flow0 + flow1) 103 | warped_img0 = warp(x[:, :3], F2) 104 | warped_img1 = warp(x[:, 3:], -F2) 105 | flow2 = self.block2(torch.cat((warped_img0, warped_img1, F2), 1)) 106 | F3 = (flow0 + flow1 + flow2) 107 | warped_img0 = warp(x[:, :3], F3) 108 | warped_img1 = warp(x[:, 3:], -F3) 109 | flow3 = self.block3(torch.cat((warped_img0, warped_img1, F3), 1)) 110 | F4 = (flow0 + flow1 + flow2 + flow3) 111 | F4 = F.interpolate(F4, scale_factor=1 / scale, mode="bilinear", 112 | align_corners=False) / scale 113 | return F4, [F1, F2, F3, F4] 114 | 115 | if __name__ == '__main__': 116 | img0 = torch.zeros(3, 3, 256, 256).float().to(device) 117 | img1 = torch.tensor(np.random.normal( 118 | 0, 1, (3, 3, 256, 256))).float().to(device) 119 | imgs = torch.cat((img0, img1), 1) 120 | flownet = IFNet() 121 | flow, _ = flownet(imgs) 122 | print(flow.shape) 123 | -------------------------------------------------------------------------------- /model/IFNet_m.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from model.warplayer import warp 5 | from model.refine import * 6 | 7 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 8 | return nn.Sequential( 9 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, kernel_size=4, stride=2, padding=1), 10 | nn.PReLU(out_planes) 11 | ) 12 | 13 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 14 | return nn.Sequential( 15 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 16 | padding=padding, dilation=dilation, bias=True), 17 | nn.PReLU(out_planes) 18 | ) 19 | 20 | class IFBlock(nn.Module): 21 | def __init__(self, in_planes, c=64): 22 | super(IFBlock, self).__init__() 23 | self.conv0 = nn.Sequential( 24 | conv(in_planes, c//2, 3, 2, 1), 25 | conv(c//2, c, 3, 2, 1), 26 | ) 27 | self.convblock = nn.Sequential( 28 | conv(c, c), 29 | conv(c, c), 30 | conv(c, c), 31 | conv(c, c), 32 | conv(c, c), 33 | conv(c, c), 34 | conv(c, c), 35 | conv(c, c), 36 | ) 37 | self.lastconv = nn.ConvTranspose2d(c, 5, 4, 2, 1) 38 | 39 | def forward(self, x, flow, scale): 40 | if scale != 1: 41 | x = F.interpolate(x, scale_factor = 1. / scale, mode="bilinear", align_corners=False) 42 | if flow != None: 43 | flow = F.interpolate(flow, scale_factor = 1. / scale, mode="bilinear", align_corners=False) * 1. / scale 44 | x = torch.cat((x, flow), 1) 45 | x = self.conv0(x) 46 | x = self.convblock(x) + x 47 | tmp = self.lastconv(x) 48 | tmp = F.interpolate(tmp, scale_factor = scale * 2, mode="bilinear", align_corners=False) 49 | flow = tmp[:, :4] * scale * 2 50 | mask = tmp[:, 4:5] 51 | return flow, mask 52 | 53 | class IFNet_m(nn.Module): 54 | def __init__(self): 55 | super(IFNet_m, self).__init__() 56 | self.block0 = IFBlock(6+1, c=240) 57 | self.block1 = IFBlock(13+4+1, c=150) 58 | self.block2 = IFBlock(13+4+1, c=90) 59 | self.block_tea = IFBlock(16+4+1, c=90) 60 | self.contextnet = Contextnet() 61 | self.unet = Unet() 62 | 63 | def forward(self, x, scale=[4,2,1], timestep=0.5, returnflow=False): 64 | timestep = (x[:, :1].clone() * 0 + 1) * timestep 65 | img0 = x[:, :3] 66 | img1 = x[:, 3:6] 67 | gt = x[:, 6:] # In inference time, gt is None 68 | flow_list = [] 69 | merged = [] 70 | mask_list = [] 71 | warped_img0 = img0 72 | warped_img1 = img1 73 | flow = None 74 | loss_distill = 0 75 | stu = [self.block0, self.block1, self.block2] 76 | for i in range(3): 77 | if flow != None: 78 | flow_d, mask_d = stu[i](torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask), 1), flow, scale=scale[i]) 79 | flow = flow + flow_d 80 | mask = mask + mask_d 81 | else: 82 | flow, mask = stu[i](torch.cat((img0, img1, timestep), 1), None, scale=scale[i]) 83 | mask_list.append(torch.sigmoid(mask)) 84 | flow_list.append(flow) 85 | warped_img0 = warp(img0, flow[:, :2]) 86 | warped_img1 = warp(img1, flow[:, 2:4]) 87 | merged_student = (warped_img0, warped_img1) 88 | merged.append(merged_student) 89 | if gt.shape[1] == 3: 90 | flow_d, mask_d = self.block_tea(torch.cat((img0, img1, timestep, warped_img0, warped_img1, mask, gt), 1), flow, scale=1) 91 | flow_teacher = flow + flow_d 92 | warped_img0_teacher = warp(img0, flow_teacher[:, :2]) 93 | warped_img1_teacher = warp(img1, flow_teacher[:, 2:4]) 94 | mask_teacher = torch.sigmoid(mask + mask_d) 95 | merged_teacher = warped_img0_teacher * mask_teacher + warped_img1_teacher * (1 - mask_teacher) 96 | else: 97 | flow_teacher = None 98 | merged_teacher = None 99 | for i in range(3): 100 | merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) 101 | if gt.shape[1] == 3: 102 | loss_mask = ((merged[i] - gt).abs().mean(1, True) > (merged_teacher - gt).abs().mean(1, True) + 0.01).float().detach() 103 | loss_distill += (((flow_teacher.detach() - flow_list[i]) ** 2).mean(1, True) ** 0.5 * loss_mask).mean() 104 | if returnflow: 105 | return flow 106 | else: 107 | c0 = self.contextnet(img0, flow[:, :2]) 108 | c1 = self.contextnet(img1, flow[:, 2:4]) 109 | tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) 110 | res = tmp[:, :3] * 2 - 1 111 | merged[2] = torch.clamp(merged[2] + res, 0, 1) 112 | return flow_list, mask_list[2], merged, flow_teacher, merged_teacher, loss_distill 113 | -------------------------------------------------------------------------------- /model/loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | import torchvision.models as models 6 | 7 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 8 | 9 | 10 | class EPE(nn.Module): 11 | def __init__(self): 12 | super(EPE, self).__init__() 13 | 14 | def forward(self, flow, gt, loss_mask): 15 | loss_map = (flow - gt.detach()) ** 2 16 | loss_map = (loss_map.sum(1, True) + 1e-6) ** 0.5 17 | return (loss_map * loss_mask) 18 | 19 | 20 | class Ternary(nn.Module): 21 | def __init__(self): 22 | super(Ternary, self).__init__() 23 | patch_size = 7 24 | out_channels = patch_size * patch_size 25 | self.w = np.eye(out_channels).reshape( 26 | (patch_size, patch_size, 1, out_channels)) 27 | self.w = np.transpose(self.w, (3, 2, 0, 1)) 28 | self.w = torch.tensor(self.w).float().to(device) 29 | 30 | def transform(self, img): 31 | patches = F.conv2d(img, self.w, padding=3, bias=None) 32 | transf = patches - img 33 | transf_norm = transf / torch.sqrt(0.81 + transf**2) 34 | return transf_norm 35 | 36 | def rgb2gray(self, rgb): 37 | r, g, b = rgb[:, 0:1, :, :], rgb[:, 1:2, :, :], rgb[:, 2:3, :, :] 38 | gray = 0.2989 * r + 0.5870 * g + 0.1140 * b 39 | return gray 40 | 41 | def hamming(self, t1, t2): 42 | dist = (t1 - t2) ** 2 43 | dist_norm = torch.mean(dist / (0.1 + dist), 1, True) 44 | return dist_norm 45 | 46 | def valid_mask(self, t, padding): 47 | n, _, h, w = t.size() 48 | inner = torch.ones(n, 1, h - 2 * padding, w - 2 * padding).type_as(t) 49 | mask = F.pad(inner, [padding] * 4) 50 | return mask 51 | 52 | def forward(self, img0, img1): 53 | img0 = self.transform(self.rgb2gray(img0)) 54 | img1 = self.transform(self.rgb2gray(img1)) 55 | return self.hamming(img0, img1) * self.valid_mask(img0, 1) 56 | 57 | 58 | class SOBEL(nn.Module): 59 | def __init__(self): 60 | super(SOBEL, self).__init__() 61 | self.kernelX = torch.tensor([ 62 | [1, 0, -1], 63 | [2, 0, -2], 64 | [1, 0, -1], 65 | ]).float() 66 | self.kernelY = self.kernelX.clone().T 67 | self.kernelX = self.kernelX.unsqueeze(0).unsqueeze(0).to(device) 68 | self.kernelY = self.kernelY.unsqueeze(0).unsqueeze(0).to(device) 69 | 70 | def forward(self, pred, gt): 71 | N, C, H, W = pred.shape[0], pred.shape[1], pred.shape[2], pred.shape[3] 72 | img_stack = torch.cat( 73 | [pred.reshape(N*C, 1, H, W), gt.reshape(N*C, 1, H, W)], 0) 74 | sobel_stack_x = F.conv2d(img_stack, self.kernelX, padding=1) 75 | sobel_stack_y = F.conv2d(img_stack, self.kernelY, padding=1) 76 | pred_X, gt_X = sobel_stack_x[:N*C], sobel_stack_x[N*C:] 77 | pred_Y, gt_Y = sobel_stack_y[:N*C], sobel_stack_y[N*C:] 78 | 79 | L1X, L1Y = torch.abs(pred_X-gt_X), torch.abs(pred_Y-gt_Y) 80 | loss = (L1X+L1Y) 81 | return loss 82 | 83 | class MeanShift(nn.Conv2d): 84 | def __init__(self, data_mean, data_std, data_range=1, norm=True): 85 | c = len(data_mean) 86 | super(MeanShift, self).__init__(c, c, kernel_size=1) 87 | std = torch.Tensor(data_std) 88 | self.weight.data = torch.eye(c).view(c, c, 1, 1) 89 | if norm: 90 | self.weight.data.div_(std.view(c, 1, 1, 1)) 91 | self.bias.data = -1 * data_range * torch.Tensor(data_mean) 92 | self.bias.data.div_(std) 93 | else: 94 | self.weight.data.mul_(std.view(c, 1, 1, 1)) 95 | self.bias.data = data_range * torch.Tensor(data_mean) 96 | self.requires_grad = False 97 | 98 | class VGGPerceptualLoss(torch.nn.Module): 99 | def __init__(self, rank=0): 100 | super(VGGPerceptualLoss, self).__init__() 101 | blocks = [] 102 | pretrained = True 103 | self.vgg_pretrained_features = models.vgg19(pretrained=pretrained).features 104 | self.normalize = MeanShift([0.485, 0.456, 0.406], [0.229, 0.224, 0.225], norm=True).cuda() 105 | for param in self.parameters(): 106 | param.requires_grad = False 107 | 108 | def forward(self, X, Y, indices=None): 109 | X = self.normalize(X) 110 | Y = self.normalize(Y) 111 | indices = [2, 7, 12, 21, 30] 112 | weights = [1.0/2.6, 1.0/4.8, 1.0/3.7, 1.0/5.6, 10/1.5] 113 | k = 0 114 | loss = 0 115 | for i in range(indices[-1]): 116 | X = self.vgg_pretrained_features[i](X) 117 | Y = self.vgg_pretrained_features[i](Y) 118 | if (i+1) in indices: 119 | loss += weights[k] * (X - Y.detach()).abs().mean() * 0.1 120 | k += 1 121 | return loss 122 | 123 | if __name__ == '__main__': 124 | img0 = torch.zeros(3, 3, 256, 256).float().to(device) 125 | img1 = torch.tensor(np.random.normal( 126 | 0, 1, (3, 3, 256, 256))).float().to(device) 127 | ternary_loss = Ternary() 128 | print(ternary_loss(img0, img1).shape) 129 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import math 4 | import time 5 | import torch 6 | import torch.distributed as dist 7 | import numpy as np 8 | import random 9 | import argparse 10 | 11 | from model.RIFE import Model 12 | from dataset import * 13 | from torch.utils.data import DataLoader, Dataset 14 | from torch.utils.tensorboard import SummaryWriter 15 | from torch.utils.data.distributed import DistributedSampler 16 | 17 | device = torch.device("cuda") 18 | 19 | log_path = 'train_log' 20 | 21 | def get_learning_rate(step): 22 | if step < 2000: 23 | mul = step / 2000. 24 | return 3e-4 * mul 25 | else: 26 | mul = np.cos((step - 2000) / (args.epoch * args.step_per_epoch - 2000.) * math.pi) * 0.5 + 0.5 27 | return (3e-4 - 3e-6) * mul + 3e-6 28 | 29 | def flow2rgb(flow_map_np): 30 | h, w, _ = flow_map_np.shape 31 | rgb_map = np.ones((h, w, 3)).astype(np.float32) 32 | normalized_flow_map = flow_map_np / (np.abs(flow_map_np).max()) 33 | 34 | rgb_map[:, :, 0] += normalized_flow_map[:, :, 0] 35 | rgb_map[:, :, 1] -= 0.5 * (normalized_flow_map[:, :, 0] + normalized_flow_map[:, :, 1]) 36 | rgb_map[:, :, 2] += normalized_flow_map[:, :, 1] 37 | return rgb_map.clip(0, 1) 38 | 39 | def train(model, local_rank): 40 | if local_rank == 0: 41 | writer = SummaryWriter('train') 42 | writer_val = SummaryWriter('validate') 43 | else: 44 | writer = None 45 | writer_val = None 46 | step = 0 47 | nr_eval = 0 48 | dataset = VimeoDataset('train') 49 | sampler = DistributedSampler(dataset) 50 | train_data = DataLoader(dataset, batch_size=args.batch_size, num_workers=8, pin_memory=True, drop_last=True, sampler=sampler) 51 | args.step_per_epoch = train_data.__len__() 52 | dataset_val = VimeoDataset('validation') 53 | val_data = DataLoader(dataset_val, batch_size=16, pin_memory=True, num_workers=8) 54 | print('training...') 55 | time_stamp = time.time() 56 | for epoch in range(args.epoch): 57 | sampler.set_epoch(epoch) 58 | for i, data in enumerate(train_data): 59 | data_time_interval = time.time() - time_stamp 60 | time_stamp = time.time() 61 | data_gpu, timestep = data 62 | data_gpu = data_gpu.to(device, non_blocking=True) / 255. 63 | timestep = timestep.to(device, non_blocking=True) 64 | imgs = data_gpu[:, :6] 65 | gt = data_gpu[:, 6:9] 66 | learning_rate = get_learning_rate(step) * args.world_size / 4 67 | pred, info = model.update(imgs, gt, learning_rate, training=True) # pass timestep if you are training RIFEm 68 | train_time_interval = time.time() - time_stamp 69 | time_stamp = time.time() 70 | if step % 200 == 1 and local_rank == 0: 71 | writer.add_scalar('learning_rate', learning_rate, step) 72 | writer.add_scalar('loss/l1', info['loss_l1'], step) 73 | writer.add_scalar('loss/tea', info['loss_tea'], step) 74 | writer.add_scalar('loss/distill', info['loss_distill'], step) 75 | if step % 1000 == 1 and local_rank == 0: 76 | gt = (gt.permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8') 77 | mask = (torch.cat((info['mask'], info['mask_tea']), 3).permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8') 78 | pred = (pred.permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8') 79 | merged_img = (info['merged_tea'].permute(0, 2, 3, 1).detach().cpu().numpy() * 255).astype('uint8') 80 | flow0 = info['flow'].permute(0, 2, 3, 1).detach().cpu().numpy() 81 | flow1 = info['flow_tea'].permute(0, 2, 3, 1).detach().cpu().numpy() 82 | for i in range(5): 83 | imgs = np.concatenate((merged_img[i], pred[i], gt[i]), 1)[:, :, ::-1] 84 | writer.add_image(str(i) + '/img', imgs, step, dataformats='HWC') 85 | writer.add_image(str(i) + '/flow', np.concatenate((flow2rgb(flow0[i]), flow2rgb(flow1[i])), 1), step, dataformats='HWC') 86 | writer.add_image(str(i) + '/mask', mask[i], step, dataformats='HWC') 87 | writer.flush() 88 | if local_rank == 0: 89 | print('epoch:{} {}/{} time:{:.2f}+{:.2f} loss_l1:{:.4e}'.format(epoch, i, args.step_per_epoch, data_time_interval, train_time_interval, info['loss_l1'])) 90 | step += 1 91 | nr_eval += 1 92 | if nr_eval % 5 == 0: 93 | evaluate(model, val_data, step, local_rank, writer_val) 94 | model.save_model(log_path, local_rank) 95 | dist.barrier() 96 | 97 | def evaluate(model, val_data, nr_eval, local_rank, writer_val): 98 | loss_l1_list = [] 99 | loss_distill_list = [] 100 | loss_tea_list = [] 101 | psnr_list = [] 102 | psnr_list_teacher = [] 103 | time_stamp = time.time() 104 | for i, data in enumerate(val_data): 105 | data_gpu, timestep = data 106 | data_gpu = data_gpu.to(device, non_blocking=True) / 255. 107 | imgs = data_gpu[:, :6] 108 | gt = data_gpu[:, 6:9] 109 | with torch.no_grad(): 110 | pred, info = model.update(imgs, gt, training=False) 111 | merged_img = info['merged_tea'] 112 | loss_l1_list.append(info['loss_l1'].cpu().numpy()) 113 | loss_tea_list.append(info['loss_tea'].cpu().numpy()) 114 | loss_distill_list.append(info['loss_distill'].cpu().numpy()) 115 | for j in range(gt.shape[0]): 116 | psnr = -10 * math.log10(torch.mean((gt[j] - pred[j]) * (gt[j] - pred[j])).cpu().data) 117 | psnr_list.append(psnr) 118 | psnr = -10 * math.log10(torch.mean((merged_img[j] - gt[j]) * (merged_img[j] - gt[j])).cpu().data) 119 | psnr_list_teacher.append(psnr) 120 | gt = (gt.permute(0, 2, 3, 1).cpu().numpy() * 255).astype('uint8') 121 | pred = (pred.permute(0, 2, 3, 1).cpu().numpy() * 255).astype('uint8') 122 | merged_img = (merged_img.permute(0, 2, 3, 1).cpu().numpy() * 255).astype('uint8') 123 | flow0 = info['flow'].permute(0, 2, 3, 1).cpu().numpy() 124 | flow1 = info['flow_tea'].permute(0, 2, 3, 1).cpu().numpy() 125 | if i == 0 and local_rank == 0: 126 | for j in range(10): 127 | imgs = np.concatenate((merged_img[j], pred[j], gt[j]), 1)[:, :, ::-1] 128 | writer_val.add_image(str(j) + '/img', imgs.copy(), nr_eval, dataformats='HWC') 129 | writer_val.add_image(str(j) + '/flow', flow2rgb(flow0[j][:, :, ::-1]), nr_eval, dataformats='HWC') 130 | 131 | eval_time_interval = time.time() - time_stamp 132 | 133 | if local_rank != 0: 134 | return 135 | writer_val.add_scalar('psnr', np.array(psnr_list).mean(), nr_eval) 136 | writer_val.add_scalar('psnr_teacher', np.array(psnr_list_teacher).mean(), nr_eval) 137 | 138 | if __name__ == "__main__": 139 | parser = argparse.ArgumentParser() 140 | parser.add_argument('--epoch', default=300, type=int) 141 | parser.add_argument('--batch_size', default=16, type=int, help='minibatch size') 142 | parser.add_argument('--local_rank', default=0, type=int, help='local rank') 143 | parser.add_argument('--world_size', default=4, type=int, help='world size') 144 | args = parser.parse_args() 145 | torch.distributed.init_process_group(backend="nccl", world_size=args.world_size) 146 | torch.cuda.set_device(args.local_rank) 147 | seed = 1234 148 | random.seed(seed) 149 | np.random.seed(seed) 150 | torch.manual_seed(seed) 151 | torch.cuda.manual_seed_all(seed) 152 | torch.backends.cudnn.benchmark = True 153 | model = Model(args.local_rank) 154 | train(model, args.local_rank) 155 | 156 | -------------------------------------------------------------------------------- /model/pytorch_msssim/__init__.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn.functional as F 3 | from math import exp 4 | import numpy as np 5 | 6 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 7 | 8 | def gaussian(window_size, sigma): 9 | gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)]) 10 | return gauss/gauss.sum() 11 | 12 | 13 | def create_window(window_size, channel=1): 14 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 15 | _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0).to(device) 16 | window = _2D_window.expand(channel, 1, window_size, window_size).contiguous() 17 | return window 18 | 19 | def create_window_3d(window_size, channel=1): 20 | _1D_window = gaussian(window_size, 1.5).unsqueeze(1) 21 | _2D_window = _1D_window.mm(_1D_window.t()) 22 | _3D_window = _2D_window.unsqueeze(2) @ (_1D_window.t()) 23 | window = _3D_window.expand(1, channel, window_size, window_size, window_size).contiguous().to(device) 24 | return window 25 | 26 | 27 | def ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): 28 | # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). 29 | if val_range is None: 30 | if torch.max(img1) > 128: 31 | max_val = 255 32 | else: 33 | max_val = 1 34 | 35 | if torch.min(img1) < -0.5: 36 | min_val = -1 37 | else: 38 | min_val = 0 39 | L = max_val - min_val 40 | else: 41 | L = val_range 42 | 43 | padd = 0 44 | (_, channel, height, width) = img1.size() 45 | if window is None: 46 | real_size = min(window_size, height, width) 47 | window = create_window(real_size, channel=channel).to(img1.device) 48 | 49 | # mu1 = F.conv2d(img1, window, padding=padd, groups=channel) 50 | # mu2 = F.conv2d(img2, window, padding=padd, groups=channel) 51 | mu1 = F.conv2d(F.pad(img1, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel) 52 | mu2 = F.conv2d(F.pad(img2, (5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=channel) 53 | 54 | mu1_sq = mu1.pow(2) 55 | mu2_sq = mu2.pow(2) 56 | mu1_mu2 = mu1 * mu2 57 | 58 | sigma1_sq = F.conv2d(F.pad(img1 * img1, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_sq 59 | sigma2_sq = F.conv2d(F.pad(img2 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu2_sq 60 | sigma12 = F.conv2d(F.pad(img1 * img2, (5, 5, 5, 5), 'replicate'), window, padding=padd, groups=channel) - mu1_mu2 61 | 62 | C1 = (0.01 * L) ** 2 63 | C2 = (0.03 * L) ** 2 64 | 65 | v1 = 2.0 * sigma12 + C2 66 | v2 = sigma1_sq + sigma2_sq + C2 67 | cs = torch.mean(v1 / v2) # contrast sensitivity 68 | 69 | ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) 70 | 71 | if size_average: 72 | ret = ssim_map.mean() 73 | else: 74 | ret = ssim_map.mean(1).mean(1).mean(1) 75 | 76 | if full: 77 | return ret, cs 78 | return ret 79 | 80 | 81 | def ssim_matlab(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None): 82 | # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh). 83 | if val_range is None: 84 | if torch.max(img1) > 128: 85 | max_val = 255 86 | else: 87 | max_val = 1 88 | 89 | if torch.min(img1) < -0.5: 90 | min_val = -1 91 | else: 92 | min_val = 0 93 | L = max_val - min_val 94 | else: 95 | L = val_range 96 | 97 | padd = 0 98 | (_, _, height, width) = img1.size() 99 | if window is None: 100 | real_size = min(window_size, height, width) 101 | window = create_window_3d(real_size, channel=1).to(img1.device) 102 | # Channel is set to 1 since we consider color images as volumetric images 103 | 104 | img1 = img1.unsqueeze(1) 105 | img2 = img2.unsqueeze(1) 106 | 107 | mu1 = F.conv3d(F.pad(img1, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1) 108 | mu2 = F.conv3d(F.pad(img2, (5, 5, 5, 5, 5, 5), mode='replicate'), window, padding=padd, groups=1) 109 | 110 | mu1_sq = mu1.pow(2) 111 | mu2_sq = mu2.pow(2) 112 | mu1_mu2 = mu1 * mu2 113 | 114 | sigma1_sq = F.conv3d(F.pad(img1 * img1, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_sq 115 | sigma2_sq = F.conv3d(F.pad(img2 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu2_sq 116 | sigma12 = F.conv3d(F.pad(img1 * img2, (5, 5, 5, 5, 5, 5), 'replicate'), window, padding=padd, groups=1) - mu1_mu2 117 | 118 | C1 = (0.01 * L) ** 2 119 | C2 = (0.03 * L) ** 2 120 | 121 | v1 = 2.0 * sigma12 + C2 122 | v2 = sigma1_sq + sigma2_sq + C2 123 | cs = torch.mean(v1 / v2) # contrast sensitivity 124 | 125 | ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2) 126 | 127 | if size_average: 128 | ret = ssim_map.mean() 129 | else: 130 | ret = ssim_map.mean(1).mean(1).mean(1) 131 | 132 | if full: 133 | return ret, cs 134 | return ret 135 | 136 | 137 | def msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=False): 138 | device = img1.device 139 | weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device) 140 | levels = weights.size()[0] 141 | mssim = [] 142 | mcs = [] 143 | for _ in range(levels): 144 | sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range) 145 | mssim.append(sim) 146 | mcs.append(cs) 147 | 148 | img1 = F.avg_pool2d(img1, (2, 2)) 149 | img2 = F.avg_pool2d(img2, (2, 2)) 150 | 151 | mssim = torch.stack(mssim) 152 | mcs = torch.stack(mcs) 153 | 154 | # Normalize (to avoid NaNs during training unstable models, not compliant with original definition) 155 | if normalize: 156 | mssim = (mssim + 1) / 2 157 | mcs = (mcs + 1) / 2 158 | 159 | pow1 = mcs ** weights 160 | pow2 = mssim ** weights 161 | # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/ 162 | output = torch.prod(pow1[:-1] * pow2[-1]) 163 | return output 164 | 165 | 166 | # Classes to re-use window 167 | class SSIM(torch.nn.Module): 168 | def __init__(self, window_size=11, size_average=True, val_range=None): 169 | super(SSIM, self).__init__() 170 | self.window_size = window_size 171 | self.size_average = size_average 172 | self.val_range = val_range 173 | 174 | # Assume 3 channel for SSIM 175 | self.channel = 3 176 | self.window = create_window(window_size, channel=self.channel) 177 | 178 | def forward(self, img1, img2): 179 | (_, channel, _, _) = img1.size() 180 | 181 | if channel == self.channel and self.window.dtype == img1.dtype: 182 | window = self.window 183 | else: 184 | window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype) 185 | self.window = window 186 | self.channel = channel 187 | 188 | _ssim = ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average) 189 | dssim = (1 - _ssim) / 2 190 | return dssim 191 | 192 | class MSSSIM(torch.nn.Module): 193 | def __init__(self, window_size=11, size_average=True, channel=3): 194 | super(MSSSIM, self).__init__() 195 | self.window_size = window_size 196 | self.size_average = size_average 197 | self.channel = channel 198 | 199 | def forward(self, img1, img2): 200 | return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average) 201 | -------------------------------------------------------------------------------- /model/oldmodel/RIFE_HDv2.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from torch.optim import AdamW 5 | import torch.optim as optim 6 | import itertools 7 | from model.warplayer import warp 8 | from torch.nn.parallel import DistributedDataParallel as DDP 9 | from model.oldmodel.IFNet_HDv2 import * 10 | import torch.nn.functional as F 11 | from model.loss import * 12 | 13 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 14 | 15 | 16 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 17 | return nn.Sequential( 18 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 19 | padding=padding, dilation=dilation, bias=True), 20 | nn.PReLU(out_planes) 21 | ) 22 | 23 | 24 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 25 | return nn.Sequential( 26 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, 27 | kernel_size=4, stride=2, padding=1, bias=True), 28 | nn.PReLU(out_planes) 29 | ) 30 | 31 | def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 32 | return nn.Sequential( 33 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 34 | padding=padding, dilation=dilation, bias=True), 35 | ) 36 | 37 | class Conv2(nn.Module): 38 | def __init__(self, in_planes, out_planes, stride=2): 39 | super(Conv2, self).__init__() 40 | self.conv1 = conv(in_planes, out_planes, 3, stride, 1) 41 | self.conv2 = conv(out_planes, out_planes, 3, 1, 1) 42 | 43 | def forward(self, x): 44 | x = self.conv1(x) 45 | x = self.conv2(x) 46 | return x 47 | 48 | c = 32 49 | 50 | class ContextNet(nn.Module): 51 | def __init__(self): 52 | super(ContextNet, self).__init__() 53 | self.conv0 = Conv2(3, c) 54 | self.conv1 = Conv2(c, c) 55 | self.conv2 = Conv2(c, 2*c) 56 | self.conv3 = Conv2(2*c, 4*c) 57 | self.conv4 = Conv2(4*c, 8*c) 58 | 59 | def forward(self, x, flow): 60 | x = self.conv0(x) 61 | x = self.conv1(x) 62 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 63 | f1 = warp(x, flow) 64 | x = self.conv2(x) 65 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", 66 | align_corners=False) * 0.5 67 | f2 = warp(x, flow) 68 | x = self.conv3(x) 69 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", 70 | align_corners=False) * 0.5 71 | f3 = warp(x, flow) 72 | x = self.conv4(x) 73 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", 74 | align_corners=False) * 0.5 75 | f4 = warp(x, flow) 76 | return [f1, f2, f3, f4] 77 | 78 | 79 | class FusionNet(nn.Module): 80 | def __init__(self): 81 | super(FusionNet, self).__init__() 82 | self.conv0 = Conv2(10, c) 83 | self.down0 = Conv2(c, 2*c) 84 | self.down1 = Conv2(4*c, 4*c) 85 | self.down2 = Conv2(8*c, 8*c) 86 | self.down3 = Conv2(16*c, 16*c) 87 | self.up0 = deconv(32*c, 8*c) 88 | self.up1 = deconv(16*c, 4*c) 89 | self.up2 = deconv(8*c, 2*c) 90 | self.up3 = deconv(4*c, c) 91 | self.conv = nn.ConvTranspose2d(c, 4, 4, 2, 1) 92 | 93 | def forward(self, img0, img1, flow, c0, c1, flow_gt): 94 | warped_img0 = warp(img0, flow[:, :2]) 95 | warped_img1 = warp(img1, flow[:, 2:4]) 96 | if flow_gt == None: 97 | warped_img0_gt, warped_img1_gt = None, None 98 | else: 99 | warped_img0_gt = warp(img0, flow_gt[:, :2]) 100 | warped_img1_gt = warp(img1, flow_gt[:, 2:4]) 101 | x = self.conv0(torch.cat((warped_img0, warped_img1, flow), 1)) 102 | s0 = self.down0(x) 103 | s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) 104 | s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) 105 | s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) 106 | x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) 107 | x = self.up1(torch.cat((x, s2), 1)) 108 | x = self.up2(torch.cat((x, s1), 1)) 109 | x = self.up3(torch.cat((x, s0), 1)) 110 | x = self.conv(x) 111 | return x, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt 112 | 113 | 114 | class Model: 115 | def __init__(self, local_rank=-1): 116 | self.flownet = IFNet() 117 | self.contextnet = ContextNet() 118 | self.fusionnet = FusionNet() 119 | self.device() 120 | self.optimG = AdamW(itertools.chain( 121 | self.flownet.parameters(), 122 | self.contextnet.parameters(), 123 | self.fusionnet.parameters()), lr=1e-6, weight_decay=1e-4) 124 | self.schedulerG = optim.lr_scheduler.CyclicLR( 125 | self.optimG, base_lr=1e-6, max_lr=1e-3, step_size_up=8000, cycle_momentum=False) 126 | self.epe = EPE() 127 | self.ter = Ternary() 128 | self.sobel = SOBEL() 129 | if local_rank != -1: 130 | self.flownet = DDP(self.flownet, device_ids=[ 131 | local_rank], output_device=local_rank) 132 | self.contextnet = DDP(self.contextnet, device_ids=[ 133 | local_rank], output_device=local_rank) 134 | self.fusionnet = DDP(self.fusionnet, device_ids=[ 135 | local_rank], output_device=local_rank) 136 | 137 | def train(self): 138 | self.flownet.train() 139 | self.contextnet.train() 140 | self.fusionnet.train() 141 | 142 | def eval(self): 143 | self.flownet.eval() 144 | self.contextnet.eval() 145 | self.fusionnet.eval() 146 | 147 | def device(self): 148 | self.flownet.to(device) 149 | self.contextnet.to(device) 150 | self.fusionnet.to(device) 151 | 152 | def load_model(self, path, rank): 153 | def convert(param): 154 | if rank == -1: 155 | return { 156 | k.replace("module.", ""): v 157 | for k, v in param.items() 158 | if "module." in k 159 | } 160 | else: 161 | return param 162 | if rank <= 0: 163 | self.flownet.load_state_dict( 164 | convert(torch.load('{}/flownet.pkl'.format(path), map_location=device))) 165 | self.contextnet.load_state_dict( 166 | convert(torch.load('{}/contextnet.pkl'.format(path), map_location=device))) 167 | self.fusionnet.load_state_dict( 168 | convert(torch.load('{}/unet.pkl'.format(path), map_location=device))) 169 | 170 | def save_model(self, path, rank): 171 | if rank == 0: 172 | torch.save(self.flownet.state_dict(), '{}/flownet.pkl'.format(path)) 173 | torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path)) 174 | torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path)) 175 | 176 | def predict(self, imgs, flow, training=True, flow_gt=None): 177 | img0 = imgs[:, :3] 178 | img1 = imgs[:, 3:] 179 | c0 = self.contextnet(img0, flow[:, :2]) 180 | c1 = self.contextnet(img1, flow[:, 2:4]) 181 | flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", 182 | align_corners=False) * 2.0 183 | refine_output, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.fusionnet( 184 | img0, img1, flow, c0, c1, flow_gt) 185 | res = torch.sigmoid(refine_output[:, :3]) * 2 - 1 186 | mask = torch.sigmoid(refine_output[:, 3:4]) 187 | merged_img = warped_img0 * mask + warped_img1 * (1 - mask) 188 | pred = merged_img + res 189 | pred = torch.clamp(pred, 0, 1) 190 | if training: 191 | return pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt 192 | else: 193 | return pred 194 | 195 | def inference(self, img0, img1, scale=1.0): 196 | imgs = torch.cat((img0, img1), 1) 197 | flow, _ = self.flownet(imgs, scale) 198 | return self.predict(imgs, flow, training=False) 199 | 200 | def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): 201 | for param_group in self.optimG.param_groups: 202 | param_group['lr'] = learning_rate 203 | if training: 204 | self.train() 205 | else: 206 | self.eval() 207 | flow, flow_list = self.flownet(imgs) 208 | pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.predict( 209 | imgs, flow, flow_gt=flow_gt) 210 | loss_ter = self.ter(pred, gt).mean() 211 | if training: 212 | with torch.no_grad(): 213 | loss_flow = torch.abs(warped_img0_gt - gt).mean() 214 | loss_mask = torch.abs( 215 | merged_img - gt).sum(1, True).float().detach() 216 | loss_mask = F.interpolate(loss_mask, scale_factor=0.5, mode="bilinear", 217 | align_corners=False).detach() 218 | flow_gt = (F.interpolate(flow_gt, scale_factor=0.5, mode="bilinear", 219 | align_corners=False) * 0.5).detach() 220 | loss_cons = 0 221 | for i in range(4): 222 | loss_cons += self.epe(flow_list[i][:, :2], flow_gt[:, :2], 1) 223 | loss_cons += self.epe(flow_list[i][:, 2:4], flow_gt[:, 2:4], 1) 224 | loss_cons = loss_cons.mean() * 0.01 225 | else: 226 | loss_cons = torch.tensor([0]) 227 | loss_flow = torch.abs(warped_img0 - gt).mean() 228 | loss_mask = 1 229 | loss_l1 = (((pred - gt) ** 2 + 1e-6) ** 0.5).mean() 230 | if training: 231 | self.optimG.zero_grad() 232 | loss_G = loss_l1 + loss_cons + loss_ter 233 | loss_G.backward() 234 | self.optimG.step() 235 | return pred, merged_img, flow, loss_l1, loss_flow, loss_cons, loss_ter, loss_mask 236 | 237 | 238 | if __name__ == '__main__': 239 | img0 = torch.zeros(3, 3, 256, 256).float().to(device) 240 | img1 = torch.tensor(np.random.normal( 241 | 0, 1, (3, 3, 256, 256))).float().to(device) 242 | imgs = torch.cat((img0, img1), 1) 243 | model = Model() 244 | model.eval() 245 | print(model.inference(imgs).shape) 246 | -------------------------------------------------------------------------------- /model/oldmodel/RIFE_HD.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import numpy as np 4 | from torch.optim import AdamW 5 | import torch.optim as optim 6 | import itertools 7 | from model.warplayer import warp 8 | from torch.nn.parallel import DistributedDataParallel as DDP 9 | from model.oldmodel.IFNet_HD import * 10 | import torch.nn.functional as F 11 | from model.loss import * 12 | 13 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 14 | 15 | 16 | def conv(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 17 | return nn.Sequential( 18 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 19 | padding=padding, dilation=dilation, bias=True), 20 | nn.PReLU(out_planes) 21 | ) 22 | 23 | 24 | def deconv(in_planes, out_planes, kernel_size=4, stride=2, padding=1): 25 | return nn.Sequential( 26 | torch.nn.ConvTranspose2d(in_channels=in_planes, out_channels=out_planes, 27 | kernel_size=4, stride=2, padding=1, bias=True), 28 | nn.PReLU(out_planes) 29 | ) 30 | 31 | def conv_woact(in_planes, out_planes, kernel_size=3, stride=1, padding=1, dilation=1): 32 | return nn.Sequential( 33 | nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, 34 | padding=padding, dilation=dilation, bias=True), 35 | ) 36 | 37 | class ResBlock(nn.Module): 38 | def __init__(self, in_planes, out_planes, stride=2): 39 | super(ResBlock, self).__init__() 40 | if in_planes == out_planes and stride == 1: 41 | self.conv0 = nn.Identity() 42 | else: 43 | self.conv0 = nn.Conv2d(in_planes, out_planes, 44 | 3, stride, 1, bias=False) 45 | self.conv1 = conv(in_planes, out_planes, 3, stride, 1) 46 | self.conv2 = conv_woact(out_planes, out_planes, 3, 1, 1) 47 | self.relu1 = nn.PReLU(1) 48 | self.relu2 = nn.PReLU(out_planes) 49 | self.fc1 = nn.Conv2d(out_planes, 16, kernel_size=1, bias=False) 50 | self.fc2 = nn.Conv2d(16, out_planes, kernel_size=1, bias=False) 51 | 52 | def forward(self, x): 53 | y = self.conv0(x) 54 | x = self.conv1(x) 55 | x = self.conv2(x) 56 | w = x.mean(3, True).mean(2, True) 57 | w = self.relu1(self.fc1(w)) 58 | w = torch.sigmoid(self.fc2(w)) 59 | x = self.relu2(x * w + y) 60 | return x 61 | 62 | c = 32 63 | 64 | class ContextNet(nn.Module): 65 | def __init__(self): 66 | super(ContextNet, self).__init__() 67 | self.conv0 = conv(3, c, 3, 2, 1) 68 | self.conv1 = ResBlock(c, c) 69 | self.conv2 = ResBlock(c, 2*c) 70 | self.conv3 = ResBlock(2*c, 4*c) 71 | self.conv4 = ResBlock(4*c, 8*c) 72 | 73 | def forward(self, x, flow): 74 | x = self.conv0(x) 75 | x = self.conv1(x) 76 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", align_corners=False) * 0.5 77 | f1 = warp(x, flow) 78 | x = self.conv2(x) 79 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", 80 | align_corners=False) * 0.5 81 | f2 = warp(x, flow) 82 | x = self.conv3(x) 83 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", 84 | align_corners=False) * 0.5 85 | f3 = warp(x, flow) 86 | x = self.conv4(x) 87 | flow = F.interpolate(flow, scale_factor=0.5, mode="bilinear", 88 | align_corners=False) * 0.5 89 | f4 = warp(x, flow) 90 | return [f1, f2, f3, f4] 91 | 92 | 93 | class FusionNet(nn.Module): 94 | def __init__(self): 95 | super(FusionNet, self).__init__() 96 | self.conv0 = conv(8, c, 3, 2, 1) 97 | self.down0 = ResBlock(c, 2*c) 98 | self.down1 = ResBlock(4*c, 4*c) 99 | self.down2 = ResBlock(8*c, 8*c) 100 | self.down3 = ResBlock(16*c, 16*c) 101 | self.up0 = deconv(32*c, 8*c) 102 | self.up1 = deconv(16*c, 4*c) 103 | self.up2 = deconv(8*c, 2*c) 104 | self.up3 = deconv(4*c, c) 105 | self.conv = nn.Conv2d(c, 16, 3, 1, 1) 106 | self.up4 = nn.PixelShuffle(2) 107 | 108 | def forward(self, img0, img1, flow, c0, c1, flow_gt): 109 | warped_img0 = warp(img0, flow) 110 | warped_img1 = warp(img1, -flow) 111 | if flow_gt == None: 112 | warped_img0_gt, warped_img1_gt = None, None 113 | else: 114 | warped_img0_gt = warp(img0, flow_gt[:, :2]) 115 | warped_img1_gt = warp(img1, flow_gt[:, 2:4]) 116 | x = self.conv0(torch.cat((warped_img0, warped_img1, flow), 1)) 117 | s0 = self.down0(x) 118 | s1 = self.down1(torch.cat((s0, c0[0], c1[0]), 1)) 119 | s2 = self.down2(torch.cat((s1, c0[1], c1[1]), 1)) 120 | s3 = self.down3(torch.cat((s2, c0[2], c1[2]), 1)) 121 | x = self.up0(torch.cat((s3, c0[3], c1[3]), 1)) 122 | x = self.up1(torch.cat((x, s2), 1)) 123 | x = self.up2(torch.cat((x, s1), 1)) 124 | x = self.up3(torch.cat((x, s0), 1)) 125 | x = self.up4(self.conv(x)) 126 | return x, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt 127 | 128 | 129 | class Model: 130 | def __init__(self, local_rank=-1): 131 | self.flownet = IFNet() 132 | self.contextnet = ContextNet() 133 | self.fusionnet = FusionNet() 134 | self.device() 135 | self.optimG = AdamW(itertools.chain( 136 | self.flownet.parameters(), 137 | self.contextnet.parameters(), 138 | self.fusionnet.parameters()), lr=1e-6, weight_decay=1e-4) 139 | self.schedulerG = optim.lr_scheduler.CyclicLR( 140 | self.optimG, base_lr=1e-6, max_lr=1e-3, step_size_up=8000, cycle_momentum=False) 141 | self.epe = EPE() 142 | self.ter = Ternary() 143 | self.sobel = SOBEL() 144 | if local_rank != -1: 145 | self.flownet = DDP(self.flownet, device_ids=[ 146 | local_rank], output_device=local_rank) 147 | self.contextnet = DDP(self.contextnet, device_ids=[ 148 | local_rank], output_device=local_rank) 149 | self.fusionnet = DDP(self.fusionnet, device_ids=[ 150 | local_rank], output_device=local_rank) 151 | 152 | def train(self): 153 | self.flownet.train() 154 | self.contextnet.train() 155 | self.fusionnet.train() 156 | 157 | def eval(self): 158 | self.flownet.eval() 159 | self.contextnet.eval() 160 | self.fusionnet.eval() 161 | 162 | def device(self): 163 | self.flownet.to(device) 164 | self.contextnet.to(device) 165 | self.fusionnet.to(device) 166 | 167 | def load_model(self, path, rank): 168 | def convert(param): 169 | if rank == -1: 170 | return { 171 | k.replace("module.", ""): v 172 | for k, v in param.items() 173 | if "module." in k 174 | } 175 | else: 176 | return param 177 | if rank <= 0: 178 | self.flownet.load_state_dict( 179 | convert(torch.load('{}/flownet.pkl'.format(path), map_location=device))) 180 | self.contextnet.load_state_dict( 181 | convert(torch.load('{}/contextnet.pkl'.format(path), map_location=device))) 182 | self.fusionnet.load_state_dict( 183 | convert(torch.load('{}/unet.pkl'.format(path), map_location=device))) 184 | 185 | def save_model(self, path, rank): 186 | if rank == 0: 187 | torch.save(self.flownet.state_dict(), '{}/flownet.pkl'.format(path)) 188 | torch.save(self.contextnet.state_dict(), '{}/contextnet.pkl'.format(path)) 189 | torch.save(self.fusionnet.state_dict(), '{}/unet.pkl'.format(path)) 190 | 191 | def predict(self, imgs, flow, training=True, flow_gt=None): 192 | img0 = imgs[:, :3] 193 | img1 = imgs[:, 3:] 194 | c0 = self.contextnet(img0, flow) 195 | c1 = self.contextnet(img1, -flow) 196 | flow = F.interpolate(flow, scale_factor=2.0, mode="bilinear", 197 | align_corners=False) * 2.0 198 | refine_output, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.fusionnet( 199 | img0, img1, flow, c0, c1, flow_gt) 200 | res = torch.sigmoid(refine_output[:, :3]) * 2 - 1 201 | mask = torch.sigmoid(refine_output[:, 3:4]) 202 | merged_img = warped_img0 * mask + warped_img1 * (1 - mask) 203 | pred = merged_img + res 204 | pred = torch.clamp(pred, 0, 1) 205 | if training: 206 | return pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt 207 | else: 208 | return pred 209 | 210 | def inference(self, img0, img1, scale=1.0): 211 | imgs = torch.cat((img0, img1), 1) 212 | flow, _ = self.flownet(imgs, scale) 213 | return self.predict(imgs, flow, training=False) 214 | 215 | def update(self, imgs, gt, learning_rate=0, mul=1, training=True, flow_gt=None): 216 | for param_group in self.optimG.param_groups: 217 | param_group['lr'] = learning_rate 218 | if training: 219 | self.train() 220 | else: 221 | self.eval() 222 | flow, flow_list = self.flownet(imgs) 223 | pred, mask, merged_img, warped_img0, warped_img1, warped_img0_gt, warped_img1_gt = self.predict( 224 | imgs, flow, flow_gt=flow_gt) 225 | loss_ter = self.ter(pred, gt).mean() 226 | if training: 227 | with torch.no_grad(): 228 | loss_flow = torch.abs(warped_img0_gt - gt).mean() 229 | loss_mask = torch.abs( 230 | merged_img - gt).sum(1, True).float().detach() 231 | loss_mask = F.interpolate(loss_mask, scale_factor=0.5, mode="bilinear", 232 | align_corners=False).detach() 233 | flow_gt = (F.interpolate(flow_gt, scale_factor=0.5, mode="bilinear", 234 | align_corners=False) * 0.5).detach() 235 | loss_cons = 0 236 | for i in range(3): 237 | loss_cons += self.epe(flow_list[i], flow_gt[:, :2], 1) 238 | loss_cons += self.epe(-flow_list[i], flow_gt[:, 2:4], 1) 239 | loss_cons = loss_cons.mean() * 0.01 240 | else: 241 | loss_cons = torch.tensor([0]) 242 | loss_flow = torch.abs(warped_img0 - gt).mean() 243 | loss_mask = 1 244 | loss_l1 = (((pred - gt) ** 2 + 1e-6) ** 0.5).mean() 245 | if training: 246 | self.optimG.zero_grad() 247 | loss_G = loss_l1 + loss_cons + loss_ter 248 | loss_G.backward() 249 | self.optimG.step() 250 | return pred, merged_img, flow, loss_l1, loss_flow, loss_cons, loss_ter, loss_mask 251 | 252 | 253 | if __name__ == '__main__': 254 | img0 = torch.zeros(3, 3, 256, 256).float().to(device) 255 | img1 = torch.tensor(np.random.normal( 256 | 0, 1, (3, 3, 256, 256))).float().to(device) 257 | imgs = torch.cat((img0, img1), 1) 258 | model = Model() 259 | model.eval() 260 | print(model.inference(imgs).shape) 261 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Real-Time Intermediate Flow Estimation for Video Frame Interpolation 2 | ## [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) | [DeepWiki](https://deepwiki.com/hzwer/ECCV2022-RIFE) 3 | 4 | ## Introduction 5 | 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. 6 | 7 | **2024.08 - We find that [4.22.lite](https://github.com/hzwer/Practical-RIFE/tree/main?tab=readme-ov-file#model-list) is quite suitable for post-processing of [some diffusion model generated videos](https://drive.google.com/drive/folders/1hSzUn10Era3JCaVz0Z5Eg4wT9R6eJ9U9?usp=sharing).** 8 | 9 | 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. 10 | 11 | 2022.7.4 - Our paper is accepted by ECCV2022. Thanks to all relevant authors, contributors and users! 12 | 13 | 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! 14 | 15 | [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) 16 | 17 | **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)** 18 | 19 | 16X interpolation results from two input images: 20 | 21 | ![Demo](./demo/I2_slomo_clipped.gif) 22 | ![Demo](./demo/D2_slomo_clipped.gif) 23 | 24 | ## Software 25 | [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_PlayKit](https://github.com/hooke007/mpv_PlayKit) | [enhancr](https://github.com/mafiosnik777/enhancr) 26 | 27 | [RIFE-App(Paid)](https://grisk.itch.io/rife-app) | [Steam-VFI(Paid)](https://store.steampowered.com/app/1692080/SVFI/) 28 | 29 | 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. 30 | 31 | [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/styler00dollar/VapourSynth-RIFE-ncnn-Vulkan) | [vs-mlrt](https://github.com/AmusementClub/vs-mlrt) 32 | 33 | 34 | 35 | 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. 36 | 37 | You may check [this pull request](https://github.com/megvii-research/ECCV2022-RIFE/pull/300) for supporting macOS. 38 | ## CLI Usage 39 | 40 | ### Installation 41 | 42 | ``` 43 | git clone git@github.com:megvii-research/ECCV2022-RIFE.git 44 | cd ECCV2022-RIFE 45 | pip3 install -r requirements.txt 46 | ``` 47 | 48 | * 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/\*) 49 | 50 | * Unzip and move the pretrained parameters to train_log/\* 51 | 52 | * This model is not reported by our paper, for our paper model please refer to [evaluation](https://github.com/hzwer/ECCV2022-RIFE#evaluation). 53 | 54 | ### Run 55 | 56 | **Video Frame Interpolation** 57 | 58 | You can use our [demo video](https://drive.google.com/file/d/1i3xlKb7ax7Y70khcTcuePi6E7crO_dFc/view?usp=sharing) or your own video. 59 | ``` 60 | python3 inference_video.py --exp=1 --video=video.mp4 61 | ``` 62 | (generate video_2X_xxfps.mp4) 63 | ``` 64 | python3 inference_video.py --exp=2 --video=video.mp4 65 | ``` 66 | (for 4X interpolation) 67 | ``` 68 | python3 inference_video.py --exp=1 --video=video.mp4 --scale=0.5 69 | ``` 70 | (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.) 71 | ``` 72 | python3 inference_video.py --exp=2 --img=input/ 73 | ``` 74 | (to read video from pngs, like input/0.png ... input/612.png, ensure that the png names are numbers) 75 | ``` 76 | python3 inference_video.py --exp=2 --video=video.mp4 --fps=60 77 | ``` 78 | (add slomo effect, the audio will be removed) 79 | ``` 80 | python3 inference_video.py --video=video.mp4 --montage --png 81 | ``` 82 | (if you want to montage the origin video and save the png format output) 83 | 84 | **Extended Application** 85 | 86 | You may refer to [#278](https://github.com/megvii-research/ECCV2022-RIFE/issues/278#event-7199085190) for **Optical Flow Estimation** and refer to [#291](https://github.com/hzwer/ECCV2022-RIFE/issues/291#issuecomment-1328685348) for **Video Stitching**. 87 | 88 | **Image Interpolation** 89 | 90 | ``` 91 | python3 inference_img.py --img img0.png img1.png --exp=4 92 | ``` 93 | (2^4=16X interpolation results) 94 | After that, you can use pngs to generate mp4: 95 | ``` 96 | 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 97 | ``` 98 | You can also use pngs to generate gif: 99 | ``` 100 | 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 101 | ``` 102 | 103 | ### Run in docker 104 | Place the pre-trained models in `train_log/\*.pkl` (as above) 105 | 106 | Building the container: 107 | ``` 108 | docker build -t rife -f docker/Dockerfile . 109 | ``` 110 | 111 | Running the container: 112 | ``` 113 | docker run --rm -it -v $PWD:/host rife:latest inference_video --exp=1 --video=untitled.mp4 --output=untitled_rife.mp4 114 | ``` 115 | ``` 116 | docker run --rm -it -v $PWD:/host rife:latest inference_img --img img0.png img1.png --exp=4 117 | ``` 118 | 119 | Using gpu acceleration (requires proper gpu drivers for docker): 120 | ``` 121 | 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 122 | ``` 123 | 124 | ## Evaluation 125 | 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. 126 | 127 | **UCF101**: Download [UCF101 dataset](https://liuziwei7.github.io/projects/VoxelFlow) at ./UCF101/ucf101_interp_ours/ 128 | 129 | **Vimeo90K**: Download [Vimeo90K dataset](http://toflow.csail.mit.edu/) at ./vimeo_interp_test 130 | 131 | **MiddleBury**: Download [MiddleBury OTHER dataset](https://vision.middlebury.edu/flow/data/) at ./other-data and ./other-gt-interp 132 | 133 | **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). 134 | ``` 135 | # RIFE 136 | python3 benchmark/UCF101.py 137 | # "PSNR: 35.282 SSIM: 0.9688" 138 | python3 benchmark/Vimeo90K.py 139 | # "PSNR: 35.615 SSIM: 0.9779" 140 | python3 benchmark/MiddleBury_Other.py 141 | # "IE: 1.956" 142 | python3 benchmark/HD.py 143 | # "PSNR: 32.14" 144 | 145 | # RIFE_m 146 | python3 benchmark/HD_multi_4X.py 147 | # "PSNR: 22.96(544*1280), 31.87(720p), 34.25(1080p)" 148 | ``` 149 | 150 | ## Training and Reproduction 151 | Download [Vimeo90K dataset](http://toflow.csail.mit.edu/). 152 | 153 | We use 16 CPUs, 4 GPUs and 20G memory for training: 154 | ``` 155 | python3 -m torch.distributed.launch --nproc_per_node=4 train.py --world_size=4 156 | ``` 157 | 158 | ## Revision History 159 | 160 | 2021.3.18 [arXiv](https://arxiv.org/pdf/2011.06294v5.pdf): Modify the main experimental data, especially the runtime related issues. 161 | 162 | 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. 163 | 164 | 2021.11.17 [arXiv](https://arxiv.org/pdf/2011.06294v11.pdf): Support arbitrary-time frame interpolation, aka RIFEm and add more experiments. 165 | 166 | ## Recommend 167 | We sincerely recommend some related papers: 168 | 169 | 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) 170 | 171 | CVPR22 - [Video Frame Interpolation with Transformer](https://openaccess.thecvf.com/content/CVPR2022/html/Lu_Video_Frame_Interpolation_With_Transformer_CVPR_2022_paper.html) 172 | 173 | 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) 174 | 175 | CVPR23 - [A Dynamic Multi-Scale Voxel Flow Network for Video Prediction](https://huxiaotaostasy.github.io/DMVFN/) 176 | 177 | CVPR23 - [Extracting Motion and Appearance via Inter-Frame Attention for Efficient Video Frame Interpolation](https://arxiv.org/abs/2303.00440) 178 | 179 | ## Citation 180 | If you think this project is helpful, please feel free to leave a star or cite our paper: 181 | 182 | ``` 183 | @inproceedings{huang2022rife, 184 | title={Real-Time Intermediate Flow Estimation for Video Frame Interpolation}, 185 | author={Huang, Zhewei and Zhang, Tianyuan and Heng, Wen and Shi, Boxin and Zhou, Shuchang}, 186 | booktitle={Proceedings of the European Conference on Computer Vision (ECCV)}, 187 | year={2022} 188 | } 189 | ``` 190 | 191 | ## Reference 192 | 193 | Optical Flow: 194 | [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) 195 | 196 | Video Interpolation: 197 | [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) 198 | -------------------------------------------------------------------------------- /inference_video.py: -------------------------------------------------------------------------------- 1 | import os 2 | import cv2 3 | import torch 4 | import argparse 5 | import numpy as np 6 | from tqdm import tqdm 7 | from torch.nn import functional as F 8 | import warnings 9 | import _thread 10 | import skvideo.io 11 | from queue import Queue, Empty 12 | from model.pytorch_msssim import ssim_matlab 13 | 14 | warnings.filterwarnings("ignore") 15 | 16 | def transferAudio(sourceVideo, targetVideo): 17 | import shutil 18 | import moviepy.editor 19 | tempAudioFileName = "./temp/audio.mkv" 20 | 21 | # split audio from original video file and store in "temp" directory 22 | if True: 23 | 24 | # clear old "temp" directory if it exits 25 | if os.path.isdir("temp"): 26 | # remove temp directory 27 | shutil.rmtree("temp") 28 | # create new "temp" directory 29 | os.makedirs("temp") 30 | # extract audio from video 31 | os.system('ffmpeg -y -i "{}" -c:a copy -vn {}'.format(sourceVideo, tempAudioFileName)) 32 | 33 | targetNoAudio = os.path.splitext(targetVideo)[0] + "_noaudio" + os.path.splitext(targetVideo)[1] 34 | os.rename(targetVideo, targetNoAudio) 35 | # combine audio file and new video file 36 | os.system('ffmpeg -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo)) 37 | 38 | if os.path.getsize(targetVideo) == 0: # if ffmpeg failed to merge the video and audio together try converting the audio to aac 39 | tempAudioFileName = "./temp/audio.m4a" 40 | os.system('ffmpeg -y -i "{}" -c:a aac -b:a 160k -vn {}'.format(sourceVideo, tempAudioFileName)) 41 | os.system('ffmpeg -y -i "{}" -i {} -c copy "{}"'.format(targetNoAudio, tempAudioFileName, targetVideo)) 42 | if (os.path.getsize(targetVideo) == 0): # if aac is not supported by selected format 43 | os.rename(targetNoAudio, targetVideo) 44 | print("Audio transfer failed. Interpolated video will have no audio") 45 | else: 46 | print("Lossless audio transfer failed. Audio was transcoded to AAC (M4A) instead.") 47 | 48 | # remove audio-less video 49 | os.remove(targetNoAudio) 50 | else: 51 | os.remove(targetNoAudio) 52 | 53 | # remove temp directory 54 | shutil.rmtree("temp") 55 | 56 | parser = argparse.ArgumentParser(description='Interpolation for a pair of images') 57 | parser.add_argument('--video', dest='video', type=str, default=None) 58 | parser.add_argument('--output', dest='output', type=str, default=None) 59 | parser.add_argument('--img', dest='img', type=str, default=None) 60 | parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video') 61 | parser.add_argument('--model', dest='modelDir', type=str, default='train_log', help='directory with trained model files') 62 | parser.add_argument('--fp16', dest='fp16', action='store_true', help='fp16 mode for faster and more lightweight inference on cards with Tensor Cores') 63 | parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video') 64 | parser.add_argument('--scale', dest='scale', type=float, default=1.0, help='Try scale=0.5 for 4k video') 65 | parser.add_argument('--skip', dest='skip', action='store_true', help='whether to remove static frames before processing') 66 | parser.add_argument('--fps', dest='fps', type=int, default=None) 67 | parser.add_argument('--png', dest='png', action='store_true', help='whether to vid_out png format vid_outs') 68 | parser.add_argument('--ext', dest='ext', type=str, default='mp4', help='vid_out video extension') 69 | parser.add_argument('--exp', dest='exp', type=int, default=1) 70 | args = parser.parse_args() 71 | assert (not args.video is None or not args.img is None) 72 | if args.skip: 73 | print("skip flag is abandoned, please refer to issue #207.") 74 | if args.UHD and args.scale==1.0: 75 | args.scale = 0.5 76 | assert args.scale in [0.25, 0.5, 1.0, 2.0, 4.0] 77 | if not args.img is None: 78 | args.png = True 79 | 80 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 81 | torch.set_grad_enabled(False) 82 | if torch.cuda.is_available(): 83 | torch.backends.cudnn.enabled = True 84 | torch.backends.cudnn.benchmark = True 85 | if(args.fp16): 86 | torch.set_default_tensor_type(torch.cuda.HalfTensor) 87 | 88 | try: 89 | try: 90 | try: 91 | from model.RIFE_HDv2 import Model 92 | model = Model() 93 | model.load_model(args.modelDir, -1) 94 | print("Loaded v2.x HD model.") 95 | except: 96 | from train_log.RIFE_HDv3 import Model 97 | model = Model() 98 | model.load_model(args.modelDir, -1) 99 | print("Loaded v3.x HD model.") 100 | except: 101 | from model.RIFE_HD import Model 102 | model = Model() 103 | model.load_model(args.modelDir, -1) 104 | print("Loaded v1.x HD model") 105 | except: 106 | from model.RIFE import Model 107 | model = Model() 108 | model.load_model(args.modelDir, -1) 109 | print("Loaded ArXiv-RIFE model") 110 | model.eval() 111 | model.device() 112 | 113 | if not args.video is None: 114 | videoCapture = cv2.VideoCapture(args.video) 115 | fps = videoCapture.get(cv2.CAP_PROP_FPS) 116 | tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT) 117 | videoCapture.release() 118 | if args.fps is None: 119 | fpsNotAssigned = True 120 | args.fps = fps * (2 ** args.exp) 121 | else: 122 | fpsNotAssigned = False 123 | videogen = skvideo.io.vreader(args.video) 124 | lastframe = next(videogen) 125 | fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') 126 | video_path_wo_ext, ext = os.path.splitext(args.video) 127 | print('{}.{}, {} frames in total, {}FPS to {}FPS'.format(video_path_wo_ext, args.ext, tot_frame, fps, args.fps)) 128 | if args.png == False and fpsNotAssigned == True: 129 | print("The audio will be merged after interpolation process") 130 | else: 131 | print("Will not merge audio because using png or fps flag!") 132 | else: 133 | videogen = [] 134 | for f in os.listdir(args.img): 135 | if 'png' in f: 136 | videogen.append(f) 137 | tot_frame = len(videogen) 138 | videogen.sort(key= lambda x:int(x[:-4])) 139 | lastframe = cv2.imread(os.path.join(args.img, videogen[0]), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy() 140 | videogen = videogen[1:] 141 | h, w, _ = lastframe.shape 142 | vid_out_name = None 143 | vid_out = None 144 | if args.png: 145 | if not os.path.exists('vid_out'): 146 | os.mkdir('vid_out') 147 | else: 148 | if args.output is not None: 149 | vid_out_name = args.output 150 | else: 151 | vid_out_name = '{}_{}X_{}fps.{}'.format(video_path_wo_ext, (2 ** args.exp), int(np.round(args.fps)), args.ext) 152 | vid_out = cv2.VideoWriter(vid_out_name, fourcc, args.fps, (w, h)) 153 | 154 | def clear_write_buffer(user_args, write_buffer): 155 | cnt = 0 156 | while True: 157 | item = write_buffer.get() 158 | if item is None: 159 | break 160 | if user_args.png: 161 | cv2.imwrite('vid_out/{:0>7d}.png'.format(cnt), item[:, :, ::-1]) 162 | cnt += 1 163 | else: 164 | vid_out.write(item[:, :, ::-1]) 165 | 166 | def build_read_buffer(user_args, read_buffer, videogen): 167 | try: 168 | for frame in videogen: 169 | if not user_args.img is None: 170 | frame = cv2.imread(os.path.join(user_args.img, frame), cv2.IMREAD_UNCHANGED)[:, :, ::-1].copy() 171 | if user_args.montage: 172 | frame = frame[:, left: left + w] 173 | read_buffer.put(frame) 174 | except: 175 | pass 176 | read_buffer.put(None) 177 | 178 | def make_inference(I0, I1, n): 179 | global model 180 | middle = model.inference(I0, I1, args.scale) 181 | if n == 1: 182 | return [middle] 183 | first_half = make_inference(I0, middle, n=n//2) 184 | second_half = make_inference(middle, I1, n=n//2) 185 | if n%2: 186 | return [*first_half, middle, *second_half] 187 | else: 188 | return [*first_half, *second_half] 189 | 190 | def pad_image(img): 191 | if(args.fp16): 192 | return F.pad(img, padding).half() 193 | else: 194 | return F.pad(img, padding) 195 | 196 | if args.montage: 197 | left = w // 4 198 | w = w // 2 199 | tmp = max(32, int(32 / args.scale)) 200 | ph = ((h - 1) // tmp + 1) * tmp 201 | pw = ((w - 1) // tmp + 1) * tmp 202 | padding = (0, pw - w, 0, ph - h) 203 | pbar = tqdm(total=tot_frame) 204 | if args.montage: 205 | lastframe = lastframe[:, left: left + w] 206 | write_buffer = Queue(maxsize=500) 207 | read_buffer = Queue(maxsize=500) 208 | _thread.start_new_thread(build_read_buffer, (args, read_buffer, videogen)) 209 | _thread.start_new_thread(clear_write_buffer, (args, write_buffer)) 210 | 211 | I1 = torch.from_numpy(np.transpose(lastframe, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255. 212 | I1 = pad_image(I1) 213 | temp = None # save lastframe when processing static frame 214 | 215 | while True: 216 | if temp is not None: 217 | frame = temp 218 | temp = None 219 | else: 220 | frame = read_buffer.get() 221 | if frame is None: 222 | break 223 | I0 = I1 224 | I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255. 225 | I1 = pad_image(I1) 226 | I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False) 227 | I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False) 228 | ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) 229 | 230 | break_flag = False 231 | if ssim > 0.996: 232 | frame = read_buffer.get() # read a new frame 233 | if frame is None: 234 | break_flag = True 235 | frame = lastframe 236 | else: 237 | temp = frame 238 | I1 = torch.from_numpy(np.transpose(frame, (2,0,1))).to(device, non_blocking=True).unsqueeze(0).float() / 255. 239 | I1 = pad_image(I1) 240 | I1 = model.inference(I0, I1, args.scale) 241 | I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False) 242 | ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) 243 | frame = (I1[0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w] 244 | 245 | if ssim < 0.2: 246 | output = [] 247 | for i in range((2 ** args.exp) - 1): 248 | output.append(I0) 249 | ''' 250 | output = [] 251 | step = 1 / (2 ** args.exp) 252 | alpha = 0 253 | for i in range((2 ** args.exp) - 1): 254 | alpha += step 255 | beta = 1-alpha 256 | 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.) 257 | ''' 258 | else: 259 | output = make_inference(I0, I1, 2**args.exp-1) if args.exp else [] 260 | 261 | if args.montage: 262 | write_buffer.put(np.concatenate((lastframe, lastframe), 1)) 263 | for mid in output: 264 | mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0))) 265 | write_buffer.put(np.concatenate((lastframe, mid[:h, :w]), 1)) 266 | else: 267 | write_buffer.put(lastframe) 268 | for mid in output: 269 | mid = (((mid[0] * 255.).byte().cpu().numpy().transpose(1, 2, 0))) 270 | write_buffer.put(mid[:h, :w]) 271 | pbar.update(1) 272 | lastframe = frame 273 | if break_flag: 274 | break 275 | 276 | if args.montage: 277 | write_buffer.put(np.concatenate((lastframe, lastframe), 1)) 278 | else: 279 | write_buffer.put(lastframe) 280 | 281 | write_buffer.put(None) 282 | 283 | import time 284 | while(not write_buffer.empty()): 285 | time.sleep(0.1) 286 | pbar.close() 287 | if not vid_out is None: 288 | vid_out.release() 289 | 290 | # move audio to new video file if appropriate 291 | if args.png == False and fpsNotAssigned == True and not args.video is None: 292 | try: 293 | transferAudio(args.video, vid_out_name) 294 | except: 295 | print("Audio transfer failed. Interpolated video will have no audio") 296 | targetNoAudio = os.path.splitext(vid_out_name)[0] + "_noaudio" + os.path.splitext(vid_out_name)[1] 297 | os.rename(targetNoAudio, vid_out_name) 298 | --------------------------------------------------------------------------------