import sys sys.path.append('core') from PIL import Image import argparse import os import time import numpy as np import torch import torch.nn.functional as F import matplotlib.pyplot as plt import os.path as osp import datasets from utils import flow_viz from utils import frame_utils from raft import RAFT from utils.utils import InputPadder, forward_interpolate import cv2 @torch.no_grad() def create_sintel_submission(model, iters=32, warm_start=False, output_path='sintel_submission'): """ Create submission for the Sintel leaderboard """ model.eval() for dstype in ['clean', 'final']: test_dataset = datasets.MpiSintel(split='test', aug_params=None, dstype=dstype) flow_prev, sequence_prev = None, None for test_id in range(len(test_dataset)): image1, image2, (sequence, frame) = test_dataset[test_id] if sequence != sequence_prev: flow_prev = None padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda()) flow_low, flow_pr = model(image1, image2, iters=iters, flow_init=flow_prev, test_mode=True) flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() if warm_start: flow_prev = forward_interpolate(flow_low[0])[None].cuda() output_dir = os.path.join(output_path, dstype, sequence) output_file = os.path.join(output_dir, 'frame%04d.flo' % (frame+1)) if not os.path.exists(output_dir): os.makedirs(output_dir) frame_utils.writeFlow(output_file, flow) sequence_prev = sequence @torch.no_grad() def create_kitti_submission(model, iters=24, output_path='kitti_submission'): """ Create submission for the Sintel leaderboard """ model.eval() test_dataset = datasets.KITTI(split='testing', aug_params=None) if not os.path.exists(output_path): os.makedirs(output_path) for test_id in range(len(test_dataset)): image1, image2, (frame_id, ) = test_dataset[test_id] padder = InputPadder(image1.shape, mode='kitti') image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda()) _, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() output_filename = os.path.join(output_path, frame_id) frame_utils.writeFlowKITTI(output_filename, flow) @torch.no_grad() def validate_chairs(model, iters=24): """ Perform evaluation on the FlyingChairs (test) split """ model.eval() epe_list = [] val_dataset = datasets.FlyingChairs(split='validation') for val_id in range(len(val_dataset)): image1, image2, flow_gt, _ = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() _, flow_pr = model(image1, image2, iters=iters, test_mode=True) epe = torch.sum((flow_pr[0].cpu() - flow_gt)**2, dim=0).sqrt() epe_list.append(epe.view(-1).numpy()) epe = np.mean(np.concatenate(epe_list)) print("Validation Chairs EPE: %f" % epe) return {'chairs': epe} @torch.no_grad() def validate_sintel(model, iters=32): """ Peform validation using the Sintel (train) split """ model.eval() results = {} for dstype in ['clean', 'final']: val_dataset = datasets.MpiSintel(split='training', dstype=dstype) epe_list = [] for val_id in range(len(val_dataset)): image1, image2, flow_gt, _ = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() epe_list.append(epe.view(-1).numpy()) epe_all = np.concatenate(epe_list) epe = np.mean(epe_all) px1 = np.mean(epe_all<1) px3 = np.mean(epe_all<3) px5 = np.mean(epe_all<5) print("Validation (%s) EPE: %f, 1px: %f, 3px: %f, 5px: %f" % (dstype, epe, px1, px3, px5)) results[dstype] = np.mean(epe_list) return results @torch.no_grad() def validate_kitti(model, iters=24): """ Peform validation using the KITTI-2015 (train) split """ model.eval() val_dataset = datasets.KITTI(split='training') out_list, epe_list = [], [] for val_id in range(len(val_dataset)): image1, image2, flow_gt, valid_gt = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape, mode='kitti') image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() mag = torch.sum(flow_gt**2, dim=0).sqrt() epe = epe.view(-1) mag = mag.view(-1) val = valid_gt.view(-1) >= 0.5 out = ((epe > 3.0) & ((epe/mag) > 0.05)).float() epe_list.append(epe[val].mean().item()) out_list.append(out[val].cpu().numpy()) epe_list = np.array(epe_list) out_list = np.concatenate(out_list) epe = np.mean(epe_list) f1 = 100 * np.mean(out_list) print("Validation KITTI: %f, %f" % (epe, f1)) return {'kitti-epe': epe, 'kitti-f1': f1} @torch.no_grad() def validate_anime(model, path=None, iters=32): """ Peform validation using the Sintel (train) split """ model.eval() results = {} for dstype in ['Frame_Anime']: val_dataset = datasets.AnimeRun(split='test', dstype=dstype) epe_list = [] epe_flat_list = [] epe_line_list = [] epe_nonocc_list = [] epe_occ_list = [] epe_ds10_list = [] epe_ds1050_list = [] epe_ds50_list = [] for val_id in range(len(val_dataset)): # if val_id % 100 == 0: # print('Testing the {}-th image pair...'.format(val_id)) image1, image2, flow_gt, occ, flat, extra_info, valid = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() mag = torch.sum(flow_gt ** 2, dim=0).sqrt() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() flo_gt_ = flow_gt.clone() valid = valid.bool() valid[:] = True # print((~valid).size(), flo_gt_[0, :, :].size()) flo_gt_[0, :, :][~valid] = 0 flo_gt_[1, :, :][~valid] = 0 if epe[valid].numpy().mean() > 50: print(extra_info[0], extra_info[1], epe[valid].numpy().mean(), flow_gt.abs().max(), (flow_gt.abs() < 0.01).sum()) if path is not None: flo_color = flow_viz.flow_to_image(flow.permute(1, 2, 0).data.numpy(), convert_to_bgr=True) flo_color = cv2.putText(flo_color, 'EPE = ' + str(np.round(epe[valid].numpy().mean(), 2)), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, fontScale=2, color=(0, 0, 255), thickness=3) flo_gt_color = flow_viz.flow_to_image(flo_gt_.permute(1, 2, 0).data.numpy(), convert_to_bgr=True) cv2.imwrite(osp.join(path, str(extra_info[0] + '_' + str(extra_info[1]) + '.jpg')), flo_color) cv2.imwrite(osp.join(path, str(extra_info[0] + '_' + str(extra_info[1]) + '_gt.jpg')), flo_gt_color) epe_nonocc = epe[(occ == 1) & valid] epe_flat = epe[(flat > 0) & valid ] epe_line = epe[(flat == 0) & valid ] epe_occ = epe[(occ == 0) & valid] epe_ds10 = epe[(mag <= 10.0) & valid] epe_ds1050 = epe[(mag > 10.0) & (mag <= 50.0) & valid] epe_ds50 = epe[(mag > 50.0) & valid] epe_list.append(epe[valid].view(-1).numpy()) epe_flat_list.append(epe_flat.view(-1).numpy()) epe_line_list.append(epe_line.view(-1).numpy()) epe_nonocc_list.append(epe_nonocc.view(-1).numpy()) epe_occ_list.append(epe_occ.view(-1).numpy()) epe_ds10_list.append(epe_ds10.view(-1).numpy()) epe_ds1050_list.append(epe_ds1050.view(-1).numpy()) epe_ds50_list.append(epe_ds50.view(-1).numpy()) epe_all = np.concatenate(epe_list) epe_flat_all = np.concatenate(epe_flat_list) epe_line_all = np.concatenate(epe_line_list) epe_nonocc_all = np.concatenate(epe_nonocc_list) epe_occ_all = np.concatenate(epe_occ_list) epe_ds10_all = np.concatenate(epe_ds10_list) epe_ds1050_all = np.concatenate(epe_ds1050_list) epe_ds50_all = np.concatenate(epe_ds50_list) epe = np.mean(epe_all) px1 = np.mean(epe_all<1) px3 = np.mean(epe_all<3) px5 = np.mean(epe_all<5) # print("Validation (%s) EPE: %f, 1px: %f, 3px: %f, 5px: %f" % (dstype, epe, px1, px3, px5)) print("Test (%s) EPE: %f, Occ: %f, Non-Occ: %f, Line: %f, Flat: %f, s<10: %f, s10-50: %f, s>50: %f" % (dstype, epe, np.mean(epe_occ_all), np.mean(epe_nonocc_all), np.mean(epe_line_all), np.mean(epe_flat_all), np.mean(epe_ds10_all), np.mean(epe_ds1050_all), np.mean(epe_ds50_all))) results[dstype] = epe return results if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--model', help="restore checkpoint") parser.add_argument('--dataset', help="dataset for evaluation") parser.add_argument('--small', action='store_true', help='use small model') parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation') args = parser.parse_args() model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model.cuda() model.eval() # create_sintel_submission(model.module, warm_start=True) # create_kitti_submission(model.module) print('Testing {}'.format(args.model)) if not osp.exists('visualize/' + args.model.replace('/', '_')): os.mkdir('visualize/' + args.model.replace('/', '_')) with torch.no_grad(): if args.dataset == 'chairs': validate_chairs(model.module) elif args.dataset == 'sintel': validate_sintel(model.module) elif args.dataset == 'anime': validate_anime(model.module, path='visualize/' + args.model.replace('/', '_')) elif args.dataset == 'kitti': validate_kitti(model.module)