scripts/evaluate_propainter.py

# -*- 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)