def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Model()
model.load_model('../train_log/HDv2', -1)
model.eval()
model.device()
path = '../../data/UCF101/ucf101_interp_ours/'
dirs = os.listdir(path)
psnr_list = []
ssim_list = []
time_list = []
# print('=========>Start Calculate PSNR and SSIM')
for d in tqdm(dirs):
img0 = (path + d + '/frame_00.png')
img1 = (path + d + '/frame_02.png')
gt = (path + d + '/frame_01_gt.png')
img0 = (torch.tensor(cv2.imread(img0).transpose(2, 0, 1) /
255.)).to(device).float().unsqueeze(0)
img1 = (torch.tensor(cv2.imread(img1).transpose(2, 0, 1) /
255.)).to(device).float().unsqueeze(0)
gt = (torch.tensor(cv2.imread(gt).transpose(2, 0, 1) /
255.)).to(device).float().unsqueeze(0)
# Avg PSNR: 35.243666269214145 SSIM: 0.9683315742368154 Time: 0.133457749911717
# inference
pred = model.inference(img0, img1)[0]
# Calculate indicators
out = pred.detach().cpu().numpy().transpose(1, 2, 0)
out = np.round(out * 255) / 255.
gt = gt[0].cpu().numpy().transpose(1, 2, 0)
psnr = compute_psnr(gt, out)
ssim = compute_ssim(gt, out)
psnr_list.append(psnr)
ssim_list.append(ssim)
# print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list)))
# print('=========>Start Calculate Inference Time')
# inference time
for i in range(100):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
pred = model.inference(img0, img1)[0]
end.record()
torch.cuda.synchronize()
time_list.append(start.elapsed_time(end))
time_list.remove(min(time_list))
time_list.remove(max(time_list))
print("Avg PSNR: {} SSIM: {} Time: {}".format(np.mean(psnr_list),
np.mean(ssim_list),
np.mean(time_list) / 100))
parser.add_argument('--rmaxcycles',
default=8,
type=int,
help='limit max number of bisectional cycles')
parser.add_argument('--model',
dest='modelDir',
type=str,
default='train_log',
help='directory with trained model files')
args = parser.parse_args()
try:
try:
from model.RIFE_HDv2 import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v2.x HD model.")
except:
from train_log.RIFE_HDv3 import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v3.x HD model.")
except:
from model.RIFE_HD import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v1.x HD model")
model.eval()
model.device()
args.scale = 0.5
assert args.scale in [0.25, 0.5, 1.0, 2.0, 4.0]
if not args.img is None:
args.png = True
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_grad_enabled(False)
if torch.cuda.is_available():
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
if(args.fp16):
torch.set_default_tensor_type(torch.cuda.HalfTensor)
try:
from model.RIFE_HDv2 import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v2.x HD model.")
except:
from model.RIFE_HD import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v1.x HD model")
model.eval()
model.device()
if not args.video is None:
videoCapture = cv2.VideoCapture(args.video)
fps = videoCapture.get(cv2.CAP_PROP_FPS)
tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
videoCapture.release()
parser.add_argument('--output', dest='output', type=str, default=None)
parser.add_argument('--img', dest='img', type=str, default=None)
parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video')
parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video')
parser.add_argument('--skip', dest='skip', action='store_true', help='whether to remove static frames before processing')
parser.add_argument('--fps', dest='fps', type=int, default=None)
parser.add_argument('--png', dest='png', action='store_true', help='whether to vid_out png format vid_outs')
parser.add_argument('--ext', dest='ext', type=str, default='mp4', help='vid_out video extension')
parser.add_argument('--exp', dest='exp', type=int, default=1)
args = parser.parse_args()
assert (not args.video is None or not args.img is None)
if not args.img is None:
args.png = True
from model.RIFE_HDv2 import Model
model = Model()
model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), -1)
model.eval()
model.device()
if not args.video is None:
videoCapture = cv2.VideoCapture(args.video)
fps = videoCapture.get(cv2.CAP_PROP_FPS)
tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
videoCapture.release()
if args.fps is None:
fpsNotAssigned = True
args.fps = fps * (2 ** args.exp)
else:
fpsNotAssigned = False
videogen = skvideo.io.vreader(args.video)
eval_time_interval = time.time() - time_stamp
if local_rank == 0:
print('eval time: {}'.format(eval_time_interval))
writer_val.add_scalar('loss_l1', np.array(loss_l1_list).mean(), nr_eval)
writer_val.add_scalar('loss_flow', np.array(loss_flow_list).mean(), nr_eval)
writer_val.add_scalar('loss_cons', np.array(loss_cons_list).mean(), nr_eval)
writer_val.add_scalar('loss_ter', np.array(loss_ter_list).mean(), nr_eval)
writer_val.add_scalar('psnr', np.array(psnr_list).mean(), nr_eval)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='slomo')
parser.add_argument('--epoch', default=300, type=int)
parser.add_argument('--batch_size', default=16, type=int, help='minibatch size')
parser.add_argument('--local_rank', default=0, type=int, help='local rank')
parser.add_argument('--world_size', default=4, type=int, help='world size')
args = parser.parse_args()
torch.distributed.init_process_group(backend="nccl", world_size=args.world_size)
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)
seed = 1234
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.benchmark = True
model = Model(args.local_rank)
model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), 0)
train(model, args.local_rank)
parser.add_argument('--ratio',
default=0,
type=float,
help='inference ratio between two images with 0 - 1 range')
parser.add_argument(
'--rthreshold',
default=0.02,
type=float,
help='returns image when actual ratio falls in given range threshold')
parser.add_argument('--rmaxcycles',
default=8,
type=int,
help='limit max number of bisectional cycles')
args = parser.parse_args()
model = Model()
model.load_model(
os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), -1)
model.eval()
model.device()
if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
img0 = cv2.imread(args.img[0], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
img1 = cv2.imread(args.img[1], cv2.IMREAD_COLOR | cv2.IMREAD_ANYDEPTH)
img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device)).unsqueeze(0)
img1 = (torch.tensor(img1.transpose(2, 0, 1)).to(device)).unsqueeze(0)
else:
img0 = cv2.imread(args.img[0])
img1 = cv2.imread(args.img[1])
img0 = (torch.tensor(img0.transpose(2, 0, 1)).to(device) /
