Spaces:
Runtime error
Runtime error
File size: 6,643 Bytes
b3f324b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 |
# 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)
|