        if self.totensor:
            Limg = img_to_tensor(Limg)
            Rimg = img_to_tensor(Rimg)
            if disp is None:
                disp = torch.tensor(
                    []
                )  # to allow dataloader batching with default collate_gn
            else:
                disp = disp_to_tensor(disp)

        return Limg, Rimg, disp, str(pairname)

    def __rmul__(self, v):
        self.rmul *= v
        self.pairnames = v * self.pairnames
        return self

    def __str__(self):
        return f"{self.__class__.__name__}_{self.split}"

    def __repr__(self):
        s = f"{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})"
        if self.rmul == 1:
            s += f"\n\tnum pairs: {len(self.pairnames)}"
        else:
            s += f"\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})"
        return s

    def _set_root(self):
        self.root = dataset_to_root[self.name]
        assert os.path.isdir(
            self.root
        ), f"could not find root directory for dataset {self.name}: {self.root}"

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")
        if osp.isfile(cache_file):
            with open(cache_file, "rb") as fid:
                self.pairnames = pickle.load(fid)[self.split]
        else:
            tosave = self._build_cache()
            os.makedirs(cache_dir, exist_ok=True)
            with open(cache_file, "wb") as fid:
                pickle.dump(tosave, fid)
            self.pairnames = tosave[self.split]


                disp = disp_to_tensor(disp)

        return Limg, Rimg, disp, str(pairname)

    def __rmul__(self, v):
        self.rmul *= v
        self.pairnames = v * self.pairnames
        return self

    def __str__(self):
        return f"{self.__class__.__name__}_{self.split}"

    def __repr__(self):
        s = f"{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})"
        if self.rmul == 1:
            s += f"\n\tnum pairs: {len(self.pairnames)}"
        else:
            s += f"\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})"
        return s

    def _set_root(self):
        self.root = dataset_to_root[self.name]
        assert os.path.isdir(
            self.root
        ), f"could not find root directory for dataset {self.name}: {self.root}"

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")
        if osp.isfile(cache_file):
            with open(cache_file, "rb") as fid:
                self.pairnames = pickle.load(fid)[self.split]
        else:
            tosave = self._build_cache()
            os.makedirs(cache_dir, exist_ok=True)
            with open(cache_file, "wb") as fid:
                pickle.dump(tosave, fid)
            self.pairnames = tosave[self.split]

class CREStereoDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "CREStereo"
        self._set_root()
        assert self.split in ["train"]


        self.pairnames = v * self.pairnames
        return self

    def __str__(self):
        return f"{self.__class__.__name__}_{self.split}"

    def __repr__(self):
        s = f"{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})"
        if self.rmul == 1:
            s += f"\n\tnum pairs: {len(self.pairnames)}"
        else:
            s += f"\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})"
        return s

    def _set_root(self):
        self.root = dataset_to_root[self.name]
        assert os.path.isdir(
            self.root
        ), f"could not find root directory for dataset {self.name}: {self.root}"

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")
        if osp.isfile(cache_file):
            with open(cache_file, "rb") as fid:
                self.pairnames = pickle.load(fid)[self.split]
        else:
            tosave = self._build_cache()
            os.makedirs(cache_dir, exist_ok=True)
            with open(cache_file, "wb") as fid:
                pickle.dump(tosave, fid)
            self.pairnames = tosave[self.split]

class CREStereoDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "CREStereo"
        self._set_root()
        assert self.split in ["train"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname + "_left.jpg"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname + "_right.jpg"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname + "_left.disp.png"
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_disparity = _read_crestereo_disp


        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 2
        assert prediction.dtype == np.float32
        outfile = os.path.join(outdir, pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        img = (prediction * 256).astype("uint16")
        Image.fromarray(img).save(outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        cmd = f'cd {outdir}/; zip -r "kitti12_results.zip" .'
        print(cmd)
        os.system(cmd)
        print(f"Done. Submission file at {outdir}/kitti12_results.zip")

class Kitti15Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Kitti15"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval", "test"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname.replace("/image_2/", "/image_3/") + "_10.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/image_2/", "/disp_occ_0/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/image_2/", "/")
        self.load_disparity = _read_kitti_disp

    def _build_cache(self):
        trainseqs = ["training/image_2/%06d" % (i) for i in range(200)]
        subtrainseqs = trainseqs[:-5]
        subvalseqs = trainseqs[-5:]
        testseqs = ["testing/image_2/%06d" % (i) for i in range(200)]
        assert (
            len(trainseqs) == 200
            and len(subtrainseqs) == 195
            and len(subvalseqs) == 5
            and len(testseqs) == 200
        ), "incorrect parsing of pairs in Kitti15"
        tosave = {
            "train": trainseqs,
            "subtrain": subtrainseqs,
            "subval": subvalseqs,
            "test": testseqs,
        }
        return tosave


    def _prepare_data(self):
        self.name = "Kitti15"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval", "test"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname.replace("/image_2/", "/image_3/") + "_10.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/image_2/", "/disp_occ_0/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/image_2/", "/")
        self.load_disparity = _read_kitti_disp

    def _build_cache(self):
        trainseqs = ["training/image_2/%06d" % (i) for i in range(200)]
        subtrainseqs = trainseqs[:-5]
        subvalseqs = trainseqs[-5:]
        testseqs = ["testing/image_2/%06d" % (i) for i in range(200)]
        assert (
            len(trainseqs) == 200
            and len(subtrainseqs) == 195
            and len(subvalseqs) == 5
            and len(testseqs) == 200
        ), "incorrect parsing of pairs in Kitti15"
        tosave = {
            "train": trainseqs,
            "subtrain": subtrainseqs,
            "subval": subvalseqs,
            "test": testseqs,
        }
        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 2
        assert prediction.dtype == np.float32
        outfile = os.path.join(outdir, "disp_0", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        img = (prediction * 256).astype("uint16")
        Image.fromarray(img).save(outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        cmd = f'cd {outdir}/; zip -r "kitti15_results.zip" disp_0'
        print(cmd)
        os.system(cmd)
        print(f"Done. Submission file at {outdir}/kitti15_results.zip")

### auxiliary functions


    def _build_cache(self):
        trainseqs = ["training/image_2/%06d" % (i) for i in range(200)]
        subtrainseqs = trainseqs[:-5]
        subvalseqs = trainseqs[-5:]
        testseqs = ["testing/image_2/%06d" % (i) for i in range(200)]
        assert (
            len(trainseqs) == 200
            and len(subtrainseqs) == 195
            and len(subvalseqs) == 5
            and len(testseqs) == 200
        ), "incorrect parsing of pairs in Kitti15"
        tosave = {
            "train": trainseqs,
            "subtrain": subtrainseqs,
            "subval": subvalseqs,
            "test": testseqs,
        }
        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 2
        assert prediction.dtype == np.float32
        outfile = os.path.join(outdir, "disp_0", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        img = (prediction * 256).astype("uint16")
        Image.fromarray(img).save(outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        cmd = f'cd {outdir}/; zip -r "kitti15_results.zip" disp_0'
        print(cmd)
        os.system(cmd)
        print(f"Done. Submission file at {outdir}/kitti15_results.zip")

### auxiliary functions

def _read_img(filename):
    # convert to RGB for scene flow finalpass data
    img = np.asarray(Image.open(filename).convert("RGB"))
    return img

def _read_booster_disp(filename):
    disp = np.load(filename)


            and len(subtrainseqs) == 195
            and len(subvalseqs) == 5
            and len(testseqs) == 200
        ), "incorrect parsing of pairs in Kitti15"
        tosave = {
            "train": trainseqs,
            "subtrain": subtrainseqs,
            "subval": subvalseqs,
            "test": testseqs,
        }
        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 2
        assert prediction.dtype == np.float32
        outfile = os.path.join(outdir, "disp_0", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        img = (prediction * 256).astype("uint16")
        Image.fromarray(img).save(outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        cmd = f'cd {outdir}/; zip -r "kitti15_results.zip" disp_0'
        print(cmd)
        os.system(cmd)
        print(f"Done. Submission file at {outdir}/kitti15_results.zip")

### auxiliary functions

def _read_img(filename):
    # convert to RGB for scene flow finalpass data
    img = np.asarray(Image.open(filename).convert("RGB"))
    return img

def _read_booster_disp(filename):
    disp = np.load(filename)
    disp[disp == 0.0] = np.inf
    return disp

def _read_png_disp(filename, coef=1.0):
    disp = np.asarray(Image.open(filename))
    disp = disp.astype(np.float32) / coef


src.croco.stereoflow.test

# 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)


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)


        "--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")


    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(


main

        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)


src.croco.stereoflow.engine

# 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)]


# 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)


# --------------------------------------------------------

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


        # 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


        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,


            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


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)]


        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)]


    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)]


    # 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))


# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

# --------------------------------------------------------
# Dataset structure for flow
# --------------------------------------------------------

import os
import os.path as osp
import pickle
import numpy as np
import struct
from PIL import Image
import json
import h5py
import torch
from torch.utils import data

from .augmentor import FlowAugmentor
from .datasets_stereo import _read_img, img_to_tensor, dataset_to_root, _read_pfm
from copy import deepcopy

dataset_to_root = deepcopy(dataset_to_root)

dataset_to_root.update(
    **{
        "TartanAir": "./data/stereoflow/TartanAir",
        "FlyingChairs": "./data/stereoflow/FlyingChairs/",
        "FlyingThings": osp.join(dataset_to_root["SceneFlow"], "FlyingThings") + "/",
        "MPISintel": "./data/stereoflow//MPI-Sintel/" + "/",
    }
)
cache_dir = "./data/stereoflow/datasets_flow_cache/"

def flow_to_tensor(disp):
    return torch.from_numpy(disp).float().permute(2, 0, 1)

class FlowDataset(data.Dataset):

    def __init__(self, split, augmentor=False, crop_size=None, totensor=True):
        self.split = split
        if not augmentor:
            assert crop_size is None
        if crop_size is not None:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = FlowAugmentor(crop_size) if augmentor else None
        self.totensor = totensor
        self.rmul = 1  # keep track of rmul
        self.has_constant_resolution = True  # whether the dataset has constant resolution or not (=> don't use batch_size>1 at test time)
        self._prepare_data()
        self._load_or_build_cache()

    def prepare_data(self):
        """
        to be defined for each dataset
        """
        raise NotImplementedError

    def __len__(self):
        return len(
            self.pairnames
        )  # each pairname is typically of the form (str, int1, int2)

    def __getitem__(self, index):
        pairname = self.pairnames[index]

        # get filenames
        img1name = self.pairname_to_img1name(pairname)
        img2name = self.pairname_to_img2name(pairname)
        flowname = (
            self.pairname_to_flowname(pairname)
            if self.pairname_to_flowname is not None
            else None
        )

        # load images and disparities
        img1 = _read_img(img1name)
        img2 = _read_img(img2name)
        flow = self.load_flow(flowname) if flowname is not None else None

        # apply augmentations
        if self.augmentor is not None:
            img1, img2, flow = self.augmentor(img1, img2, flow, self.name)

        if self.totensor:
            img1 = img_to_tensor(img1)
            img2 = img_to_tensor(img2)
            if flow is not None:
                flow = flow_to_tensor(flow)
            else:
                flow = torch.tensor(
                    []
                )  # to allow dataloader batching with default collate_gn
            pairname = str(
                pairname
            )  # transform potential tuple to str to be able to batch it

        return img1, img2, flow, pairname

    def __rmul__(self, v):
        self.rmul *= v
        self.pairnames = v * self.pairnames
        return self

    def __str__(self):
        return f"{self.__class__.__name__}_{self.split}"

    def __repr__(self):
        s = f"{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})"
        if self.rmul == 1:
            s += f"\n\tnum pairs: {len(self.pairnames)}"
        else:
            s += f"\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})"
        return s

    def _set_root(self):
        self.root = dataset_to_root[self.name]
        assert os.path.isdir(
            self.root
        ), f"could not find root directory for dataset {self.name}: {self.root}"

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")
        if osp.isfile(cache_file):
            with open(cache_file, "rb") as fid:
                self.pairnames = pickle.load(fid)[self.split]
        else:
            tosave = self._build_cache()
            os.makedirs(cache_dir, exist_ok=True)
            with open(cache_file, "wb") as fid:
                pickle.dump(tosave, fid)
            self.pairnames = tosave[self.split]

class TartanAirDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "TartanAir"
        self._set_root()
        assert self.split in ["train"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname[0], "image_left/{:06d}_left.png".format(pairname[1])
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname[0], "image_left/{:06d}_left.png".format(pairname[2])
        )
        self.pairname_to_flowname = lambda pairname: osp.join(
            self.root,
            pairname[0],
            "flow/{:06d}_{:06d}_flow.npy".format(pairname[1], pairname[2]),
        )
        self.pairname_to_str = lambda pairname: os.path.join(
            pairname[0][pairname[0].find("/") + 1 :],
            "{:06d}_{:06d}".format(pairname[1], pairname[2]),
        )
        self.load_flow = _read_numpy_flow

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        pairs = [
            (osp.join(s, s, difficulty, Pxxx), int(a[:6]), int(a[:6]) + 1)
            for s in seqs
            for difficulty in ["Easy", "Hard"]
            for Pxxx in sorted(os.listdir(osp.join(self.root, s, s, difficulty)))
            for a in sorted(
                os.listdir(osp.join(self.root, s, s, difficulty, Pxxx, "image_left/"))
            )[:-1]
        ]
        assert len(pairs) == 306268, "incorrect parsing of pairs in TartanAir"
        tosave = {"train": pairs}
        return tosave

class FlyingChairsDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "FlyingChairs"
        self._set_root()
        assert self.split in ["train", "val"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, "data", pairname + "_img1.ppm"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, "data", pairname + "_img2.ppm"
        )
        self.pairname_to_flowname = lambda pairname: osp.join(
            self.root, "data", pairname + "_flow.flo"
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_flow = _read_flo_file

    def _build_cache(self):
        split_file = osp.join(self.root, "chairs_split.txt")
        split_list = np.loadtxt(split_file, dtype=np.int32)
        trainpairs = ["{:05d}".format(i) for i in np.where(split_list == 1)[0] + 1]
        valpairs = ["{:05d}".format(i) for i in np.where(split_list == 2)[0] + 1]
        assert (
            len(trainpairs) == 22232 and len(valpairs) == 640
        ), "incorrect parsing of pairs in MPI-Sintel"
        tosave = {"train": trainpairs, "val": valpairs}
        return tosave

class FlyingThingsDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "FlyingThings"
        self._set_root()
        assert self.split in [
            f"{set_}_{pass_}pass{camstr}"
            for set_ in ["train", "test", "test1024"]
            for camstr in ["", "_rightcam"]
            for pass_ in ["clean", "final", "all"]
        ]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root,
            f"frames_{pairname[3]}pass",
            pairname[0].replace("into_future", "").replace("into_past", ""),
            "{:04d}.png".format(pairname[1]),
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root,
            f"frames_{pairname[3]}pass",
            pairname[0].replace("into_future", "").replace("into_past", ""),
            "{:04d}.png".format(pairname[2]),
        )
        self.pairname_to_flowname = lambda pairname: osp.join(
            self.root,
            "optical_flow",
            pairname[0],
            "OpticalFlowInto{f:s}_{i:04d}_{c:s}.pfm".format(
                f="Future" if "future" in pairname[0] else "Past",
                i=pairname[1],
                c="L" if "left" in pairname[0] else "R",
            ),
        )
        self.pairname_to_str = lambda pairname: os.path.join(
            pairname[3] + "pass",
            pairname[0],
            "Into{f:s}_{i:04d}_{c:s}".format(
                f="Future" if "future" in pairname[0] else "Past",
                i=pairname[1],
                c="L" if "left" in pairname[0] else "R",
            ),
        )
        self.load_flow = _read_pfm_flow

    def _build_cache(self):
        tosave = {}
        # train and test splits for the different passes
        for set_ in ["train", "test"]:
            sroot = osp.join(self.root, "optical_flow", set_.upper())
            fname_to_i = lambda f: int(
                f[len("OpticalFlowIntoFuture_") : -len("_L.pfm")]
            )
            pp = [
                (osp.join(set_.upper(), d, s, "into_future/left"), fname_to_i(fname))
                for d in sorted(os.listdir(sroot))
                for s in sorted(os.listdir(osp.join(sroot, d)))
                for fname in sorted(
                    os.listdir(osp.join(sroot, d, s, "into_future/left"))
                )[:-1]
            ]
            pairs = [(a, i, i + 1) for a, i in pp]
            pairs += [(a.replace("into_future", "into_past"), i + 1, i) for a, i in pp]
            assert (
                len(pairs) == {"train": 40302, "test": 7866}[set_]
            ), "incorrect parsing of pairs Flying Things"
            for cam in ["left", "right"]:
                camstr = "" if cam == "left" else f"_{cam}cam"
                for pass_ in ["final", "clean"]:
                    tosave[f"{set_}_{pass_}pass{camstr}"] = [
                        (a.replace("left", cam), i, j, pass_) for a, i, j in pairs
                    ]
                tosave[f"{set_}_allpass{camstr}"] = (
                    tosave[f"{set_}_cleanpass{camstr}"]
                    + tosave[f"{set_}_finalpass{camstr}"]
                )
        # test1024: this is the same split as unimatch 'validation' split
        # see https://github.com/autonomousvision/unimatch/blob/master/dataloader/flow/datasets.py#L229
        test1024_nsamples = 1024
        alltest_nsamples = len(tosave["test_cleanpass"])  # 7866
        stride = alltest_nsamples // test1024_nsamples
        remove = alltest_nsamples % test1024_nsamples
        for cam in ["left", "right"]:
            camstr = "" if cam == "left" else f"_{cam}cam"
            for pass_ in ["final", "clean"]:
                tosave[f"test1024_{pass_}pass{camstr}"] = sorted(
                    tosave[f"test_{pass_}pass{camstr}"]
                )[:-remove][
                    ::stride
                ]  # warning, it was not sorted before
            assert (
                len(tosave["test1024_cleanpass"]) == 1024
            ), "incorrect parsing of pairs in Flying Things"
            tosave[f"test1024_allpass{camstr}"] = (
                tosave[f"test1024_cleanpass{camstr}"]
                + tosave[f"test1024_finalpass{camstr}"]
            )
        return tosave

class MPISintelDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "MPISintel"
        self._set_root()
        assert self.split in [
            s + "_" + p
            for s in ["train", "test", "subval", "subtrain"]
            for p in ["cleanpass", "finalpass", "allpass"]
        ]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname[0], "frame_{:04d}.png".format(pairname[1])
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname[0], "frame_{:04d}.png".format(pairname[1] + 1)
        )
        self.pairname_to_flowname = lambda pairname: (
            None
            if pairname[0].startswith("test/")
            else osp.join(
                self.root,
                pairname[0].replace("/clean/", "/flow/").replace("/final/", "/flow/"),
                "frame_{:04d}.flo".format(pairname[1]),
            )
        )
        self.pairname_to_str = lambda pairname: osp.join(
            pairname[0], "frame_{:04d}".format(pairname[1])
        )
        self.load_flow = _read_flo_file

    def _build_cache(self):
        trainseqs = sorted(os.listdir(self.root + "training/clean"))
        trainpairs = [
            (osp.join("training/clean", s), i)
            for s in trainseqs
            for i in range(1, len(os.listdir(self.root + "training/clean/" + s)))
        ]
        subvalseqs = ["temple_2", "temple_3"]
        subtrainseqs = [s for s in trainseqs if s not in subvalseqs]
        subvalpairs = [(p, i) for p, i in trainpairs if any(s in p for s in subvalseqs)]
        subtrainpairs = [
            (p, i) for p, i in trainpairs if any(s in p for s in subtrainseqs)
        ]
        testseqs = sorted(os.listdir(self.root + "test/clean"))
        testpairs = [
            (osp.join("test/clean", s), i)
            for s in testseqs
            for i in range(1, len(os.listdir(self.root + "test/clean/" + s)))
        ]
        assert (
            len(trainpairs) == 1041
            and len(testpairs) == 552
            and len(subvalpairs) == 98
            and len(subtrainpairs) == 943
        ), "incorrect parsing of pairs in MPI-Sintel"
        tosave = {}
        tosave["train_cleanpass"] = trainpairs
        tosave["test_cleanpass"] = testpairs
        tosave["subval_cleanpass"] = subvalpairs
        tosave["subtrain_cleanpass"] = subtrainpairs
        for t in ["train", "test", "subval", "subtrain"]:
            tosave[t + "_finalpass"] = [
                (p.replace("/clean/", "/final/"), i)
                for p, i in tosave[t + "_cleanpass"]
            ]
            tosave[t + "_allpass"] = tosave[t + "_cleanpass"] + tosave[t + "_finalpass"]
        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, _time):
        assert prediction.shape[2] == 2
        outfile = os.path.join(
            outdir, "submission", self.pairname_to_str(pairname) + ".flo"
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowFile(prediction, outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test_allpass"
        bundle_exe = "/nfs/data/ffs-3d/datasets/StereoFlow/MPI-Sintel/bundler/linux-x64/bundler"  # eg <bundle_exe> <path_to_results_for_clean> <path_to_results_for_final> <output/bundled.lzma>
        if os.path.isfile(bundle_exe):
            cmd = f'{bundle_exe} "{outdir}/submission/test/clean/" "{outdir}/submission/test/final" "{outdir}/submission/bundled.lzma"'
            print(cmd)
            os.system(cmd)
            print(f'Done. Submission file at: "{outdir}/submission/bundled.lzma"')
        else:
            print("Could not find bundler executable for submission.")
            print("Please download it and run:")
            print(
                f'<bundle_exe> "{outdir}/submission/test/clean/" "{outdir}/submission/test/final" "{outdir}/submission/bundled.lzma"'
            )

class SpringDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Spring"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root,
            pairname[0],
            pairname[1],
            "frame_" + pairname[3],
            "frame_{:s}_{:04d}.png".format(pairname[3], pairname[4]),
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root,
            pairname[0],
            pairname[1],
            "frame_" + pairname[3],
            "frame_{:s}_{:04d}.png".format(
                pairname[3], pairname[4] + (1 if pairname[2] == "FW" else -1)
            ),
        )
        self.pairname_to_flowname = lambda pairname: (
            None
            if pairname[0] == "test"
            else osp.join(
                self.root,
                pairname[0],
                pairname[1],
                f"flow_{pairname[2]}_{pairname[3]}",
                f"flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}.flo5",
            )
        )
        self.pairname_to_str = lambda pairname: osp.join(
            pairname[0],
            pairname[1],
            f"flow_{pairname[2]}_{pairname[3]}",
            f"flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}",
        )
        self.load_flow = _read_hdf5_flow

    def _build_cache(self):
        # train
        trainseqs = sorted(os.listdir(osp.join(self.root, "train")))
        trainpairs = []
        for leftright in ["left", "right"]:
            for fwbw in ["FW", "BW"]:
                trainpairs += [
                    (
                        "train",
                        s,
                        fwbw,
                        leftright,
                        int(f[len(f"flow_{fwbw}_{leftright}_") : -len(".flo5")]),
                    )
                    for s in trainseqs
                    for f in sorted(
                        os.listdir(
                            osp.join(self.root, "train", s, f"flow_{fwbw}_{leftright}")
                        )
                    )
                ]
        # test
        testseqs = sorted(os.listdir(osp.join(self.root, "test")))
        testpairs = []
        for leftright in ["left", "right"]:
            testpairs += [
                (
                    "test",
                    s,
                    "FW",
                    leftright,
                    int(f[len(f"frame_{leftright}_") : -len(".png")]),
                )
                for s in testseqs
                for f in sorted(
                    os.listdir(osp.join(self.root, "test", s, f"frame_{leftright}"))
                )[:-1]
            ]
            testpairs += [
                (
                    "test",
                    s,
                    "BW",
                    leftright,
                    int(f[len(f"frame_{leftright}_") : -len(".png")]) + 1,
                )
                for s in testseqs
                for f in sorted(
                    os.listdir(osp.join(self.root, "test", s, f"frame_{leftright}"))
                )[:-1]
            ]
        # subtrain / subval
        subtrainpairs = [p for p in trainpairs if p[1] != "0041"]
        subvalpairs = [p for p in trainpairs if p[1] == "0041"]
        assert (
            len(trainpairs) == 19852
            and len(testpairs) == 3960
            and len(subtrainpairs) == 19472
            and len(subvalpairs) == 380
        ), "incorrect parsing of pairs in Spring"
        tosave = {
            "train": trainpairs,
            "test": testpairs,
            "subtrain": subtrainpairs,
            "subval": subvalpairs,
        }
        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 3
        assert prediction.shape[2] == 2
        assert prediction.dtype == np.float32
        outfile = osp.join(
            outdir,
            pairname[0],
            pairname[1],
            f"flow_{pairname[2]}_{pairname[3]}",
            f"flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}.flo5",
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlo5File(prediction, outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        exe = "{self.root}/flow_subsampling"
        if os.path.isfile(exe):
            cmd = f'cd "{outdir}/test"; {exe} .'
            print(cmd)
            os.system(cmd)
            pr
# ... truncated ...


import torch
from torch.utils import data

from .augmentor import FlowAugmentor
from .datasets_stereo import _read_img, img_to_tensor, dataset_to_root, _read_pfm
from copy import deepcopy

dataset_to_root = deepcopy(dataset_to_root)

dataset_to_root.update(
    **{
        "TartanAir": "./data/stereoflow/TartanAir",
        "FlyingChairs": "./data/stereoflow/FlyingChairs/",
        "FlyingThings": osp.join(dataset_to_root["SceneFlow"], "FlyingThings") + "/",
        "MPISintel": "./data/stereoflow//MPI-Sintel/" + "/",
    }
)
cache_dir = "./data/stereoflow/datasets_flow_cache/"

def flow_to_tensor(disp):
    return torch.from_numpy(disp).float().permute(2, 0, 1)

class FlowDataset(data.Dataset):

    def __init__(self, split, augmentor=False, crop_size=None, totensor=True):
        self.split = split
        if not augmentor:
            assert crop_size is None
        if crop_size is not None:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = FlowAugmentor(crop_size) if augmentor else None
        self.totensor = totensor
        self.rmul = 1  # keep track of rmul
        self.has_constant_resolution = True  # whether the dataset has constant resolution or not (=> don't use batch_size>1 at test time)
        self._prepare_data()
        self._load_or_build_cache()

    def prepare_data(self):


FlowDataset

from .datasets_stereo import _read_img, img_to_tensor, dataset_to_root, _read_pfm
from copy import deepcopy

dataset_to_root = deepcopy(dataset_to_root)

dataset_to_root.update(
    **{
        "TartanAir": "./data/stereoflow/TartanAir",
        "FlyingChairs": "./data/stereoflow/FlyingChairs/",
        "FlyingThings": osp.join(dataset_to_root["SceneFlow"], "FlyingThings") + "/",
        "MPISintel": "./data/stereoflow//MPI-Sintel/" + "/",
    }
)
cache_dir = "./data/stereoflow/datasets_flow_cache/"

def flow_to_tensor(disp):
    return torch.from_numpy(disp).float().permute(2, 0, 1)

class FlowDataset(data.Dataset):

    def __init__(self, split, augmentor=False, crop_size=None, totensor=True):
        self.split = split
        if not augmentor:
            assert crop_size is None
        if crop_size is not None:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = FlowAugmentor(crop_size) if augmentor else None
        self.totensor = totensor
        self.rmul = 1  # keep track of rmul
        self.has_constant_resolution = True  # whether the dataset has constant resolution or not (=> don't use batch_size>1 at test time)
        self._prepare_data()
        self._load_or_build_cache()

    def prepare_data(self):
        """
        to be defined for each dataset
        """
        raise NotImplementedError

    def __len__(self):
        return len(
            self.pairnames
        )  # each pairname is typically of the form (str, int1, int2)

    def __getitem__(self, index):
        pairname = self.pairnames[index]

        # get filenames
        img1name = self.pairname_to_img1name(pairname)
        img2name = self.pairname_to_img2name(pairname)
        flowname = (
            self.pairname_to_flowname(pairname)
            if self.pairname_to_flowname is not None
            else None
        )

        # load images and disparities
        img1 = _read_img(img1name)
        img2 = _read_img(img2name)
        flow = self.load_flow(flowname) if flowname is not None else None

        # apply augmentations
        if self.augmentor is not None:
            img1, img2, flow = self.augmentor(img1, img2, flow, self.name)

        if self.totensor:
            img1 = img_to_tensor(img1)
            img2 = img_to_tensor(img2)
            if flow is not None:
                flow = flow_to_tensor(flow)
            else:
                flow = torch.tensor(
                    []
                )  # to allow dataloader batching with default collate_gn
            pairname = str(
                pairname
            )  # transform potential tuple to str to be able to batch it

        return img1, img2, flow, pairname

    def __rmul__(self, v):
        self.rmul *= v
        self.pairnames = v * self.pairnames
        return self

    def __str__(self):
        return f"{self.__class__.__name__}_{self.split}"

    def __repr__(self):
        s = f"{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})"
        if self.rmul == 1:
            s += f"\n\tnum pairs: {len(self.pairnames)}"
        else:
            s += f"\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})"
        return s

    def _set_root(self):
        self.root = dataset_to_root[self.name]
        assert os.path.isdir(
            self.root
        ), f"could not find root directory for dataset {self.name}: {self.root}"

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")
        if osp.isfile(cache_file):
            with open(cache_file, "rb") as fid:
                self.pairnames = pickle.load(fid)[self.split]
        else:
            tosave = self._build_cache()
            os.makedirs(cache_dir, exist_ok=True)
            with open(cache_file, "wb") as fid:
                pickle.dump(tosave, fid)
            self.pairnames = tosave[self.split]

class TartanAirDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "TartanAir"
        self._set_root()
        assert self.split in ["train"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname[0], "image_left/{:06d}_left.png".format(pairname[1])
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname[0], "image_left/{:06d}_left.png".format(pairname[2])
        )
        self.pairname_to_flowname = lambda pairname: osp.join(
            self.root,
            pairname[0],
            "flow/{:06d}_{:06d}_flow.npy".format(pairname[1], pairname[2]),
        )
        self.pairname_to_str = lambda pairname: os.path.join(