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

# --------------------------------------------------------
# Dataset structure for stereo
# --------------------------------------------------------

import sys, os
import os.path as osp
import pickle
import numpy as np
from PIL import Image
import json
import h5py
from glob import glob
import cv2

import torch
from torch.utils import data

from .augmentor import StereoAugmentor

dataset_to_root = {
    "CREStereo": "./data/stereoflow//crenet_stereo_trainset/stereo_trainset/crestereo/",
    "SceneFlow": "./data/stereoflow//SceneFlow/",
    "ETH3DLowRes": "./data/stereoflow/eth3d_lowres/",
    "Booster": "./data/stereoflow/booster_gt/",
    "Middlebury2021": "./data/stereoflow/middlebury/2021/data/",
    "Middlebury2014": "./data/stereoflow/middlebury/2014/",
    "Middlebury2006": "./data/stereoflow/middlebury/2006/",
    "Middlebury2005": "./data/stereoflow/middlebury/2005/train/",
    "MiddleburyEval3": "./data/stereoflow/middlebury/MiddEval3/",
    "Spring": "./data/stereoflow/spring/",
    "Kitti15": "./data/stereoflow/kitti-stereo-2015/",
    "Kitti12": "./data/stereoflow/kitti-stereo-2012/",
}
cache_dir = "./data/stereoflow/datasets_stereo_cache/"

in1k_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
in1k_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

def img_to_tensor(img):
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = (img - in1k_mean) / in1k_std
    return img

def disp_to_tensor(disp):
    return torch.from_numpy(disp)[None, :, :]

class StereoDataset(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:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)

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

        # get filenames
        Limgname = self.pairname_to_Limgname(pairname)
        Rimgname = self.pairname_to_Rimgname(pairname)
        Ldispname = (
            self.pairname_to_Ldispname(pairname)
            if self.pairname_to_Ldispname is not None
            else None
        )

        # load images and disparities
        Limg = _read_img(Limgname)
        Rimg = _read_img(Rimgname)
        disp = self.load_disparity(Ldispname) if Ldispname is not None else None

        # sanity check
        if disp is not None:
            assert np.all(disp > 0) or self.name == "Spring", (
                self.name,
                pairname,
                Ldispname,
            )

        # apply augmentations
        if self.augmentor is not None:
            Limg, Rimg, disp = self.augmentor(Limg, Rimg, disp, self.name)

        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]

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

    def _build_cache(self):
        allpairs = [
            s + "/" + f[: -len("_left.jpg")]
            for s in sorted(os.listdir(self.root))
            for f in sorted(os.listdir(self.root + "/" + s))
            if f.endswith("_left.jpg")
        ]
        assert len(allpairs) == 200000, "incorrect parsing of pairs in CreStereo"
        tosave = {"train": allpairs}
        return tosave

class SceneFlowDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "SceneFlow"
        self._set_root()
        assert self.split in [
            "train_finalpass",
            "train_cleanpass",
            "train_allpass",
            "test_finalpass",
            "test_cleanpass",
            "test_allpass",
            "test1of100_cleanpass",
            "test1of100_finalpass",
        ]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname
        ).replace("/left/", "/right/")
        self.pairname_to_Ldispname = (
            lambda pairname: osp.join(self.root, pairname)
            .replace("/frames_finalpass/", "/disparity/")
            .replace("/frames_cleanpass/", "/disparity/")[:-4]
            + ".pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_sceneflow_disp

    def _build_cache(self):
        trainpairs = []
        # driving
        pairs = sorted(glob(self.root + "Driving/frames_finalpass/*/*/*/left/*.png"))
        pairs = list(map(lambda x: x[len(self.root) :], pairs))
        assert len(pairs) == 4400, "incorrect parsing of pairs in SceneFlow"
        trainpairs += pairs
        # monkaa
        pairs = sorted(glob(self.root + "Monkaa/frames_finalpass/*/left/*.png"))
        pairs = list(map(lambda x: x[len(self.root) :], pairs))
        assert len(pairs) == 8664, "incorrect parsing of pairs in SceneFlow"
        trainpairs += pairs
        # flyingthings
        pairs = sorted(
            glob(self.root + "FlyingThings/frames_finalpass/TRAIN/*/*/left/*.png")
        )
        pairs = list(map(lambda x: x[len(self.root) :], pairs))
        assert len(pairs) == 22390, "incorrect parsing of pairs in SceneFlow"
        trainpairs += pairs
        assert len(trainpairs) == 35454, "incorrect parsing of pairs in SceneFlow"
        testpairs = sorted(
            glob(self.root + "FlyingThings/frames_finalpass/TEST/*/*/left/*.png")
        )
        testpairs = list(map(lambda x: x[len(self.root) :], testpairs))
        assert len(testpairs) == 4370, "incorrect parsing of pairs in SceneFlow"
        test1of100pairs = testpairs[::100]
        assert len(test1of100pairs) == 44, "incorrect parsing of pairs in SceneFlow"
        # all
        tosave = {
            "train_finalpass": trainpairs,
            "train_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    trainpairs,
                )
            ),
            "test_finalpass": testpairs,
            "test_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    testpairs,
                )
            ),
            "test1of100_finalpass": test1of100pairs,
            "test1of100_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    test1of100pairs,
                )
            ),
        }
        tosave["train_allpass"] = tosave["train_finalpass"] + tosave["train_cleanpass"]
        tosave["test_allpass"] = tosave["test_finalpass"] + tosave["test_cleanpass"]
        return tosave

class Md21Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2021"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname.replace("/im0", "/im1")
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp0.pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury_disp

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            # trainpairs += [s+'/im0.png'] # we should remove it, it is included as such in other lightings
            trainpairs += [
                s + "/ambient/" + b + "/" + a
                for b in sorted(os.listdir(osp.join(self.root, s, "ambient")))
                for a in sorted(os.listdir(osp.join(self.root, s, "ambient", b)))
                if a.startswith("im0")
            ]
        assert len(trainpairs) == 355
        subtrainpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in seqs[:-2])
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in seqs[-2:])
        ]
        assert (
            len(subtrainpairs) == 335 and len(subvalpairs) == 20
        ), "incorrect parsing of pairs in Middlebury 2021"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md14Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2014"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "disp0.pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury_disp
        self.has_constant_resolution = False

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            trainpairs += [s + "/im1.png", s + "/im1E.png", s + "/im1L.png"]
        assert len(trainpairs) == 138
        valseqs = ["Umbrella-imperfect", "Vintage-perfect"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 132 and len(subvalpairs) == 6
        ), "incorrect parsing of pairs in Middlebury 2014"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md06Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2006"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "view5.png"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp1.png"
        )
        self.load_disparity = _read_middlebury20052006_disp
        self.has_constant_resolution = False

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            for i in ["Illum1", "Illum2", "Illum3"]:
                for e in ["Exp0", "Exp1", "Exp2"]:
                    trainpairs.append(osp.join(s, i, e, "view1.png"))
        assert len(trainpairs) == 189
        valseqs = ["Rocks1", "Wood2"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 171 and len(subvalpairs) == 18
        ), "incorrect parsing of pairs in Middlebury 2006"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md05Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2005"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "view5.png"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp1.png"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury20052006_disp

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            for i in ["Illum1", "Illum2", "Illum3"]:
                for e in ["Exp0", "Exp1", "Exp2"]:
                    trainpairs.append(osp.join(s, i, e, "view1.png"))
        assert len(trainpairs) == 54, "incorrect parsing of pairs in Middlebury 2005"
        valseqs = ["Reindeer"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 45 and len(subvalpairs) == 9
        ), "incorrect parsing of pairs in Middlebury 2005"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class MdEval3Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "MiddleburyEval3"
        self._set_root()
        assert self.split in [
            s + "_" + r
            for s in ["train", "subtrain", "subval", "test", "all"]
            for r in ["full", "half", "quarter"]
        ]
        if self.split.endswith("_full"):
            self.root = self.root.replace("/MiddEval3", "/MiddEval3_F")
        elif self.split.endswith("_half"):
            self.root = self.root.replace("/MiddEval3", "/MiddEval3_H")
        else:
            assert self.split.endswith("_quarter")
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname, "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname, "im1.png"
        )
        self.pairname_to_Ldispname = lambda pairname: (
            None
            if pairname.startswith("test")
            else osp.join(self.root, pairname, "disp0GT.pfm")
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_disparity = _read_middlebury_disp
        # for submission only
        self.submission_methodname = "CroCo-Stereo"
        self.submission_sresolution = (
            "F"
            if self.split.endswith("_full")
            else ("H" if self.split.endswith("_half") else "Q")
        )

    def _build_cache(self):
        trainpairs = ["train/" + s for s in sorted(os.listdir(self.root + "train/"))]
        testpairs = ["test/" + s for s in sorted(os.listdir(self.root + "test/"))]
        subvalpairs = trainpairs[-1:]
        subtrainpairs = trainpairs[:-1]
        allpairs = trainpairs + testpairs
        assert (
            len(trainpairs) == 15
            and len(testpairs) == 15
            and len(subvalpairs) == 1
            and len(subtrainpairs) == 14
            and len(allpairs) == 30
        ), "incorrect parsing of pairs in Middlebury Eval v3"
        tosave = {}
        for r in ["full", "half", "quarter"]:
            tosave.update(
                **{
                    "train_" + r: trainpairs,
                    "subtrain_" + r: subtrainpairs,
                    "subval_" + r: subvalpairs,
                    "test_" + r: testpairs,
                    "all_" + r: allpairs,
                }
            )
        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("/")[0].replace("train", "training")
            + self.submission_sresolution,
            pairname.split("/")[1],
            "disp0" + self.submission_methodname + ".pfm",
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writePFM(outfile, prediction)
        timefile = os.path.join(
            os.path.dirname(outfile), "time" + self.submission_methodname + ".txt"
        )
        with open(timefile, "w") as fid:
            fid.write(str(time))

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

class ETH3DLowResDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "ETH3DLowRes"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval", "all"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname, "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname, "im1.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: (
                None
                if pairname.startswith("test/")
                else osp.join(
                    self.root, pairname.replace("train/", "train_gt/"), "disp0GT.pfm"
                )
# ... truncated ...


    "CREStereo": "./data/stereoflow//crenet_stereo_trainset/stereo_trainset/crestereo/",
    "SceneFlow": "./data/stereoflow//SceneFlow/",
    "ETH3DLowRes": "./data/stereoflow/eth3d_lowres/",
    "Booster": "./data/stereoflow/booster_gt/",
    "Middlebury2021": "./data/stereoflow/middlebury/2021/data/",
    "Middlebury2014": "./data/stereoflow/middlebury/2014/",
    "Middlebury2006": "./data/stereoflow/middlebury/2006/",
    "Middlebury2005": "./data/stereoflow/middlebury/2005/train/",
    "MiddleburyEval3": "./data/stereoflow/middlebury/MiddEval3/",
    "Spring": "./data/stereoflow/spring/",
    "Kitti15": "./data/stereoflow/kitti-stereo-2015/",
    "Kitti12": "./data/stereoflow/kitti-stereo-2012/",
}
cache_dir = "./data/stereoflow/datasets_stereo_cache/"

in1k_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
in1k_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

def img_to_tensor(img):
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = (img - in1k_mean) / in1k_std
    return img

def disp_to_tensor(disp):
    return torch.from_numpy(disp)[None, :, :]

class StereoDataset(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:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)


    "Middlebury2006": "./data/stereoflow/middlebury/2006/",
    "Middlebury2005": "./data/stereoflow/middlebury/2005/train/",
    "MiddleburyEval3": "./data/stereoflow/middlebury/MiddEval3/",
    "Spring": "./data/stereoflow/spring/",
    "Kitti15": "./data/stereoflow/kitti-stereo-2015/",
    "Kitti12": "./data/stereoflow/kitti-stereo-2012/",
}
cache_dir = "./data/stereoflow/datasets_stereo_cache/"

in1k_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
in1k_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

def img_to_tensor(img):
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = (img - in1k_mean) / in1k_std
    return img

def disp_to_tensor(disp):
    return torch.from_numpy(disp)[None, :, :]

class StereoDataset(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:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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):


StereoDataset

    "Kitti15": "./data/stereoflow/kitti-stereo-2015/",
    "Kitti12": "./data/stereoflow/kitti-stereo-2012/",
}
cache_dir = "./data/stereoflow/datasets_stereo_cache/"

in1k_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
in1k_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

def img_to_tensor(img):
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = (img - in1k_mean) / in1k_std
    return img

def disp_to_tensor(disp):
    return torch.from_numpy(disp)[None, :, :]

class StereoDataset(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:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)

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

        # get filenames
        Limgname = self.pairname_to_Limgname(pairname)
        Rimgname = self.pairname_to_Rimgname(pairname)
        Ldispname = (
            self.pairname_to_Ldispname(pairname)
            if self.pairname_to_Ldispname is not None
            else None
        )

        # load images and disparities
        Limg = _read_img(Limgname)
        Rimg = _read_img(Rimgname)
        disp = self.load_disparity(Ldispname) if Ldispname is not None else None

        # sanity check
        if disp is not None:
            assert np.all(disp > 0) or self.name == "Spring", (
                self.name,
                pairname,
                Ldispname,
            )

        # apply augmentations
        if self.augmentor is not None:
            Limg, Rimg, disp = self.augmentor(Limg, Rimg, disp, self.name)

        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]

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


CREStereoDataset

    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

    def _build_cache(self):
        allpairs = [
            s + "/" + f[: -len("_left.jpg")]
            for s in sorted(os.listdir(self.root))
            for f in sorted(os.listdir(self.root + "/" + s))
            if f.endswith("_left.jpg")
        ]
        assert len(allpairs) == 200000, "incorrect parsing of pairs in CreStereo"
        tosave = {"train": allpairs}
        return tosave

class SceneFlowDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "SceneFlow"
        self._set_root()
        assert self.split in [
            "train_finalpass",
            "train_cleanpass",
            "train_allpass",
            "test_finalpass",
            "test_cleanpass",
            "test_allpass",
            "test1of100_cleanpass",
            "test1of100_finalpass",
        ]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname


SceneFlowDataset

            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

    def _build_cache(self):
        allpairs = [
            s + "/" + f[: -len("_left.jpg")]
            for s in sorted(os.listdir(self.root))
            for f in sorted(os.listdir(self.root + "/" + s))
            if f.endswith("_left.jpg")
        ]
        assert len(allpairs) == 200000, "incorrect parsing of pairs in CreStereo"
        tosave = {"train": allpairs}
        return tosave

class SceneFlowDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "SceneFlow"
        self._set_root()
        assert self.split in [
            "train_finalpass",
            "train_cleanpass",
            "train_allpass",
            "test_finalpass",
            "test_cleanpass",
            "test_allpass",
            "test1of100_cleanpass",
            "test1of100_finalpass",
        ]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname
        ).replace("/left/", "/right/")
        self.pairname_to_Ldispname = (
            lambda pairname: osp.join(self.root, pairname)
            .replace("/frames_finalpass/", "/disparity/")
            .replace("/frames_cleanpass/", "/disparity/")[:-4]
            + ".pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_sceneflow_disp

    def _build_cache(self):
        trainpairs = []
        # driving
        pairs = sorted(glob(self.root + "Driving/frames_finalpass/*/*/*/left/*.png"))
        pairs = list(map(lambda x: x[len(self.root) :], pairs))
        assert len(pairs) == 4400, "incorrect parsing of pairs in SceneFlow"
        trainpairs += pairs
        # monkaa
        pairs = sorted(glob(self.root + "Monkaa/frames_finalpass/*/left/*.png"))
        pairs = list(map(lambda x: x[len(self.root) :], pairs))
        assert len(pairs) == 8664, "incorrect parsing of pairs in SceneFlow"
        trainpairs += pairs
        # flyingthings
        pairs = sorted(
            glob(self.root + "FlyingThings/frames_finalpass/TRAIN/*/*/left/*.png")
        )
        pairs = list(map(lambda x: x[len(self.root) :], pairs))
        assert len(pairs) == 22390, "incorrect parsing of pairs in SceneFlow"
        trainpairs += pairs
        assert len(trainpairs) == 35454, "incorrect parsing of pairs in SceneFlow"
        testpairs = sorted(
            glob(self.root + "FlyingThings/frames_finalpass/TEST/*/*/left/*.png")
        )
        testpairs = list(map(lambda x: x[len(self.root) :], testpairs))
        assert len(testpairs) == 4370, "incorrect parsing of pairs in SceneFlow"
        test1of100pairs = testpairs[::100]
        assert len(test1of100pairs) == 44, "incorrect parsing of pairs in SceneFlow"
        # all
        tosave = {
            "train_finalpass": trainpairs,
            "train_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    trainpairs,
                )
            ),
            "test_finalpass": testpairs,
            "test_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    testpairs,
                )
            ),
            "test1of100_finalpass": test1of100pairs,
            "test1of100_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    test1of100pairs,
                )
            ),
        }
        tosave["train_allpass"] = tosave["train_finalpass"] + tosave["train_cleanpass"]
        tosave["test_allpass"] = tosave["test_finalpass"] + tosave["test_cleanpass"]
        return tosave

class Md21Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2021"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname.replace("/im0", "/im1")
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp0.pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury_disp

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))


            "test_finalpass": testpairs,
            "test_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    testpairs,
                )
            ),
            "test1of100_finalpass": test1of100pairs,
            "test1of100_cleanpass": list(
                map(
                    lambda x: x.replace("frames_finalpass", "frames_cleanpass"),
                    test1of100pairs,
                )
            ),
        }
        tosave["train_allpass"] = tosave["train_finalpass"] + tosave["train_cleanpass"]
        tosave["test_allpass"] = tosave["test_finalpass"] + tosave["test_cleanpass"]
        return tosave

class Md21Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2021"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname.replace("/im0", "/im1")
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp0.pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury_disp

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            # trainpairs += [s+'/im0.png'] # we should remove it, it is included as such in other lightings
            trainpairs += [
                s + "/ambient/" + b + "/" + a
                for b in sorted(os.listdir(osp.join(self.root, s, "ambient")))
                for a in sorted(os.listdir(osp.join(self.root, s, "ambient", b)))
                if a.startswith("im0")
            ]
        assert len(trainpairs) == 355
        subtrainpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in seqs[:-2])
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in seqs[-2:])
        ]
        assert (
            len(subtrainpairs) == 335 and len(subvalpairs) == 20
        ), "incorrect parsing of pairs in Middlebury 2021"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md14Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2014"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "disp0.pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury_disp
        self.has_constant_resolution = False

    def _build_cache(self):


            trainpairs += [
                s + "/ambient/" + b + "/" + a
                for b in sorted(os.listdir(osp.join(self.root, s, "ambient")))
                for a in sorted(os.listdir(osp.join(self.root, s, "ambient", b)))
                if a.startswith("im0")
            ]
        assert len(trainpairs) == 355
        subtrainpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in seqs[:-2])
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in seqs[-2:])
        ]
        assert (
            len(subtrainpairs) == 335 and len(subvalpairs) == 20
        ), "incorrect parsing of pairs in Middlebury 2021"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md14Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2014"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "disp0.pfm"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury_disp
        self.has_constant_resolution = False

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            trainpairs += [s + "/im1.png", s + "/im1E.png", s + "/im1L.png"]
        assert len(trainpairs) == 138
        valseqs = ["Umbrella-imperfect", "Vintage-perfect"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 132 and len(subvalpairs) == 6
        ), "incorrect parsing of pairs in Middlebury 2014"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md06Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2006"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "view5.png"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp1.png"
        )
        self.load_disparity = _read_middlebury20052006_disp
        self.has_constant_resolution = False

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))


        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            trainpairs += [s + "/im1.png", s + "/im1E.png", s + "/im1L.png"]
        assert len(trainpairs) == 138
        valseqs = ["Umbrella-imperfect", "Vintage-perfect"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 132 and len(subvalpairs) == 6
        ), "incorrect parsing of pairs in Middlebury 2014"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md06Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2006"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "view5.png"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp1.png"
        )
        self.load_disparity = _read_middlebury20052006_disp
        self.has_constant_resolution = False

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            for i in ["Illum1", "Illum2", "Illum3"]:
                for e in ["Exp0", "Exp1", "Exp2"]:
                    trainpairs.append(osp.join(s, i, e, "view1.png"))
        assert len(trainpairs) == 189
        valseqs = ["Rocks1", "Wood2"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 171 and len(subvalpairs) == 18
        ), "incorrect parsing of pairs in Middlebury 2006"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md05Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2005"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "view5.png"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp1.png"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury20052006_disp

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))


        for s in seqs:
            for i in ["Illum1", "Illum2", "Illum3"]:
                for e in ["Exp0", "Exp1", "Exp2"]:
                    trainpairs.append(osp.join(s, i, e, "view1.png"))
        assert len(trainpairs) == 189
        valseqs = ["Rocks1", "Wood2"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 171 and len(subvalpairs) == 18
        ), "incorrect parsing of pairs in Middlebury 2006"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class Md05Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Middlebury2005"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "view5.png"
        )
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, pairname.split("/")[0], "disp1.png"
        )
        self.pairname_to_str = lambda pairname: pairname[:-4]
        self.load_disparity = _read_middlebury20052006_disp

    def _build_cache(self):
        seqs = sorted(os.listdir(self.root))
        trainpairs = []
        for s in seqs:
            for i in ["Illum1", "Illum2", "Illum3"]:
                for e in ["Exp0", "Exp1", "Exp2"]:
                    trainpairs.append(osp.join(s, i, e, "view1.png"))
        assert len(trainpairs) == 54, "incorrect parsing of pairs in Middlebury 2005"
        valseqs = ["Reindeer"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 45 and len(subvalpairs) == 9
        ), "incorrect parsing of pairs in Middlebury 2005"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class MdEval3Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "MiddleburyEval3"
        self._set_root()
        assert self.split in [
            s + "_" + r
            for s in ["train", "subtrain", "subval", "test", "all"]
            for r in ["full", "half", "quarter"]
        ]
        if self.split.endswith("_full"):
            self.root = self.root.replace("/MiddEval3", "/MiddEval3_F")
        elif self.split.endswith("_half"):
            self.root = self.root.replace("/MiddEval3", "/MiddEval3_H")
        else:
            assert self.split.endswith("_quarter")
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname, "im0.png"


        for s in seqs:
            for i in ["Illum1", "Illum2", "Illum3"]:
                for e in ["Exp0", "Exp1", "Exp2"]:
                    trainpairs.append(osp.join(s, i, e, "view1.png"))
        assert len(trainpairs) == 54, "incorrect parsing of pairs in Middlebury 2005"
        valseqs = ["Reindeer"]
        assert all(s in seqs for s in valseqs)
        subtrainpairs = [
            p for p in trainpairs if not any(p.startswith(s + "/") for s in valseqs)
        ]
        subvalpairs = [
            p for p in trainpairs if any(p.startswith(s + "/") for s in valseqs)
        ]
        assert (
            len(subtrainpairs) == 45 and len(subvalpairs) == 9
        ), "incorrect parsing of pairs in Middlebury 2005"
        tosave = {"train": trainpairs, "subtrain": subtrainpairs, "subval": subvalpairs}
        return tosave

class MdEval3Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "MiddleburyEval3"
        self._set_root()
        assert self.split in [
            s + "_" + r
            for s in ["train", "subtrain", "subval", "test", "all"]
            for r in ["full", "half", "quarter"]
        ]
        if self.split.endswith("_full"):
            self.root = self.root.replace("/MiddEval3", "/MiddEval3_F")
        elif self.split.endswith("_half"):
            self.root = self.root.replace("/MiddEval3", "/MiddEval3_H")
        else:
            assert self.split.endswith("_quarter")
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname, "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname, "im1.png"
        )
        self.pairname_to_Ldispname = lambda pairname: (
            None
            if pairname.startswith("test")
            else osp.join(self.root, pairname, "disp0GT.pfm")
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_disparity = _read_middlebury_disp
        # for submission only
        self.submission_methodname = "CroCo-Stereo"
        self.submission_sresolution = (
            "F"
            if self.split.endswith("_full")
            else ("H" if self.split.endswith("_half") else "Q")
        )

    def _build_cache(self):
        trainpairs = ["train/" + s for s in sorted(os.listdir(self.root + "train/"))]
        testpairs = ["test/" + s for s in sorted(os.listdir(self.root + "test/"))]
        subvalpairs = trainpairs[-1:]
        subtrainpairs = trainpairs[:-1]
        allpairs = trainpairs + testpairs
        assert (
            len(trainpairs) == 15
            and len(testpairs) == 15
            and len(subvalpairs) == 1
            and len(subtrainpairs) == 14
            and len(allpairs) == 30
        ), "incorrect parsing of pairs in Middlebury Eval v3"
        tosave = {}
        for r in ["full", "half", "quarter"]:
            tosave.update(
                **{
                    "train_" + r: trainpairs,
                    "subtrain_" + r: subtrainpairs,
                    "subval_" + r: subvalpairs,
                    "test_" + r: testpairs,
                    "all_" + r: allpairs,
                }
            )
        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("/")[0].replace("train", "training")
            + self.submission_sresolution,
            pairname.split("/")[1],
            "disp0" + self.submission_methodname + ".pfm",
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writePFM(outfile, prediction)
        timefile = os.path.join(
            os.path.dirname(outfile), "time" + self.submission_methodname + ".txt"
        )
        with open(timefile, "w") as fid:
            fid.write(str(time))

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

class ETH3DLowResDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "ETH3DLowRes"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval", "all"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname, "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname, "im1.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: (
                None
                if pairname.startswith("test/")


            pairname.split("/")[0].replace("train", "training")
            + self.submission_sresolution,
            pairname.split("/")[1],
            "disp0" + self.submission_methodname + ".pfm",
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writePFM(outfile, prediction)
        timefile = os.path.join(
            os.path.dirname(outfile), "time" + self.submission_methodname + ".txt"
        )
        with open(timefile, "w") as fid:
            fid.write(str(time))

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

class ETH3DLowResDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "ETH3DLowRes"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval", "all"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname, "im0.png"
        )
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname, "im1.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: (
                None
                if pairname.startswith("test/")
                else osp.join(
                    self.root, pairname.replace("train/", "train_gt/"), "disp0GT.pfm"
                )
            )
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_disparity = _read_eth3d_disp
        self.has_constant_resolution = False

    def _build_cache(self):
        trainpairs = ["train/" + s for s in sorted(os.listdir(self.root + "train/"))]
        testpairs = ["test/" + s for s in sorted(os.listdir(self.root + "test/"))]
        assert (
            len(trainpairs) == 27 and len(testpairs) == 20
        ), "incorrect parsing of pairs in ETH3D Low Res"
        subvalpairs = [
            "train/delivery_area_3s",
            "train/electro_3l",
            "train/playground_3l",
        ]
        assert all(p in trainpairs for p in subvalpairs)
        subtrainpairs = [p for p in trainpairs if not p in subvalpairs]
        assert (
            len(subvalpairs) == 3 and len(subtrainpairs) == 24
        ), "incorrect parsing of pairs in ETH3D Low Res"
        tosave = {
            "train": trainpairs,
            "test": testpairs,
            "subtrain": subtrainpairs,
            "subval": subvalpairs,
            "all": trainpairs + testpairs,
        }
        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, "low_res_two_view", pairname.split("/")[1] + ".pfm"
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writePFM(outfile, prediction)
        timefile = outfile[:-4] + ".txt"
        with open(timefile, "w") as fid:
            fid.write("runtime " + str(time))

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

class BoosterDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Booster"
        self._set_root()
        assert self.split in [
            "train_balanced",
            "test_balanced",
            "subtrain_balanced",
            "subval_balanced",
        ]  # we use only the balanced version
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname
        ).replace("/camera_00/", "/camera_02/")
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "../disp_00.npy"
        )  # same images with different colors, same gt per sequence


BoosterDataset

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 2
        assert prediction.dtype == np.float32
        outfile = os.path.join(
            outdir, "low_res_two_view", pairname.split("/")[1] + ".pfm"
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writePFM(outfile, prediction)
        timefile = outfile[:-4] + ".txt"
        with open(timefile, "w") as fid:
            fid.write("runtime " + str(time))

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

class BoosterDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Booster"
        self._set_root()
        assert self.split in [
            "train_balanced",
            "test_balanced",
            "subtrain_balanced",
            "subval_balanced",
        ]  # we use only the balanced version
        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
        self.pairname_to_Rimgname = lambda pairname: osp.join(
            self.root, pairname
        ).replace("/camera_00/", "/camera_02/")
        self.pairname_to_Ldispname = lambda pairname: osp.join(
            self.root, osp.dirname(pairname), "../disp_00.npy"
        )  # same images with different colors, same gt per sequence
        self.pairname_to_str = lambda pairname: pairname[:-4].replace(
            "/camera_00/", "/"
        )
        self.load_disparity = _read_booster_disp

    def _build_cache(self):
        trainseqs = sorted(os.listdir(self.root + "train/balanced"))
        trainpairs = [
            "train/balanced/" + s + "/camera_00/" + imname
            for s in trainseqs
            for imname in sorted(
                os.listdir(self.root + "train/balanced/" + s + "/camera_00/")
            )
        ]
        testpairs = [
            "test/balanced/" + s + "/camera_00/" + imname
            for s in sorted(os.listdir(self.root + "test/balanced"))
            for imname in sorted(
                os.listdir(self.root + "test/balanced/" + s + "/camera_00/")
            )
        ]
        assert len(trainpairs) == 228 and len(testpairs) == 191
        subtrainpairs = [p for p in trainpairs if any(s in p for s in trainseqs[:-2])]
        subvalpairs = [p for p in trainpairs if any(s in p for s in trainseqs[-2:])]
        # warning: if we do validation split, we should split scenes!!!
        tosave = {
            "train_balanced": trainpairs,
            "test_balanced": testpairs,
            "subtrain_balanced": subtrainpairs,
            "subval_balanced": subvalpairs,
        }
        return tosave

class SpringDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Spring"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname + ".png"
        )
        self.pairname_to_Rimgname = (
            lambda pairname: osp.join(self.root, pairname + ".png")
            .replace("frame_right", "<frame_right>")
            .replace("frame_left", "frame_right")
            .replace("<frame_right>", "frame_left")
        )
        self.pairname_to_Ldispname = lambda pairname: (
            None
            if pairname.startswith("test")


SpringDataset

        testpairs = [
            "test/balanced/" + s + "/camera_00/" + imname
            for s in sorted(os.listdir(self.root + "test/balanced"))
            for imname in sorted(
                os.listdir(self.root + "test/balanced/" + s + "/camera_00/")
            )
        ]
        assert len(trainpairs) == 228 and len(testpairs) == 191
        subtrainpairs = [p for p in trainpairs if any(s in p for s in trainseqs[:-2])]
        subvalpairs = [p for p in trainpairs if any(s in p for s in trainseqs[-2:])]
        # warning: if we do validation split, we should split scenes!!!
        tosave = {
            "train_balanced": trainpairs,
            "test_balanced": testpairs,
            "subtrain_balanced": subtrainpairs,
            "subval_balanced": subvalpairs,
        }
        return tosave

class SpringDataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Spring"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval"]
        self.pairname_to_Limgname = lambda pairname: osp.join(
            self.root, pairname + ".png"
        )
        self.pairname_to_Rimgname = (
            lambda pairname: osp.join(self.root, pairname + ".png")
            .replace("frame_right", "<frame_right>")
            .replace("frame_left", "frame_right")
            .replace("<frame_right>", "frame_left")
        )
        self.pairname_to_Ldispname = lambda pairname: (
            None
            if pairname.startswith("test")
            else osp.join(self.root, pairname + ".dsp5")
            .replace("frame_left", "disp1_left")
            .replace("frame_right", "disp1_right")
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_disparity = _read_hdf5_disp

    def _build_cache(self):
        trainseqs = sorted(os.listdir(osp.join(self.root, "train")))
        trainpairs = [
            osp.join("train", s, "frame_left", f[:-4])
            for s in trainseqs
            for f in sorted(os.listdir(osp.join(self.root, "train", s, "frame_left")))
        ]
        testseqs = sorted(os.listdir(osp.join(self.root, "test")))
        testpairs = [
            osp.join("test", s, "frame_left", f[:-4])
            for s in testseqs
            for f in sorted(os.listdir(osp.join(self.root, "test", s, "frame_left")))
        ]
        testpairs += [p.replace("frame_left", "frame_right") for p in testpairs]
        """maxnorm = {'0001': 32.88, '0002': 228.5, '0004': 298.2, '0005': 142.5, '0006': 113.6, '0007': 27.3, '0008': 554.5, '0009': 155.6, '0010': 126.1, '0011': 87.6, '0012': 303.2, '0013': 24.14, '0014': 82.56, '0015': 98.44, '0016': 156.9, '0017': 28.17, '0018': 21.03, '0020': 178.0, '0021': 58.06, '0022': 354.2, '0023': 8.79, '0024': 97.06, '0025': 55.16, '0026': 91.9, '0027': 156.6, '0030': 200.4, '0032': 58.66, '0033': 373.5, '0036': 149.4, '0037': 5.625, '0038': 37.0, '0039': 12.2, '0041': 453.5, '0043': 457.0, '0044': 379.5, '0045': 161.8, '0047': 105.44} # => let'use 0041"""
        subtrainpairs = [p for p in trainpairs if p.split("/")[1] != "0041"]
        subvalpairs = [p for p in trainpairs if p.split("/")[1] == "0041"]
        assert (
            len(trainpairs) == 5000
            and len(testpairs) == 2000
            and len(subtrainpairs) == 4904
            and len(subvalpairs) == 96
        ), "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 == 2
        assert prediction.dtype == np.float32
        outfile = (
            os.path.join(outdir, pairname + ".dsp5")
            .replace("frame_left", "disp1_left")
            .replace("frame_right", "disp1_right")
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeDsp5File(prediction, outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        exe = "{self.root}/disp1_subsampling"
        if os.path.isfile(exe):
            cmd = f'cd "{outdir}/test"; {exe} .'
            print(cmd)
            os.system(cmd)
        else:
            print("Could not find disp1_subsampling executable for submission.")
            print("Please download it and run:")
            print(f'cd "{outdir}/test"; <disp1_subsampling_exe> .')

class Kitti12Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Kitti12"
        self._set_root()
        assert self.split in ["train", "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("/colored_0/", "/colored_1/") + "_10.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/colored_0/", "/disp_occ/") + "_10.png"
            )


            os.path.join(outdir, pairname + ".dsp5")
            .replace("frame_left", "disp1_left")
            .replace("frame_right", "disp1_right")
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeDsp5File(prediction, outfile)

    def finalize_submission(self, outdir):
        assert self.split == "test"
        exe = "{self.root}/disp1_subsampling"
        if os.path.isfile(exe):
            cmd = f'cd "{outdir}/test"; {exe} .'
            print(cmd)
            os.system(cmd)
        else:
            print("Could not find disp1_subsampling executable for submission.")
            print("Please download it and run:")
            print(f'cd "{outdir}/test"; <disp1_subsampling_exe> .')

class Kitti12Dataset(StereoDataset):

    def _prepare_data(self):
        self.name = "Kitti12"
        self._set_root()
        assert self.split in ["train", "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("/colored_0/", "/colored_1/") + "_10.png"
        )
        self.pairname_to_Ldispname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/colored_0/", "/disp_occ/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/colored_0/", "/")
        self.load_disparity = _read_kitti_disp

    def _build_cache(self):
        trainseqs = ["training/colored_0/%06d" % (i) for i in range(194)]
        testseqs = ["testing/colored_0/%06d" % (i) for i in range(195)]
        assert (
            len(trainseqs) == 194 and len(testseqs) == 195
        ), "incorrect parsing of pairs in Kitti12"
        tosave = {"train": trainseqs, "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, 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"
            )


        ), "incorrect parsing of pairs in Kitti12"
        tosave = {"train": trainseqs, "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, 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 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


    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
    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm


        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
    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)


        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
    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)

def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)


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
    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)

def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)

def _read_kitti_disp(filename):
    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp


    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
    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)

def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)

def _read_kitti_disp(filename):
    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp
_read_middlebury_disp = _read_pfm_disp
_read_carla_disp = _read_pfm_disp
_read_tartanair_disp = _read_npy_disp


def _read_png_disp(filename, coef=1.0):
    disp = np.asarray(Image.open(filename))
    disp = disp.astype(np.float32) / coef
    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)

def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)

def _read_kitti_disp(filename):
    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp
_read_middlebury_disp = _read_pfm_disp
_read_carla_disp = _read_pfm_disp
_read_tartanair_disp = _read_npy_disp

def _read_hdf5_disp(filename):
    disp = np.asarray(h5py.File(filename)["disparity"])
    disp[np.isnan(disp)] = np.inf  # make invalid values as +inf


    disp[disp == 0.0] = np.inf
    return disp

def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)

def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)

def _read_kitti_disp(filename):
    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp
_read_middlebury_disp = _read_pfm_disp
_read_carla_disp = _read_pfm_disp
_read_tartanair_disp = _read_npy_disp

def _read_hdf5_disp(filename):
    disp = np.asarray(h5py.File(filename)["disparity"])
    disp[np.isnan(disp)] = np.inf  # make invalid values as +inf
    # disp[disp==0.0] = np.inf # make invalid values as +inf
    return disp.astype(np.float32)


def _read_pfm_disp(filename):
    disp = np.ascontiguousarray(_read_pfm(filename)[0])
    disp[disp <= 0] = (
        np.inf
    )  # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
    return disp

def _read_npy_disp(filename):
    return np.load(filename)

def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)

def _read_kitti_disp(filename):
    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp
_read_middlebury_disp = _read_pfm_disp
_read_carla_disp = _read_pfm_disp
_read_tartanair_disp = _read_npy_disp

def _read_hdf5_disp(filename):
    disp = np.asarray(h5py.File(filename)["disparity"])
    disp[np.isnan(disp)] = np.inf  # make invalid values as +inf
    # disp[disp==0.0] = np.inf # make invalid values as +inf
    return disp.astype(np.float32)

import re

def _read_pfm(file):


def _read_crestereo_disp(filename):
    return _read_png_disp(filename, coef=32.0)

def _read_middlebury20052006_disp(filename):
    return _read_png_disp(filename, coef=1.0)

def _read_kitti_disp(filename):
    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp
_read_middlebury_disp = _read_pfm_disp
_read_carla_disp = _read_pfm_disp
_read_tartanair_disp = _read_npy_disp

def _read_hdf5_disp(filename):
    disp = np.asarray(h5py.File(filename)["disparity"])
    disp[np.isnan(disp)] = np.inf  # make invalid values as +inf
    # disp[disp==0.0] = np.inf # make invalid values as +inf
    return disp.astype(np.float32)

import re

def _read_pfm(file):
    file = open(file, "rb")

    color = None
    width = None
    height = None
    scale = None
    endian = None

    header = file.readline().rstrip()
    if header.decode("ascii") == "PF":
        color = True
    elif header.decode("ascii") == "Pf":
        color = False
    else:


    return _read_png_disp(filename, coef=256.0)

_read_sceneflow_disp = _read_pfm_disp
_read_eth3d_disp = _read_pfm_disp
_read_middlebury_disp = _read_pfm_disp
_read_carla_disp = _read_pfm_disp
_read_tartanair_disp = _read_npy_disp

def _read_hdf5_disp(filename):
    disp = np.asarray(h5py.File(filename)["disparity"])
    disp[np.isnan(disp)] = np.inf  # make invalid values as +inf
    # disp[disp==0.0] = np.inf # make invalid values as +inf
    return disp.astype(np.float32)

import re

def _read_pfm(file):
    file = open(file, "rb")

    color = None
    width = None
    height = None
    scale = None
    endian = None

    header = file.readline().rstrip()
    if header.decode("ascii") == "PF":
        color = True
    elif header.decode("ascii") == "Pf":
        color = False
    else:
        raise Exception("Not a PFM file.")

    dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
    if dim_match:
        width, height = list(map(int, dim_match.groups()))
    else:
        raise Exception("Malformed PFM header.")

    scale = float(file.readline().decode("ascii").rstrip())
    if scale < 0:  # little-endian
        endian = "<"
        scale = -scale
    else:
        endian = ">"  # big-endian

    data = np.fromfile(file, endian + "f")
    shape = (height, width, 3) if color else (height, width)

    data = np.reshape(data, shape)
    data = np.flipud(data)
    return data, scale

def writePFM(file, image, scale=1):
    file = open(file, "wb")

    color = None

    if image.dtype.name != "float32":
        raise Exception("Image dtype must be float32.")

    image = np.flipud(image)

    if len(image.shape) == 3 and image.shape[2] == 3:  # color image
        color = True
    elif (
        len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
    ):  # greyscale
        color = False
    else:
        raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")


writePFM

    if dim_match:
        width, height = list(map(int, dim_match.groups()))
    else:
        raise Exception("Malformed PFM header.")

    scale = float(file.readline().decode("ascii").rstrip())
    if scale < 0:  # little-endian
        endian = "<"
        scale = -scale
    else:
        endian = ">"  # big-endian

    data = np.fromfile(file, endian + "f")
    shape = (height, width, 3) if color else (height, width)

    data = np.reshape(data, shape)
    data = np.flipud(data)
    return data, scale

def writePFM(file, image, scale=1):
    file = open(file, "wb")

    color = None

    if image.dtype.name != "float32":
        raise Exception("Image dtype must be float32.")

    image = np.flipud(image)

    if len(image.shape) == 3 and image.shape[2] == 3:  # color image
        color = True
    elif (
        len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
    ):  # greyscale
        color = False
    else:
        raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")

    file.write("PF\n" if color else "Pf\n".encode())
    file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))

    endian = image.dtype.byteorder

    if endian == "<" or endian == "=" and sys.byteorder == "little":
        scale = -scale

    file.write("%f\n".encode() % scale)

    image.tofile(file)

def writeDsp5File(disp, filename):
    with h5py.File(filename, "w") as f:
        f.create_dataset("disparity", data=disp, compression="gzip", compression_opts=5)

# disp visualization

def vis_disparity(disp, m=None, M=None):
    if m is None:
        m = disp.min()
    if M is None:
        M = disp.max()
    disp_vis = (disp - m) / (M - m) * 255.0
    disp_vis = disp_vis.astype("uint8")
    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
    return disp_vis


    elif (
        len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1
    ):  # greyscale
        color = False
    else:
        raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")

    file.write("PF\n" if color else "Pf\n".encode())
    file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))

    endian = image.dtype.byteorder

    if endian == "<" or endian == "=" and sys.byteorder == "little":
        scale = -scale

    file.write("%f\n".encode() % scale)

    image.tofile(file)

def writeDsp5File(disp, filename):
    with h5py.File(filename, "w") as f:
        f.create_dataset("disparity", data=disp, compression="gzip", compression_opts=5)

# disp visualization

def vis_disparity(disp, m=None, M=None):
    if m is None:
        m = disp.min()
    if M is None:
        M = disp.max()
    disp_vis = (disp - m) / (M - m) * 255.0
    disp_vis = disp_vis.astype("uint8")
    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
    return disp_vis

# dataset getter

def get_train_dataset_stereo(dataset_str, augmentor=True, crop_size=None):


    file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))

    endian = image.dtype.byteorder

    if endian == "<" or endian == "=" and sys.byteorder == "little":
        scale = -scale

    file.write("%f\n".encode() % scale)

    image.tofile(file)

def writeDsp5File(disp, filename):
    with h5py.File(filename, "w") as f:
        f.create_dataset("disparity", data=disp, compression="gzip", compression_opts=5)

# disp visualization

def vis_disparity(disp, m=None, M=None):
    if m is None:
        m = disp.min()
    if M is None:
        M = disp.max()
    disp_vis = (disp - m) / (M - m) * 255.0
    disp_vis = disp_vis.astype("uint8")
    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
    return disp_vis

# dataset getter

def get_train_dataset_stereo(dataset_str, augmentor=True, crop_size=None):
    dataset_str = dataset_str.replace("(", "Dataset(")
    if augmentor:
        dataset_str = dataset_str.replace(")", ", augmentor=True)")
    if crop_size is not None:
        dataset_str = dataset_str.replace(
            ")", ", crop_size={:s})".format(str(crop_size))
        )
    return eval(dataset_str)

def get_test_datasets_stereo(dataset_str):
    dataset_str = dataset_str.replace("(", "Dataset(")
    return [eval(s) for s in dataset_str.split("+")]


        f.create_dataset("disparity", data=disp, compression="gzip", compression_opts=5)

# disp visualization

def vis_disparity(disp, m=None, M=None):
    if m is None:
        m = disp.min()
    if M is None:
        M = disp.max()
    disp_vis = (disp - m) / (M - m) * 255.0
    disp_vis = disp_vis.astype("uint8")
    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
    return disp_vis

# dataset getter

def get_train_dataset_stereo(dataset_str, augmentor=True, crop_size=None):
    dataset_str = dataset_str.replace("(", "Dataset(")
    if augmentor:
        dataset_str = dataset_str.replace(")", ", augmentor=True)")
    if crop_size is not None:
        dataset_str = dataset_str.replace(
            ")", ", crop_size={:s})".format(str(crop_size))
        )
    return eval(dataset_str)

def get_test_datasets_stereo(dataset_str):
    dataset_str = dataset_str.replace("(", "Dataset(")
    return [eval(s) for s in dataset_str.split("+")]


    disp_vis = (disp - m) / (M - m) * 255.0
    disp_vis = disp_vis.astype("uint8")
    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
    return disp_vis

# dataset getter

def get_train_dataset_stereo(dataset_str, augmentor=True, crop_size=None):
    dataset_str = dataset_str.replace("(", "Dataset(")
    if augmentor:
        dataset_str = dataset_str.replace(")", ", augmentor=True)")
    if crop_size is not None:
        dataset_str = dataset_str.replace(
            ")", ", crop_size={:s})".format(str(crop_size))
        )
    return eval(dataset_str)

def get_test_datasets_stereo(dataset_str):
    dataset_str = dataset_str.replace("(", "Dataset(")
    return [eval(s) for s in dataset_str.split("+")]


}
cache_dir = "./data/stereoflow/datasets_stereo_cache/"

in1k_mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
in1k_std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)

def img_to_tensor(img):
    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
    img = (img - in1k_mean) / in1k_std
    return img

def disp_to_tensor(disp):
    return torch.from_numpy(disp)[None, :, :]

class StereoDataset(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:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)

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

        # get filenames
        Limgname = self.pairname_to_Limgname(pairname)
        Rimgname = self.pairname_to_Rimgname(pairname)
        Ldispname = (
            self.pairname_to_Ldispname(pairname)
            if self.pairname_to_Ldispname is not None
            else None


    return torch.from_numpy(disp)[None, :, :]

class StereoDataset(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:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)

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

        # get filenames
        Limgname = self.pairname_to_Limgname(pairname)
        Rimgname = self.pairname_to_Rimgname(pairname)
        Ldispname = (
            self.pairname_to_Ldispname(pairname)
            if self.pairname_to_Ldispname is not None
            else None
        )

        # load images and disparities
        Limg = _read_img(Limgname)
        Rimg = _read_img(Rimgname)
        disp = self.load_disparity(Ldispname) if Ldispname is not None else None


        self.split = split
        if not augmentor:
            assert crop_size is None
        if crop_size:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)

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

        # get filenames
        Limgname = self.pairname_to_Limgname(pairname)
        Rimgname = self.pairname_to_Rimgname(pairname)
        Ldispname = (
            self.pairname_to_Ldispname(pairname)
            if self.pairname_to_Ldispname is not None
            else None
        )

        # load images and disparities
        Limg = _read_img(Limgname)
        Rimg = _read_img(Rimgname)
        disp = self.load_disparity(Ldispname) if Ldispname is not None else None

        # sanity check
        if disp is not None:


        if crop_size:
            assert augmentor
        self.crop_size = crop_size
        self.augmentor_str = augmentor
        self.augmentor = StereoAugmentor(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)

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

        # get filenames
        Limgname = self.pairname_to_Limgname(pairname)
        Rimgname = self.pairname_to_Rimgname(pairname)
        Ldispname = (
            self.pairname_to_Ldispname(pairname)
            if self.pairname_to_Ldispname is not None
            else None
        )

        # load images and disparities
        Limg = _read_img(Limgname)
        Rimg = _read_img(Rimgname)
        disp = self.load_disparity(Ldispname) if Ldispname is not None else None

        # sanity check
        if disp is not None:
            assert np.all(disp > 0) or self.name == "Spring", (
                self.name,
                pairname,
                Ldispname,
            )

        # apply augmentations
        if self.augmentor is not None:
            Limg, Rimg, disp = self.augmentor(Limg, Rimg, disp, self.name)

        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(


                pairname,
                Ldispname,
            )

        # apply augmentations
        if self.augmentor is not None:
            Limg, Rimg, disp = self.augmentor(Limg, Rimg, disp, self.name)

        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 self.augmentor is not None:
            Limg, Rimg, disp = self.augmentor(Limg, Rimg, disp, self.name)

        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]