evaluate_flow_completion.py

# -*- coding: utf-8 -*-
import sys 
sys.path.append(".") 

import cv2
import os
import numpy as np
import argparse
from PIL import Image

import torch
from torch.utils.data import DataLoader

from core.dataset import TestDataset
from model.modules.flow_comp_raft import RAFT_bi
from model.recurrent_flow_completion import RecurrentFlowCompleteNet

from RAFT.utils.flow_viz_pt import flow_to_image

import cvbase
import imageio
from time import time

import warnings
warnings.filterwarnings("ignore")

def create_dir(dir):
    """Creates a directory if not exist.
    """
    if not os.path.exists(dir):
        os.makedirs(dir)

def save_flows(output, videoFlowF, videoFlowB):
    # create_dir(os.path.join(output, 'forward_flo'))
    # create_dir(os.path.join(output, 'backward_flo'))
    create_dir(os.path.join(output, 'forward_png'))
    create_dir(os.path.join(output, 'backward_png'))
    N = videoFlowF.shape[-1]
    for i in range(N):
        forward_flow = videoFlowF[..., i]
        backward_flow = videoFlowB[..., i]
        forward_flow_vis = cvbase.flow2rgb(forward_flow)
        backward_flow_vis = cvbase.flow2rgb(backward_flow)
        # cvbase.write_flow(forward_flow, os.path.join(output,  'forward_flo', '{:05d}.flo'.format(i)))
        # cvbase.write_flow(backward_flow, os.path.join(output,  'backward_flo', '{:05d}.flo'.format(i)))
        forward_flow_vis = (forward_flow_vis*255.0).astype(np.uint8)
        backward_flow_vis = (backward_flow_vis*255.0).astype(np.uint8)
        imageio.imwrite(os.path.join(output,  'forward_png', '{:05d}.png'.format(i)), forward_flow_vis)
        imageio.imwrite(os.path.join(output,  'backward_png', '{:05d}.png'.format(i)), backward_flow_vis)

def tensor2np(array):
    array = torch.stack(array, dim=-1).squeeze(0).permute(1, 2, 0, 3).cpu().numpy()
    return array

def main_worker(args):
    # set up datasets and data loader
    args.size = (args.width, args.height)
    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()

    total_frame_epe = []
    time_all = []

    print('Start evaluation...')
    # create results directory
    result_path = os.path.join('results_flow', f'{args.dataset}')
    if not os.path.exists(result_path):
        os.makedirs(result_path)

    eval_summary = open(os.path.join(result_path, f"{args.dataset}_metrics.txt"), "w")

    for index, items in enumerate(test_loader):
        frames, masks, flows_f, flows_b, video_name, frames_PIL = items
        local_masks = masks.float().to(device)

        video_length = frames.size(1)
        
        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=20)

        torch.cuda.synchronize()
        time_start = time()

        # flow_length = flows_f.size(1)
        # f_stride = 30
        # pred_flows_f = []
        # pred_flows_b = []
        # suffix = flow_length%f_stride
        # last = flow_length//f_stride
        # for f in range(0, flow_length, f_stride):
        #     gt_flows_bi_i = (gt_flows_bi[0][:,f:f+f_stride], gt_flows_bi[1][:,f:f+f_stride])
        #     pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi_i, local_masks[:,f:f+f_stride+1])
        #     pred_flows_f_i, pred_flows_b_i = fix_flow_complete.combine_flow(gt_flows_bi_i, pred_flows_bi, local_masks[:,f:f+f_stride+1])
        #     pred_flows_f.append(pred_flows_f_i)
        #     pred_flows_b.append(pred_flows_b_i)
        # pred_flows_f = torch.cat(pred_flows_f, dim=1)
        # pred_flows_b = torch.cat(pred_flows_b, dim=1)
        # pred_flows_bi = (pred_flows_f, pred_flows_b)

        pred_flows_bi, _ = fix_flow_complete.forward_bidirect_flow(gt_flows_bi, local_masks)
        pred_flows_bi = fix_flow_complete.combine_flow(gt_flows_bi, pred_flows_bi, local_masks)

        torch.cuda.synchronize()
        time_i = time() - time_start
        time_i = time_i*1.0/frames.size(1)

        time_all = time_all+[time_i]*frames.size(1)

        cur_video_epe = []
        
        epe1 = torch.mean(torch.sum((flows_f - pred_flows_bi[0].cpu())**2, dim=2).sqrt())
        epe2 = torch.mean(torch.sum((flows_b - pred_flows_bi[1].cpu())**2, dim=2).sqrt())

        cur_video_epe.append(epe1.numpy())
        cur_video_epe.append(epe2.numpy())

        total_frame_epe = total_frame_epe+[epe1.numpy()]*flows_f.size(1)
        total_frame_epe = total_frame_epe+[epe2.numpy()]*flows_f.size(1)

        cur_epe = sum(cur_video_epe) / len(cur_video_epe)
        avg_time = sum(time_all) / len(time_all)
        print(
            f'[{index+1:3}/{len(test_loader)}] Name: {str(video_name):25} | EPE: {cur_epe:.4f} | Time: {avg_time:.4f}'
        )
        eval_summary.write(
            f'[{index+1:3}/{len(test_loader)}] Name: {str(video_name):25} | EPE: {cur_epe:.4f} | Time: {avg_time:.4f}\n'
        )

        # saving images for evaluating warpping errors
        if args.save_results:
            forward_flows = pred_flows_bi[0].cpu().permute(1,0,2,3,4)
            backward_flows = pred_flows_bi[1].cpu().permute(1,0,2,3,4)
            # forward_flows = flows_f.cpu().permute(1,0,2,3,4)
            # backward_flows = flows_b.cpu().permute(1,0,2,3,4)
            videoFlowF = list(forward_flows)
            videoFlowB = list(backward_flows)

            videoFlowF = tensor2np(videoFlowF)
            videoFlowB = tensor2np(videoFlowB)

            save_frame_path = os.path.join(result_path, video_name[0])
            save_flows(save_frame_path, videoFlowF, videoFlowB)

    avg_frame_epe = sum(total_frame_epe) / len(total_frame_epe)

    print(f'Finish evaluation... Average Frame EPE: {avg_frame_epe:.4f} | | Time: {avg_time:.4f}')
    eval_summary.write(f'Finish evaluation... Average Frame EPE: {avg_frame_epe:.4f} | | Time: {avg_time:.4f}\n')
    eval_summary.close()
    
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('--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('--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("--raft_iter", type=int, default=20)
    parser.add_argument('--save_results', action='store_true')
    parser.add_argument('--num_workers', default=4, type=int)
    args = parser.parse_args()
    main_worker(args)