cascade_evaluation.py

import os
import warnings

import torch
import torchvision
import torchvision.prototype.models.depth.stereo
import utils
from torch.nn import functional as F
from train import make_eval_loader

from utils.metrics import AVAILABLE_METRICS
from visualization import make_prediction_image_side_to_side


def get_args_parser(add_help=True):
    import argparse

    parser = argparse.ArgumentParser(description="PyTorch Stereo Matching Evaluation", add_help=add_help)
    parser.add_argument("--dataset", type=str, default="middlebury2014-train", help="dataset to use")
    parser.add_argument("--dataset-root", type=str, default="", help="root of the dataset")

    parser.add_argument("--checkpoint", type=str, default="", help="path to weights")
    parser.add_argument("--weights", type=str, default=None, help="torchvision API weight")
    parser.add_argument(
        "--model",
        type=str,
        default="crestereo_base",
        help="which model to use if not speciffying a training checkpoint",
    )
    parser.add_argument("--img-folder", type=str, default="images")

    parser.add_argument("--batch-size", type=int, default=1, help="batch size")
    parser.add_argument("--workers", type=int, default=0, help="number of workers")

    parser.add_argument("--eval-size", type=int, nargs="+", default=[384, 512], help="resize size")
    parser.add_argument(
        "--norm-mean", type=float, nargs="+", default=[0.5, 0.5, 0.5], help="mean for image normalization"
    )
    parser.add_argument(
        "--norm-std", type=float, nargs="+", default=[0.5, 0.5, 0.5], help="std for image normalization"
    )
    parser.add_argument(
        "--use-grayscale", action="store_true", help="use grayscale images instead of RGB", default=False
    )
    parser.add_argument("--max-disparity", type=float, default=None, help="maximum disparity")
    parser.add_argument(
        "--interpolation-strategy",
        type=str,
        default="bilinear",
        help="interpolation strategy",
        choices=["bilinear", "bicubic", "mixed"],
    )

    parser.add_argument("--n_iterations", nargs="+", type=int, default=[10], help="number of recurent iterations")
    parser.add_argument("--n_cascades", nargs="+", type=int, default=[1], help="number of cascades")
    parser.add_argument(
        "--metrics",
        type=str,
        nargs="+",
        default=["mae", "rmse", "1px", "3px", "5px", "relepe"],
        help="metrics to log",
        choices=AVAILABLE_METRICS,
    )
    parser.add_argument("--mixed-precision", action="store_true", help="use mixed precision training")

    parser.add_argument("--world-size", type=int, default=1, help="number of distributed processes")
    parser.add_argument("--dist-url", type=str, default="env://", help="url used to set up distributed training")
    parser.add_argument("--device", type=str, default="cuda", help="device to use for training")

    parser.add_argument("--save-images", action="store_true", help="save images of the predictions")
    parser.add_argument("--padder-type", type=str, default="kitti", help="padder type", choices=["kitti", "sintel"])

    return parser


def cascade_inference(model, image_left, image_right, iterations, cascades):
    # check that image size is divisible by 16 * (2 ** (cascades - 1))
    for image in [image_left, image_right]:
        if image.shape[-2] % ((2 ** (cascades - 1))) != 0:
            raise ValueError(
                f"image height is not divisible by {16 * (2 ** (cascades - 1))}. Image shape: {image.shape[-2]}"
            )

        if image.shape[-1] % ((2 ** (cascades - 1))) != 0:
            raise ValueError(
                f"image width is not divisible by {16 * (2 ** (cascades - 1))}. Image shape: {image.shape[-2]}"
            )

    left_image_pyramid = [image_left]
    right_image_pyramid = [image_right]
    for idx in range(0, cascades - 1):
        ds_factor = int(2 ** (idx + 1))
        ds_shape = (image_left.shape[-2] // ds_factor, image_left.shape[-1] // ds_factor)
        left_image_pyramid += F.interpolate(image_left, size=ds_shape, mode="bilinear", align_corners=True).unsqueeze(0)
        right_image_pyramid += F.interpolate(image_right, size=ds_shape, mode="bilinear", align_corners=True).unsqueeze(
            0
        )

    flow_init = None
    for left_image, right_image in zip(reversed(left_image_pyramid), reversed(right_image_pyramid)):
        flow_pred = model(left_image, right_image, flow_init, num_iters=iterations)
        # flow pred is a list
        flow_init = flow_pred[-1]

    return flow_init


@torch.inference_mode()
def _evaluate(
    model,
    args,
    val_loader,
    *,
    padder_mode,
    print_freq=10,
    writer=None,
    step=None,
    iterations=10,
    cascades=1,
    batch_size=None,
    header=None,
    save_images=False,
    save_path="",
):
    """Helper function to compute various metrics (epe, etc.) for a model on a given dataset.
    We process as many samples as possible with ddp.
    """
    model.eval()
    header = header or "Test:"
    device = torch.device(args.device)
    metric_logger = utils.MetricLogger(delimiter="  ")

    iterations = iterations or args.recurrent_updates

    logger = utils.MetricLogger()
    for meter_name in args.metrics:
        logger.add_meter(meter_name, fmt="{global_avg:.4f}")
    if "fl-all" not in args.metrics:
        logger.add_meter("fl-all", fmt="{global_avg:.4f}")

    num_processed_samples = 0
    with torch.cuda.amp.autocast(enabled=args.mixed_precision, dtype=torch.float16):
        batch_idx = 0
        for blob in metric_logger.log_every(val_loader, print_freq, header):
            image_left, image_right, disp_gt, valid_disp_mask = (x.to(device) for x in blob)
            padder = utils.InputPadder(image_left.shape, mode=padder_mode)
            image_left, image_right = padder.pad(image_left, image_right)

            disp_pred = cascade_inference(model, image_left, image_right, iterations, cascades)
            disp_pred = disp_pred[:, :1, :, :]
            disp_pred = padder.unpad(disp_pred)

            if save_images:
                if args.distributed:
                    rank_prefix = args.rank
                else:
                    rank_prefix = 0
                make_prediction_image_side_to_side(
                    disp_pred, disp_gt, valid_disp_mask, save_path, prefix=f"batch_{rank_prefix}_{batch_idx}"
                )

            metrics, _ = utils.compute_metrics(disp_pred, disp_gt, valid_disp_mask, metrics=logger.meters.keys())
            num_processed_samples += image_left.shape[0]
            for name in metrics:
                logger.meters[name].update(metrics[name], n=1)

            batch_idx += 1

    num_processed_samples = utils.reduce_across_processes(num_processed_samples) / args.world_size

    print("Num_processed_samples: ", num_processed_samples)
    if (
        hasattr(val_loader.dataset, "__len__")
        and len(val_loader.dataset) != num_processed_samples
        and torch.distributed.get_rank() == 0
    ):
        warnings.warn(
            f"Number of processed samples {num_processed_samples} is different"
            f"from the dataset size {len(val_loader.dataset)}. This may happen if"
            "the dataset is not divisible by the batch size. Try lowering the batch size for more accurate results."
        )

    if writer is not None and args.rank == 0:
        for meter_name, meter_value in logger.meters.items():
            scalar_name = f"{meter_name} {header}"
            writer.add_scalar(scalar_name, meter_value.avg, step)

    logger.synchronize_between_processes()
    print(header, logger)

    logger_metrics = {k: v.global_avg for k, v in logger.meters.items()}
    return logger_metrics


def evaluate(model, loader, args, writer=None, step=None):
    os.makedirs(args.img_folder, exist_ok=True)
    checkpoint_name = os.path.basename(args.checkpoint) or args.weights
    image_checkpoint_folder = os.path.join(args.img_folder, checkpoint_name)

    metrics = {}
    base_image_folder = os.path.join(image_checkpoint_folder, args.dataset)
    os.makedirs(base_image_folder, exist_ok=True)

    for n_cascades in args.n_cascades:
        for n_iters in args.n_iterations:

            config = f"{n_cascades}c_{n_iters}i"
            config_image_folder = os.path.join(base_image_folder, config)
            os.makedirs(config_image_folder, exist_ok=True)

            metrics[config] = _evaluate(
                model,
                args,
                loader,
                padder_mode=args.padder_type,
                header=f"{args.dataset} evaluation@ size:{args.eval_size} n_cascades:{n_cascades} n_iters:{n_iters}",
                batch_size=args.batch_size,
                writer=writer,
                step=step,
                iterations=n_iters,
                cascades=n_cascades,
                save_path=config_image_folder,
                save_images=args.save_images,
            )

    metric_log = []
    metric_log_dict = {}
    # print the final results
    for config in metrics:
        config_tokens = config.split("_")
        config_iters = config_tokens[1][:-1]
        config_cascades = config_tokens[0][:-1]

        metric_log_dict[config_cascades] = metric_log_dict.get(config_cascades, {})
        metric_log_dict[config_cascades][config_iters] = metrics[config]

        evaluation_str = f"{args.dataset} evaluation@ size:{args.eval_size} n_cascades:{config_cascades} recurrent_updates:{config_iters}"
        metrics_str = f"Metrics: {metrics[config]}"
        metric_log.extend([evaluation_str, metrics_str])

        print(evaluation_str)
        print(metrics_str)

    eval_log_name = f"{checkpoint_name.replace('.pth', '')}_eval.log"
    print("Saving eval log to: ", eval_log_name)
    with open(eval_log_name, "w") as f:
        f.write(f"Dataset: {args.dataset} @size: {args.eval_size}:\n")
        # write the dict line by line for each key, and each value in the keys
        for config_cascades in metric_log_dict:
            f.write("{\n")
            f.write(f"\t{config_cascades}: {{\n")
            for config_iters in metric_log_dict[config_cascades]:
                # convert every metric to 4 decimal places
                metrics = metric_log_dict[config_cascades][config_iters]
                metrics = {k: float(f"{v:.3f}") for k, v in metrics.items()}
                f.write(f"\t\t{config_iters}: {metrics}\n")
            f.write("\t},\n")
            f.write("}\n")


def load_checkpoint(args):
    utils.setup_ddp(args)

    if not args.weights:
        checkpoint = torch.load(args.checkpoint, map_location=torch.device("cpu"), weights_only=True)
        if "model" in checkpoint:
            experiment_args = checkpoint["args"]
            model = torchvision.prototype.models.depth.stereo.__dict__[experiment_args.model](weights=None)
            model.load_state_dict(checkpoint["model"])
        else:
            model = torchvision.prototype.models.depth.stereo.__dict__[args.model](weights=None)
            model.load_state_dict(checkpoint)

        # set the appropriate devices
        if args.distributed and args.device == "cpu":
            raise ValueError("The device must be cuda if we want to run in distributed mode using torchrun")
        device = torch.device(args.device)
    else:
        model = torchvision.prototype.models.depth.stereo.__dict__[args.model](weights=args.weights)

    # convert to DDP if need be
    if args.distributed:
        model = model.to(args.device)
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
    else:
        model.to(device)

    return model


def main(args):
    model = load_checkpoint(args)
    loader = make_eval_loader(args.dataset, args)
    evaluate(model, loader, args)


if __name__ == "__main__":
    args = get_args_parser().parse_args()
    main(args)