# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
# --------------------------------------------------------
# Main test function
# --------------------------------------------------------
import os
import argparse
import pickle
from PIL import Image
import numpy as np
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader
import utils.misc as misc
from models.croco_downstream import CroCoDownstreamBinocular
from models.head_downstream import PixelwiseTaskWithDPT
from stereoflow.criterion import *
from stereoflow.datasets_stereo import get_test_datasets_stereo
from stereoflow.datasets_flow import get_test_datasets_flow
from stereoflow.engine import tiled_pred
from stereoflow.datasets_stereo import vis_disparity
from stereoflow.datasets_flow import flowToColor
def get_args_parser():
parser = argparse.ArgumentParser('Test CroCo models on stereo/flow', add_help=False)
# important argument
parser.add_argument('--model', required=True, type=str, help='Path to the model to evaluate')
parser.add_argument('--dataset', required=True, type=str, help="test dataset (there can be multiple dataset separated by a +)")
# tiling
parser.add_argument('--tile_conf_mode', type=str, default='', help='Weights for the tiling aggregation based on confidence (empty means use the formula from the loaded checkpoint')
parser.add_argument('--tile_overlap', type=float, default=0.7, help='overlap between tiles')
# save (it will automatically go to <model_path>_<dataset_str>/<tile_str>_<save>)
parser.add_argument('--save', type=str, nargs='+', default=[],
help='what to save: \
metrics (pickle file), \
pred (raw prediction save as torch tensor), \
visu (visualization in png of each prediction), \
err10 (visualization in png of the error clamp at 10 for each prediction), \
submission (submission file)')
# other (no impact)
parser.add_argument('--num_workers', default=4, type=int)
return parser
def _load_model_and_criterion(model_path, do_load_metrics, device):
print('loading model from', model_path)
assert os.path.isfile(model_path)
ckpt = torch.load(model_path, 'cpu')
ckpt_args = ckpt['args']
task = ckpt_args.task
tile_conf_mode = ckpt_args.tile_conf_mode
num_channels = {'stereo': 1, 'flow': 2}[task]
with_conf = eval(ckpt_args.criterion).with_conf
if with_conf: num_channels += 1
print('head: PixelwiseTaskWithDPT()')
head = PixelwiseTaskWithDPT()
head.num_channels = num_channels
print('croco_args:', ckpt_args.croco_args)
model = CroCoDownstreamBinocular(head, **ckpt_args.croco_args)
msg = model.load_state_dict(ckpt['model'], strict=True)
model.eval()
model = model.to(device)
if do_load_metrics:
if task=='stereo':
metrics = StereoDatasetMetrics().to(device)
else:
metrics = FlowDatasetMetrics().to(device)
else:
metrics = None
return model, metrics, ckpt_args.crop, with_conf, task, tile_conf_mode
def _save_batch(pred, gt, pairnames, dataset, task, save, outdir, time, submission_dir=None):
for i in range(len(pairnames)):
pairname = eval(pairnames[i]) if pairnames[i].startswith('(') else pairnames[i] # unbatch pairname
fname = os.path.join(outdir, dataset.pairname_to_str(pairname))
os.makedirs(os.path.dirname(fname), exist_ok=True)
predi = pred[i,...]
if gt is not None: gti = gt[i,...]
if 'pred' in save:
torch.save(predi.squeeze(0).cpu(), fname+'_pred.pth')
if 'visu' in save:
if task=='stereo':
disparity = predi.permute((1,2,0)).squeeze(2).cpu().numpy()
m,M = None
if gt is not None:
mask = torch.isfinite(gti)
m = gt[mask].min()
M = gt[mask].max()
img_disparity = vis_disparity(disparity, m=m, M=M)
Image.fromarray(img_disparity).save(fname+'_pred.png')
else:
# normalize flowToColor according to the maxnorm of gt (or prediction if not available)
flowNorm = torch.sqrt(torch.sum( (gti if gt is not None else predi)**2, dim=0)).max().item()
imgflow = flowToColor(predi.permute((1,2,0)).cpu().numpy(), maxflow=flowNorm)
Image.fromarray(imgflow).save(fname+'_pred.png')
if 'err10' in save:
assert gt is not None
L2err = torch.sqrt(torch.sum( (gti-predi)**2, dim=0))
valid = torch.isfinite(gti[0,:,:])
L2err[~valid] = 0.0
L2err = torch.clamp(L2err, max=10.0)
red = (L2err*255.0/10.0).to(dtype=torch.uint8)[:,:,None]
zer = torch.zeros_like(red)
imgerr = torch.cat( (red,zer,zer), dim=2).cpu().numpy()
Image.fromarray(imgerr).save(fname+'_err10.png')
if 'submission' in save:
assert submission_dir is not None
predi_np = predi.permute(1,2,0).squeeze(2).cpu().numpy() # transform into HxWx2 for flow or HxW for stereo
dataset.submission_save_pairname(pairname, predi_np, submission_dir, time)
def main(args):
# load the pretrained model and metrics
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
model, metrics, cropsize, with_conf, task, tile_conf_mode = _load_model_and_criterion(args.model, 'metrics' in args.save, device)
if args.tile_conf_mode=='': args.tile_conf_mode = tile_conf_mode
# load the datasets
datasets = (get_test_datasets_stereo if task=='stereo' else get_test_datasets_flow)(args.dataset)
dataloaders = [DataLoader(dataset, batch_size=1, shuffle=False, num_workers=args.num_workers, pin_memory=True, drop_last=False) for dataset in datasets]
# run
for i,dataloader in enumerate(dataloaders):
dataset = datasets[i]
dstr = args.dataset.split('+')[i]
outdir = args.model+'_'+misc.filename(dstr)
if 'metrics' in args.save and len(args.save)==1:
fname = os.path.join(outdir, f'conf_{args.tile_conf_mode}_overlap_{args.tile_overlap}.pkl')
if os.path.isfile(fname) and len(args.save)==1:
print(' metrics already compute in '+fname)
with open(fname, 'rb') as fid:
results = pickle.load(fid)
for k,v in results.items():
print('{:s}: {:.3f}'.format(k, v))
continue
if 'submission' in args.save:
dirname = f'submission_conf_{args.tile_conf_mode}_overlap_{args.tile_overlap}'
submission_dir = os.path.join(outdir, dirname)
else:
submission_dir = None
print('')
print('saving {:s} in {:s}'.format('+'.join(args.save), outdir))
print(repr(dataset))
if metrics is not None:
metrics.reset()
for data_iter_step, (image1, image2, gt, pairnames) in enumerate(tqdm(dataloader)):
do_flip = (task=='stereo' and dstr.startswith('Spring') and any("right" in p for p in pairnames)) # we flip the images and will flip the prediction after as we assume img1 is on the left
image1 = image1.to(device, non_blocking=True)
image2 = image2.to(device, non_blocking=True)
gt = gt.to(device, non_blocking=True) if gt.numel()>0 else None # special case for test time
if do_flip:
assert all("right" in p for p in pairnames)
image1 = image1.flip(dims=[3]) # this is already the right frame, let's flip it
image2 = image2.flip(dims=[3])
gt = gt # that is ok
with torch.inference_mode():
pred, _, _, time = tiled_pred(model, None, image1, image2, None if dataset.name=='Spring' else gt, conf_mode=args.tile_conf_mode, overlap=args.tile_overlap, crop=cropsize, with_conf=with_conf, return_time=True)
if do_flip:
pred = pred.flip(dims=[3])
if metrics is not None:
metrics.add_batch(pred, gt)
if any(k in args.save for k in ['pred','visu','err10','submission']):
_save_batch(pred, gt, pairnames, dataset, task, args.save, outdir, time, submission_dir=submission_dir)
# print
if metrics is not None:
results = metrics.get_results()
for k,v in results.items():
print('{:s}: {:.3f}'.format(k, v))
# save if needed
if 'metrics' in args.save:
os.makedirs(os.path.dirname(fname), exist_ok=True)
with open(fname, 'wb') as fid:
pickle.dump(results, fid)
print('metrics saved in', fname)
# finalize submission if needed
if 'submission' in args.save:
dataset.finalize_submission(submission_dir)
if __name__ == '__main__':
args = get_args_parser()
args = args.parse_args()
main(args)