File size: 6,652 Bytes
04fbff5 |
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 |
import os
import cv2
import glob
import torch
import numpy as np
from tqdm import tqdm
from omegaconf import OmegaConf
from vbench.utils import load_dimension_info
from vbench.third_party.amt.utils.utils import (
img2tensor, tensor2img,
check_dim_and_resize
)
from vbench.third_party.amt.utils.build_utils import build_from_cfg
from vbench.third_party.amt.utils.utils import InputPadder
class FrameProcess:
def __init__(self):
pass
def get_frames(self, video_path):
frame_list = []
video = cv2.VideoCapture(video_path)
while video.isOpened():
success, frame = video.read()
if success:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # convert to rgb
frame_list.append(frame)
else:
break
video.release()
assert frame_list != []
return frame_list
def get_frames_from_img_folder(self, img_folder):
exts = ['jpg', 'png', 'jpeg', 'bmp', 'tif',
'tiff', 'JPG', 'PNG', 'JPEG', 'BMP',
'TIF', 'TIFF']
frame_list = []
imgs = sorted([p for p in glob.glob(os.path.join(img_folder, "*")) if os.path.splitext(p)[1][1:] in exts])
# imgs = sorted(glob.glob(os.path.join(img_folder, "*.png")))
for img in imgs:
frame = cv2.imread(img, cv2.IMREAD_COLOR)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame_list.append(frame)
assert frame_list != []
return frame_list
def extract_frame(self, frame_list, start_from=0):
extract = []
for i in range(start_from, len(frame_list), 2):
extract.append(frame_list[i])
return extract
class MotionSmoothness:
def __init__(self, config, ckpt, device):
self.device = device
self.config = config
self.ckpt = ckpt
self.niters = 1
self.initialization()
self.load_model()
def load_model(self):
cfg_path = self.config
ckpt_path = self.ckpt
network_cfg = OmegaConf.load(cfg_path).network
network_name = network_cfg.name
print(f'Loading [{network_name}] from [{ckpt_path}]...')
self.model = build_from_cfg(network_cfg)
ckpt = torch.load(ckpt_path)
self.model.load_state_dict(ckpt['state_dict'])
self.model = self.model.to(self.device)
self.model.eval()
def initialization(self):
if self.device == 'cuda':
self.anchor_resolution = 1024 * 512
self.anchor_memory = 1500 * 1024**2
self.anchor_memory_bias = 2500 * 1024**2
self.vram_avail = torch.cuda.get_device_properties(self.device).total_memory
print("VRAM available: {:.1f} MB".format(self.vram_avail / 1024 ** 2))
else:
# Do not resize in cpu mode
self.anchor_resolution = 8192*8192
self.anchor_memory = 1
self.anchor_memory_bias = 0
self.vram_avail = 1
self.embt = torch.tensor(1/2).float().view(1, 1, 1, 1).to(self.device)
self.fp = FrameProcess()
def motion_score(self, video_path):
iters = int(self.niters)
# get inputs
if video_path.endswith('.mp4'):
frames = self.fp.get_frames(video_path)
elif os.path.isdir(video_path):
frames = self.fp.get_frames_from_img_folder(video_path)
else:
raise NotImplementedError
frame_list = self.fp.extract_frame(frames, start_from=0)
# print(f'Loading [images] from [{video_path}], the number of images = [{len(frame_list)}]')
inputs = [img2tensor(frame).to(self.device) for frame in frame_list]
assert len(inputs) > 1, f"The number of input should be more than one (current {len(inputs)})"
inputs = check_dim_and_resize(inputs)
h, w = inputs[0].shape[-2:]
scale = self.anchor_resolution / (h * w) * np.sqrt((self.vram_avail - self.anchor_memory_bias) / self.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(inputs[0].shape, padding)
inputs = padder.pad(*inputs)
# ----------------------- Interpolater -----------------------
# print(f'Start frame interpolation:')
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(self.device)
in_1 = in_1.to(self.device)
with torch.no_grad():
imgt_pred = self.model(in_0, in_1, self.embt, scale_factor=scale, eval=True)['imgt_pred']
outputs += [imgt_pred.cpu(), in_1.cpu()]
inputs = outputs
# ----------------------- cal_vfi_score -----------------------
outputs = padder.unpad(*outputs)
outputs = [tensor2img(out) for out in outputs]
vfi_score = self.vfi_score(frames, outputs)
norm = (255.0 - vfi_score)/255.0
return norm
def vfi_score(self, ori_frames, interpolate_frames):
ori = self.fp.extract_frame(ori_frames, start_from=1)
interpolate = self.fp.extract_frame(interpolate_frames, start_from=1)
scores = []
for i in range(len(interpolate)):
scores.append(self.get_diff(ori[i], interpolate[i]))
return np.mean(np.array(scores))
def get_diff(self, img1, img2):
img = cv2.absdiff(img1, img2)
return np.mean(img)
def motion_smoothness(motion, video_list):
sim = []
video_results = []
for video_path in tqdm(video_list):
score_per_video = motion.motion_score(video_path)
video_results.append({'video_path': video_path, 'video_results': score_per_video})
sim.append(score_per_video)
avg_score = np.mean(sim)
return avg_score, video_results
def compute_motion_smoothness(json_dir, device, submodules_list):
config = submodules_list["config"] # pretrained/amt_model/AMT-S.yaml
ckpt = submodules_list["ckpt"] # pretrained/amt_model/amt-s.pth
motion = MotionSmoothness(config, ckpt, device)
video_list, _ = load_dimension_info(json_dir, dimension='motion_smoothness', lang='en')
all_results, video_results = motion_smoothness(motion, video_list)
return all_results, video_results
|