mirror of
https://github.com/sczhou/ProPainter.git
synced 2026-10-21 06:09:53 -05:00
281 lines
12 KiB
Python
281 lines
12 KiB
Python
# -*- 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)
|