vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/duster/croco/stereoflow/engine.py

@@ -0,0 +1,280 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+# --------------------------------------------------------
+# Main function for training one epoch or testing
+# --------------------------------------------------------
+
+import math
+import sys
+from typing import Iterable
+import numpy as np
+import torch
+import torchvision
+
+from utils import misc as misc
+
+
+def split_prediction_conf(predictions, with_conf=False):
+    if not with_conf:
+        return predictions, None
+    conf = predictions[:,-1:,:,:]
+    predictions = predictions[:,:-1,:,:]
+    return predictions, conf
+
+def train_one_epoch(model: torch.nn.Module, criterion: torch.nn.Module, metrics: torch.nn.Module,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    log_writer=None, print_freq = 20,
+                    args=None):
+    model.train(True)
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    details = {}
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    if args.img_per_epoch:
+        iter_per_epoch = args.img_per_epoch // args.batch_size + int(args.img_per_epoch % args.batch_size > 0)
+        assert len(data_loader) >= iter_per_epoch, 'Dataset is too small for so many iterations'
+        len_data_loader = iter_per_epoch
+    else:
+        len_data_loader, iter_per_epoch = len(data_loader), None
+
+    for data_iter_step, (image1, image2, gt, pairname) in enumerate(metric_logger.log_every(data_loader, print_freq, header, max_iter=iter_per_epoch)):
+        
+        image1 = image1.to(device, non_blocking=True)
+        image2 = image2.to(device, non_blocking=True)
+        gt = gt.to(device, non_blocking=True)
+        
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            misc.adjust_learning_rate(optimizer, data_iter_step / len_data_loader + epoch, args)
+
+        with torch.cuda.amp.autocast(enabled=bool(args.amp)):
+            prediction = model(image1, image2)
+            prediction, conf = split_prediction_conf(prediction, criterion.with_conf)
+            batch_metrics = metrics(prediction.detach(), gt)
+            loss = criterion(prediction, gt) if conf is None else criterion(prediction, gt, conf)
+            
+        loss_value = loss.item()
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+        
+        metric_logger.update(loss=loss_value)
+        for k,v in batch_metrics.items():
+            metric_logger.update(**{k: v.item()})
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        #if args.dsitributed: loss_value_reduce = misc.all_reduce_mean(loss_value)
+        time_to_log = ((data_iter_step + 1) % (args.tboard_log_step * accum_iter) == 0 or data_iter_step == len_data_loader-1)
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        if log_writer is not None and time_to_log:
+            epoch_1000x = int((data_iter_step / len_data_loader + epoch) * 1000)
+            # We use epoch_1000x as the x-axis in tensorboard. This calibrates different curves when batch size changes.
+            log_writer.add_scalar('train/loss', loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+            for k,v in batch_metrics.items():
+                log_writer.add_scalar('train/'+k, v.item(), epoch_1000x)
+
+    # gather the stats from all processes
+    #if args.distributed: metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def validate_one_epoch(model: torch.nn.Module,
+                   criterion: torch.nn.Module,
+                   metrics: torch.nn.Module,
+                   data_loaders: list[Iterable],
+                   device: torch.device,
+                   epoch: int,
+                   log_writer=None,
+                   args=None):
+
+    model.eval()
+    metric_loggers = []
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 20
+
+    conf_mode = args.tile_conf_mode
+    crop = args.crop
+    
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+
+    results = {}
+    dnames = []
+    image1, image2, gt, prediction = None, None, None, None
+    for didx, data_loader in enumerate(data_loaders):
+        dname = str(data_loader.dataset)
+        dnames.append(dname)
+        metric_loggers.append(misc.MetricLogger(delimiter="  "))
+        for data_iter_step, (image1, image2, gt, pairname) in enumerate(metric_loggers[didx].log_every(data_loader, print_freq, header)):
+            image1 = image1.to(device, non_blocking=True)
+            image2 = image2.to(device, non_blocking=True)
+            gt = gt.to(device, non_blocking=True)
+            if dname.startswith('Spring'):
+                assert gt.size(2)==image1.size(2)*2 and gt.size(3)==image1.size(3)*2
+                gt = (gt[:,:,0::2,0::2] + gt[:,:,0::2,1::2] + gt[:,:,1::2,0::2] + gt[:,:,1::2,1::2] ) / 4.0 # we approximate the gt based on the 2x upsampled ones
+
+            with torch.inference_mode():
+                prediction, tiled_loss, c = tiled_pred(model, criterion, image1, image2, gt, conf_mode=conf_mode, overlap=args.val_overlap, crop=crop, with_conf=criterion.with_conf)
+                batch_metrics = metrics(prediction.detach(), gt)
+                loss = criterion(prediction.detach(), gt) if not criterion.with_conf else criterion(prediction.detach(), gt, c)
+                loss_value = loss.item()
+                metric_loggers[didx].update(loss_tiled=tiled_loss.item())
+                metric_loggers[didx].update(**{f'loss': loss_value})
+                for k,v in batch_metrics.items():
+                    metric_loggers[didx].update(**{dname+'_' + k: v.item()})
+        
+    results = {k: meter.global_avg for ml in metric_loggers for k, meter in ml.meters.items()}
+    if len(dnames)>1:
+        for k in batch_metrics.keys():
+            results['AVG_'+k] = sum(results[dname+'_'+k] for dname in dnames) / len(dnames)
+            
+    if log_writer is not None :
+        epoch_1000x = int((1 + epoch) * 1000)
+        for k,v in results.items():
+            log_writer.add_scalar('val/'+k, v, epoch_1000x)
+
+    print("Averaged stats:", results)
+    return results
+
+import torch.nn.functional as F
+def _resize_img(img, new_size):
+    return F.interpolate(img, size=new_size, mode='bicubic', align_corners=False)
+def _resize_stereo_or_flow(data, new_size):
+    assert data.ndim==4
+    assert data.size(1) in [1,2]
+    scale_x = new_size[1]/float(data.size(3))
+    out = F.interpolate(data, size=new_size, mode='bicubic', align_corners=False)
+    out[:,0,:,:] *= scale_x
+    if out.size(1)==2:
+        scale_y = new_size[0]/float(data.size(2))        
+        out[:,1,:,:] *= scale_y
+        print(scale_x, new_size, data.shape)
+    return out
+    
+
+@torch.no_grad()
+def tiled_pred(model, criterion, img1, img2, gt,
+               overlap=0.5, bad_crop_thr=0.05,
+               downscale=False, crop=512, ret='loss',
+               conf_mode='conf_expsigmoid_10_5', with_conf=False, 
+               return_time=False):
+                     
+    # for each image, we are going to run inference on many overlapping patches
+    # then, all predictions will be weighted-averaged
+    if gt is not None:
+        B, C, H, W = gt.shape
+    else:
+        B, _, H, W = img1.shape
+        C = model.head.num_channels-int(with_conf)
+    win_height, win_width = crop[0], crop[1]
+    
+    # upscale to be larger than the crop
+    do_change_scale =  H<win_height or W<win_width
+    if do_change_scale: 
+        upscale_factor = max(win_width/W, win_height/W)
+        original_size = (H,W)
+        new_size = (round(H*upscale_factor),round(W*upscale_factor))
+        img1 = _resize_img(img1, new_size)
+        img2 = _resize_img(img2, new_size)
+        # resize gt just for the computation of tiled losses
+        if gt is not None: gt = _resize_stereo_or_flow(gt, new_size)
+        H,W = img1.shape[2:4]
+        
+    if conf_mode.startswith('conf_expsigmoid_'): # conf_expsigmoid_30_10
+        beta, betasigmoid = map(float, conf_mode[len('conf_expsigmoid_'):].split('_'))
+    elif conf_mode.startswith('conf_expbeta'): # conf_expbeta3
+        beta = float(conf_mode[len('conf_expbeta'):])
+    else:
+        raise NotImplementedError(f"conf_mode {conf_mode} is not implemented")
+
+    def crop_generator():
+        for sy in _overlapping(H, win_height, overlap):
+          for sx in _overlapping(W, win_width, overlap):
+            yield sy, sx, sy, sx, True
+
+    # keep track of weighted sum of prediction*weights and weights
+    accu_pred = img1.new_zeros((B, C, H, W)) # accumulate the weighted sum of predictions 
+    accu_conf = img1.new_zeros((B, H, W)) + 1e-16 # accumulate the weights 
+    accu_c = img1.new_zeros((B, H, W)) # accumulate the weighted sum of confidences ; not so useful except for computing some losses
+
+    tiled_losses = []
+    
+    if return_time:
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+
+    for sy1, sx1, sy2, sx2, aligned in crop_generator():
+        # compute optical flow there
+        pred =  model(_crop(img1,sy1,sx1), _crop(img2,sy2,sx2))
+        pred, predconf = split_prediction_conf(pred, with_conf=with_conf)
+        
+        if gt is not None: gtcrop = _crop(gt,sy1,sx1)
+        if criterion is not None and gt is not None: 
+            tiled_losses.append( criterion(pred, gtcrop).item() if predconf is None else criterion(pred, gtcrop, predconf).item() )
+        
+        if conf_mode.startswith('conf_expsigmoid_'):
+            conf = torch.exp(- beta * 2 * (torch.sigmoid(predconf / betasigmoid) - 0.5)).view(B,win_height,win_width)
+        elif conf_mode.startswith('conf_expbeta'):
+            conf = torch.exp(- beta * predconf).view(B,win_height,win_width)
+        else:
+            raise NotImplementedError
+                        
+        accu_pred[...,sy1,sx1] += pred * conf[:,None,:,:]
+        accu_conf[...,sy1,sx1] += conf
+        accu_c[...,sy1,sx1] += predconf.view(B,win_height,win_width) * conf 
+        
+    pred = accu_pred / accu_conf[:, None,:,:]
+    c = accu_c / accu_conf
+    assert not torch.any(torch.isnan(pred))
+
+    if return_time:
+        end.record()
+        torch.cuda.synchronize()
+        time = start.elapsed_time(end)/1000.0 # this was in milliseconds
+
+    if do_change_scale:
+        pred = _resize_stereo_or_flow(pred, original_size)
+    
+    if return_time:
+        return pred, torch.mean(torch.tensor(tiled_losses)), c, time
+    return pred, torch.mean(torch.tensor(tiled_losses)), c
+
+
+def _overlapping(total, window, overlap=0.5):
+    assert total >= window and 0 <= overlap < 1, (total, window, overlap)
+    num_windows = 1 + int(np.ceil( (total - window) / ((1-overlap) * window) ))
+    offsets = np.linspace(0, total-window, num_windows).round().astype(int)
+    yield from (slice(x, x+window) for x in offsets)
+
+def _crop(img, sy, sx):
+    B, THREE, H, W = img.shape
+    if 0 <= sy.start and sy.stop <= H and 0 <= sx.start and sx.stop <= W:
+        return img[:,:,sy,sx]
+    l, r = max(0,-sx.start), max(0,sx.stop-W)
+    t, b = max(0,-sy.start), max(0,sy.stop-H)
+    img = torch.nn.functional.pad(img, (l,r,t,b), mode='constant')
+    return img[:, :, slice(sy.start+t,sy.stop+t), slice(sx.start+l,sx.stop+l)]

vismatch/third_party/duster/croco/stereoflow/test.py

@@ -0,0 +1,216 @@
+# 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)