# -*- coding: utf-8 -*-
import sys
sys.path.append(".")
import os
import cv2
import numpy as np
import argparse
from PIL import Image
import torch.nn.functional as F
import torch
from torch.utils.data import DataLoader
from model.modules.flow_comp_raft import RAFT_bi
from model.recurrent_flow_completion import RecurrentFlowCompleteNet
from model.propainter import InpaintGenerator
# from core.dataset import TestDataset
from core.dataset import TestDataset
from core.metrics import calc_psnr_and_ssim, calculate_i3d_activations, calculate_vfid, init_i3d_model
from time import time
import warnings
warnings.filterwarnings("ignore")
# sample reference frames from the whole video
def get_ref_index(neighbor_ids, length, ref_stride=10):
ref_index = []
for i in range(0, length, ref_stride):
if i not in neighbor_ids:
ref_index.append(i)
return ref_index
def main_worker(args):
args.size = (args.width, args.height)
w, h = args.size
# set up datasets and data loader
assert (args.dataset == 'davis') or args.dataset == 'youtube-vos', \
f"{args.dataset} dataset is not supported"
test_dataset = TestDataset(vars(args))
test_loader = DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=args.num_workers)
# set up models
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
fix_raft = RAFT_bi(args.raft_model_path, device)
fix_flow_complete = RecurrentFlowCompleteNet(args.fc_model_path)
for p in fix_flow_complete.parameters():
p.requires_grad = False
fix_flow_complete.to(device)
fix_flow_complete.eval()
model = InpaintGenerator(model_path=args.propainter_model_path).to(device)
model.eval()
time_all = []
print('Start evaluation ...')
if args.task == 'video_completion':
result_path = os.path.join(f'results_eval',
f'{args.dataset}_rs_{args.ref_stride}_nl_{args.neighbor_length}_video_completion')
if not os.path.exists(result_path):
os.makedirs(result_path, exist_ok=True)
eval_summary = open(os.path.join(result_path, f"{args.dataset}_metrics.txt"),"w")
total_frame_psnr = []
total_frame_ssim = []
output_i3d_activations = []
real_i3d_activations = []
i3d_model = init_i3d_model('weights/i3d_rgb_imagenet.pt')
else:
result_path = os.path.join(f'results_eval',
f'{args.dataset}_rs_{args.ref_stride}_nl_{args.neighbor_length}_object_removal')
if not os.path.exists(result_path):
os.makedirs(result_path, exist_ok=True)
if not os.path.exists(result_path):
os.makedirs(result_path)
for index, items in enumerate(test_loader):
torch.cuda.empty_cache()
# frames, masks, video_name, frames_PIL = items
frames, masks, flows_f, flows_b, video_name, frames_PIL = items
video_name = video_name[0]
print('Processing:', video_name)
video_length = frames.size(1)
frames, masks = frames.to(device), masks.to(device)
masked_frames = frames * (1 - masks)
torch.cuda.synchronize()
time_start = time()
with torch.no_grad():
# ---- compute flow ----
if args.load_flow:
gt_flows_bi = (flows_f.to(device), flows_b.to(device))
else:
short_len = 60
if frames.size(1) > short_len:
gt_flows_f_list, gt_flows_b_list = [], []
for f in range(0, video_length, short_len):
end_f = min(video_length, f + short_len)
if f == 0:
flows_f, flows_b = fix_raft(frames[:,f:end_f], iters=args.raft_iter)
else:
flows_f, flows_b = fix_raft(frames[:,f-1:end_f], iters=args.raft_iter)
gt_flows_f_list.append(flows_f)
gt_flows_b_list.append(flows_b)
gt_flows_f = torch.cat(gt_flows_f_list, dim=1)
gt_flows_b = torch.cat(gt_flows_b_list, dim=1)
gt_flows_bi = (gt_flows_f, gt_flows_b)
else:
gt_flows_bi = fix_raft(frames, iters=args.raft_iter)
# ---- complete flow ----
pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, masks)
pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, masks)
# ---- temporal propagation ----
prop_imgs, updated_local_masks = model.img_propagation(masked_frames, pred_flows_bi, masks, 'nearest')
b, t, _, _, _ = masks.size()
updated_masks = updated_local_masks.view(b, t, 1, h, w)
updated_frames = frames * (1-masks) + prop_imgs.view(b, t, 3, h, w) * masks # merge
del gt_flows_bi, frames, updated_local_masks
if not args.load_flow:
torch.cuda.empty_cache()
ori_frames = frames_PIL
ori_frames = [
ori_frames[i].squeeze().cpu().numpy() for i in range(video_length)
]
comp_frames = [None] * video_length
# complete holes by our model
neighbor_stride = args.neighbor_length // 2
for f in range(0, video_length, neighbor_stride):
neighbor_ids = [
i for i in range(max(0, f - neighbor_stride),
min(video_length, f + neighbor_stride + 1))
]
ref_ids = get_ref_index(neighbor_ids, video_length, args.ref_stride)
selected_imgs = updated_frames[:, neighbor_ids + ref_ids, :, :, :]
selected_masks = masks[:, neighbor_ids + ref_ids, :, :, :]
selected_update_masks = updated_masks[:, neighbor_ids + ref_ids, :, :, :]
selected_pred_flows_bi = (pred_flows_bi[0][:, neighbor_ids[:-1], :, :, :], pred_flows_bi[1][:, neighbor_ids[:-1], :, :, :])
with torch.no_grad():
l_t = len(neighbor_ids)
pred_img = model(selected_imgs, selected_pred_flows_bi, selected_masks, selected_update_masks, l_t)
pred_img = pred_img.view(-1, 3, h, w)
pred_img = (pred_img + 1) / 2
pred_img = pred_img.cpu().permute(0, 2, 3, 1).numpy() * 255
binary_masks = masks[0, neighbor_ids, :, :, :].cpu().permute(
0, 2, 3, 1).numpy().astype(np.uint8)
for i in range(len(neighbor_ids)):
idx = neighbor_ids[i]
img = np.array(pred_img[i]).astype(np.uint8) * binary_masks[i] \
+ ori_frames[idx] * (1 - binary_masks[i])
if comp_frames[idx] is None:
comp_frames[idx] = img
else:
comp_frames[idx] = comp_frames[idx].astype(
np.float32) * 0.5 + img.astype(np.float32) * 0.5
torch.cuda.synchronize()
time_i = time() - time_start
time_i = time_i*1.0/video_length
time_all.append(time_i)
if args.task == 'video_completion':
# calculate metrics
cur_video_psnr = []
cur_video_ssim = []
comp_PIL = [] # to calculate VFID
frames_PIL = []
for ori, comp in zip(ori_frames, comp_frames):
psnr, ssim = calc_psnr_and_ssim(ori, comp)
cur_video_psnr.append(psnr)
cur_video_ssim.append(ssim)
total_frame_psnr.append(psnr)
total_frame_ssim.append(ssim)
frames_PIL.append(Image.fromarray(ori.astype(np.uint8)))
comp_PIL.append(Image.fromarray(comp.astype(np.uint8)))
# saving i3d activations
frames_i3d, comp_i3d = calculate_i3d_activations(frames_PIL,
comp_PIL,
i3d_model,
device=device)
real_i3d_activations.append(frames_i3d)
output_i3d_activations.append(comp_i3d)
cur_psnr = sum(cur_video_psnr) / len(cur_video_psnr)
cur_ssim = sum(cur_video_ssim) / len(cur_video_ssim)
avg_psnr = sum(total_frame_psnr) / len(total_frame_psnr)
avg_ssim = sum(total_frame_ssim) / len(total_frame_ssim)
avg_time = sum(time_all) / len(time_all)
print(
f'[{index+1:3}/{len(test_loader)}] Name: {str(video_name):25} | PSNR/SSIM: {cur_psnr:.4f}/{cur_ssim:.4f} \
| Avg PSNR/SSIM: {avg_psnr:.4f}/{avg_ssim:.4f} | Time: {avg_time:.4f}'
)
eval_summary.write(
f'[{index+1:3}/{len(test_loader)}] Name: {str(video_name):25} | PSNR/SSIM: {cur_psnr:.4f}/{cur_ssim:.4f} \
| Avg PSNR/SSIM: {avg_psnr:.4f}/{avg_ssim:.4f} | Time: {avg_time:.4f}\n'
)
else:
avg_time = sum(time_all) / len(time_all)
print(
f'[{index+1:3}/{len(test_loader)}] Name: {str(video_name):25} | Time: {avg_time:.4f}'
)
# saving images for evaluating warpping errors
if args.save_results:
save_frame_path = os.path.join(result_path, video_name)
if not os.path.exists(save_frame_path):
os.makedirs(save_frame_path, exist_ok=False)
for i, frame in enumerate(comp_frames):
cv2.imwrite(
os.path.join(save_frame_path,
str(i).zfill(5) + '.png'),
cv2.cvtColor(frame.astype(np.uint8), cv2.COLOR_RGB2BGR))
if args.task == 'video_completion':
avg_frame_psnr = sum(total_frame_psnr) / len(total_frame_psnr)
avg_frame_ssim = sum(total_frame_ssim) / len(total_frame_ssim)
fid_score = calculate_vfid(real_i3d_activations, output_i3d_activations)
print('Finish evaluation... Average Frame PSNR/SSIM/VFID: '
f'{avg_frame_psnr:.2f}/{avg_frame_ssim:.4f}/{fid_score:.3f} | Time: {avg_time:.4f}')
eval_summary.write(
'Finish evaluation... Average Frame PSNR/SSIM/VFID: '
f'{avg_frame_psnr:.2f}/{avg_frame_ssim:.4f}/{fid_score:.3f} | Time: {avg_time:.4f}')
eval_summary.close()
else:
print('Finish evaluation... Time: {avg_time:.4f}')
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--height', type=int, default=240)
parser.add_argument('--width', type=int, default=432)
parser.add_argument("--ref_stride", type=int, default=10)
parser.add_argument("--neighbor_length", type=int, default=20)
parser.add_argument("--raft_iter", type=int, default=20)
parser.add_argument('--task', default='video_completion', choices=['object_removal', 'video_completion'])
parser.add_argument('--raft_model_path', default='weights/raft-things.pth', type=str)
parser.add_argument('--fc_model_path', default='weights/recurrent_flow_completion.pth', type=str)
parser.add_argument('--propainter_model_path', default='weights/ProPainter.pth', type=str)
parser.add_argument('--dataset', choices=['davis', 'youtube-vos'], type=str)
parser.add_argument('--video_root', default='dataset_root', type=str)
parser.add_argument('--mask_root', default='mask_root', type=str)
parser.add_argument('--flow_root', default='flow_ground_truth_root', type=str)
parser.add_argument('--load_flow', default=False, type=bool)
parser.add_argument('--save_results', action='store_true')
parser.add_argument('--num_workers', default=4, type=int)
args = parser.parse_args()
main_worker(args)