|
|
import os
|
|
|
import torch
|
|
|
from torch.nn import functional as F
|
|
|
|
|
|
from .ssim import ssim_matlab
|
|
|
|
|
|
from .RIFE_HDv3 import Model
|
|
|
|
|
|
def get_frame(frames, frame_no):
|
|
|
if frame_no >= frames.shape[1]:
|
|
|
return None
|
|
|
frame = (frames[:, frame_no] + 1) /2
|
|
|
frame = frame.clip(0., 1.)
|
|
|
return frame
|
|
|
|
|
|
def add_frame(frames, frame, h, w):
|
|
|
frame = (frame * 2) - 1
|
|
|
frame = frame.clip(-1., 1.)
|
|
|
frame = frame.squeeze(0)
|
|
|
frame = frame[:, :h, :w]
|
|
|
frame = frame.unsqueeze(1)
|
|
|
frames.append(frame.cpu())
|
|
|
|
|
|
def process_frames(model, device, frames, exp):
|
|
|
pos = 0
|
|
|
output_frames = []
|
|
|
|
|
|
lastframe = get_frame(frames, 0)
|
|
|
_, h, w = lastframe.shape
|
|
|
scale = 1
|
|
|
fp16 = False
|
|
|
|
|
|
def make_inference(I0, I1, n):
|
|
|
middle = model.inference(I0, I1, scale)
|
|
|
if n == 1:
|
|
|
return [middle]
|
|
|
first_half = make_inference(I0, middle, n=n//2)
|
|
|
second_half = make_inference(middle, I1, n=n//2)
|
|
|
if n%2:
|
|
|
return [*first_half, middle, *second_half]
|
|
|
else:
|
|
|
return [*first_half, *second_half]
|
|
|
|
|
|
tmp = max(32, int(32 / scale))
|
|
|
ph = ((h - 1) // tmp + 1) * tmp
|
|
|
pw = ((w - 1) // tmp + 1) * tmp
|
|
|
padding = (0, pw - w, 0, ph - h)
|
|
|
|
|
|
def pad_image(img):
|
|
|
if(fp16):
|
|
|
return F.pad(img, padding).half()
|
|
|
else:
|
|
|
return F.pad(img, padding)
|
|
|
|
|
|
I1 = lastframe.to(device, non_blocking=True).unsqueeze(0)
|
|
|
I1 = pad_image(I1)
|
|
|
temp = None
|
|
|
|
|
|
while True:
|
|
|
if temp is not None:
|
|
|
frame = temp
|
|
|
temp = None
|
|
|
else:
|
|
|
pos += 1
|
|
|
frame = get_frame(frames, pos)
|
|
|
if frame is None:
|
|
|
break
|
|
|
I0 = I1
|
|
|
I1 = frame.to(device, non_blocking=True).unsqueeze(0)
|
|
|
I1 = pad_image(I1)
|
|
|
I0_small = F.interpolate(I0, (32, 32), mode='bilinear', align_corners=False)
|
|
|
I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
|
|
|
ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
|
|
|
|
|
|
break_flag = False
|
|
|
if ssim > 0.996 or pos > 100:
|
|
|
pos += 1
|
|
|
frame = get_frame(frames, pos)
|
|
|
if frame is None:
|
|
|
break_flag = True
|
|
|
frame = lastframe
|
|
|
else:
|
|
|
temp = frame
|
|
|
I1 = frame.to(device, non_blocking=True).unsqueeze(0)
|
|
|
I1 = pad_image(I1)
|
|
|
I1 = model.inference(I0, I1, scale)
|
|
|
I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False)
|
|
|
ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3])
|
|
|
frame = I1[0][:, :h, :w]
|
|
|
|
|
|
if ssim < 0.2:
|
|
|
output = []
|
|
|
for _ in range((2 ** exp) - 1):
|
|
|
output.append(I0)
|
|
|
else:
|
|
|
output = make_inference(I0, I1, 2**exp-1) if exp else []
|
|
|
|
|
|
add_frame(output_frames, lastframe, h, w)
|
|
|
for mid in output:
|
|
|
add_frame(output_frames, mid, h, w)
|
|
|
lastframe = frame
|
|
|
if break_flag:
|
|
|
break
|
|
|
|
|
|
add_frame(output_frames, lastframe, h, w)
|
|
|
return torch.cat( output_frames, dim=1)
|
|
|
|
|
|
def temporal_interpolation(model_path, frames, exp, device ="cuda"):
|
|
|
|
|
|
model = Model()
|
|
|
model.load_model(model_path, -1, device=device)
|
|
|
|
|
|
model.eval()
|
|
|
model.to(device=device)
|
|
|
|
|
|
with torch.no_grad():
|
|
|
output = process_frames(model, device, frames.float(), exp)
|
|
|
|
|
|
return output
|
|
|
|
