189 lines
6.3 KiB
Python
189 lines
6.3 KiB
Python
# this script is modified from https://github.com/MCG-NKU/AMT/blob/main/demos/demo_2x.py
|
|
from json import load
|
|
import os
|
|
import cv2
|
|
import sys
|
|
import glob
|
|
import torch
|
|
import argparse
|
|
import numpy as np
|
|
import os.path as osp
|
|
from warnings import warn
|
|
from omegaconf import OmegaConf
|
|
from torchvision.utils import make_grid
|
|
sys.path.append('.')
|
|
from utils.utils import (
|
|
read, write,
|
|
img2tensor, tensor2img,
|
|
check_dim_and_resize
|
|
)
|
|
from utils.build_utils import build_from_cfg
|
|
from utils.utils import InputPadder
|
|
|
|
|
|
def init(device="cuda"):
|
|
|
|
'''
|
|
initialize the device and the anchor resolution.
|
|
'''
|
|
|
|
if device == 'cuda':
|
|
anchor_resolution = 1024 * 512
|
|
anchor_memory = 1500 * 1024**2
|
|
anchor_memory_bias = 2500 * 1024**2
|
|
vram_avail = torch.cuda.get_device_properties(device).total_memory
|
|
print("VRAM available: {:.1f} MB".format(vram_avail / 1024 ** 2))
|
|
else:
|
|
# Do not resize in cpu mode
|
|
anchor_resolution = 8192*8192
|
|
anchor_memory = 1
|
|
anchor_memory_bias = 0
|
|
vram_avail = 1
|
|
|
|
return anchor_resolution, anchor_memory, anchor_memory_bias, vram_avail
|
|
|
|
def get_input_video_from_path(input_path, device="cuda"):
|
|
|
|
'''
|
|
Get the input video from the input_path.
|
|
|
|
params:
|
|
input_path: str, the path of the input video.
|
|
devices: str, the device to run the model.
|
|
returns:
|
|
inputs: list, the list of the input frames.
|
|
scale: float, the scale of the input frames.
|
|
padder: InputPadder, the padder to pad the input frames.
|
|
'''
|
|
|
|
anchor_resolution, anchor_memory, anchor_memory_bias, vram_avail = init(device)
|
|
|
|
if osp.splitext(input_path)[-1] in ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv',
|
|
'.webm', '.MP4', '.AVI', '.MOV', '.MKV', '.FLV',
|
|
'.WMV', '.WEBM']:
|
|
|
|
vcap = cv2.VideoCapture(input_path)
|
|
|
|
inputs = []
|
|
w = int(vcap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
h = int(vcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
scale = anchor_resolution / (h * w) * np.sqrt((vram_avail - anchor_memory_bias) / anchor_memory)
|
|
scale = 1 if scale > 1 else scale
|
|
scale = 1 / np.floor(1 / np.sqrt(scale) * 16) * 16
|
|
if scale < 1:
|
|
print(f"Due to the limited VRAM, the video will be scaled by {scale:.2f}")
|
|
padding = int(16 / scale)
|
|
padder = InputPadder((h, w), padding)
|
|
while True:
|
|
ret, frame = vcap.read()
|
|
if ret is False:
|
|
break
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
frame_t = img2tensor(frame).to(device)
|
|
frame_t = padder.pad(frame_t)
|
|
inputs.append(frame_t)
|
|
print(f'Loading the [video] from {input_path}, the number of frames [{len(inputs)}]')
|
|
else:
|
|
raise TypeError("Input should be a video.")
|
|
|
|
return inputs, scale, padder
|
|
|
|
|
|
def load_model(cfg_path, ckpt_path, device="cuda"):
|
|
|
|
'''
|
|
load the frame interpolation model.
|
|
'''
|
|
|
|
network_cfg = OmegaConf.load(cfg_path).network
|
|
network_name = network_cfg.name
|
|
print(f'Loading [{network_name}] from [{ckpt_path}]...')
|
|
model = build_from_cfg(network_cfg)
|
|
ckpt = torch.load(ckpt_path)
|
|
model.load_state_dict(ckpt['state_dict'])
|
|
model = model.to(device)
|
|
model.eval()
|
|
return model
|
|
|
|
def interpolater(model, inputs, scale, padder, iters=1):
|
|
|
|
'''
|
|
interpolating with the interpolation model.
|
|
|
|
params:
|
|
model: nn.Module, the frame interpolation model.
|
|
inputs: list, the list of the input frames.
|
|
scale: float, the scale of the input frames.
|
|
iters: int, the number of iterations of interpolation. The final frames model generating is 2 ** iters * (m - 1) + 1 and m is input frames.
|
|
returns:
|
|
outputs: list, the list of the output frames.
|
|
'''
|
|
|
|
print(f'Start frame interpolation:')
|
|
embt = torch.tensor(1/2).float().view(1, 1, 1, 1).to(device)
|
|
|
|
for i in range(iters):
|
|
print(f'Iter {i+1}. input_frames={len(inputs)} output_frames={2*len(inputs)-1}')
|
|
outputs = [inputs[0]]
|
|
for in_0, in_1 in zip(inputs[:-1], inputs[1:]):
|
|
in_0 = in_0.to(device)
|
|
in_1 = in_1.to(device)
|
|
with torch.no_grad():
|
|
imgt_pred = model(in_0, in_1, embt, scale_factor=scale, eval=True)['imgt_pred']
|
|
outputs += [imgt_pred.cpu(), in_1.cpu()]
|
|
inputs = outputs
|
|
|
|
outputs = padder.unpad(*outputs)
|
|
|
|
return outputs
|
|
|
|
def write(outputs, output_path, frame_rate=30):
|
|
'''
|
|
write results to the output_path.
|
|
'''
|
|
|
|
if osp.exists(output_path) is False:
|
|
os.makedirs(output_path)
|
|
|
|
|
|
size = outputs[0].shape[2:][::-1]
|
|
|
|
video_count = len(glob.glob(os.path.join(output_path, "*.mp4")))
|
|
|
|
save_video_path = f'{output_path}/demo_{video_count:04d}.mp4'
|
|
writer = cv2.VideoWriter(save_video_path, cv2.VideoWriter_fourcc(*"mp4v"),
|
|
frame_rate, size)
|
|
|
|
for i, imgt_pred in enumerate(outputs):
|
|
imgt_pred = tensor2img(imgt_pred)
|
|
imgt_pred = cv2.cvtColor(imgt_pred, cv2.COLOR_RGB2BGR)
|
|
writer.write(imgt_pred)
|
|
print(f"Demo video is saved to [{save_video_path}]")
|
|
|
|
writer.release()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('-c', '--config', type=str, default='cfgs/AMT-G.yaml')
|
|
parser.add_argument('-p', '--ckpt', type=str, default='amt-g.pth')
|
|
parser.add_argument('-n', '--niters', type=int, default=6)
|
|
parser.add_argument('-i', '--input', default="test.mp4")
|
|
parser.add_argument('-o', '--output_path', type=str, default='results')
|
|
parser.add_argument('-r', '--frame_rate', type=int, default=30)
|
|
|
|
args = parser.parse_args()
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
cfg_path = args.config
|
|
ckpt_path = args.ckpt
|
|
input_path = args.input
|
|
output_path = args.output_path
|
|
iters = int(args.niters)
|
|
frame_rate = int(args.frame_rate)
|
|
|
|
inputs, scale, padder = get_input_video_from_path(input_path, device)
|
|
model = load_model(cfg_path, ckpt_path, device)
|
|
outputs = interpolater(model, inputs, scale, padder, iters)
|
|
write(outputs, output_path, frame_rate)
|