# 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 AMT_G = { 'name': 'networks.AMT-G.Model', 'params':{ 'corr_radius': 3, 'corr_lvls': 4, 'num_flows': 5, } } 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(ckpt_path, device="cuda"): ''' load the frame interpolation model. ''' network_cfg = AMT_G 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, input_path, 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] _, file_name_with_extension = os.path.split(input_path) file_name, _ = os.path.splitext(file_name_with_extension) save_video_path = f'{output_path}/output_{file_name}.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('--ckpt', type=str, default='amt-g.pth', help="The pretrained model.") parser.add_argument('--niters', type=int, default=1, help="Iter of Interpolation. The number of frames will be double after per iter.") parser.add_argument('--input', default="test.mp4", help="Input video.") parser.add_argument('--output_path', type=str, default='results', help="Output path.") parser.add_argument('--frame_rate', type=int, default=30, help="Frames rate of the output video.") args = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 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(ckpt_path, device) outputs = interpolater(model, inputs, scale, padder, iters) write(outputs, input_path, output_path, frame_rate)