TartanAirDataset

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

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

class TartanAirDataset(FlowDataset):

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

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

class FlyingChairsDataset(FlowDataset):

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


FlyingChairsDataset

            "{:06d}_{:06d}".format(pairname[1], pairname[2]),
        )
        self.load_flow = _read_numpy_flow

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

class FlyingChairsDataset(FlowDataset):

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

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

class FlyingThingsDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "FlyingThings"
        self._set_root()
        assert self.split in [
            f"{set_}_{pass_}pass{camstr}"
            for set_ in ["train", "test", "test1024"]
            for camstr in ["", "_rightcam"]
            for pass_ in ["clean", "final", "all"]
        ]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root,
            f"frames_{pairname[3]}pass",
            pairname[0].replace("into_future", "").replace("into_past", ""),
            "{:04d}.png".format(pairname[1]),
        )
        self.pairname_to_img2name = lambda pairname: osp.join(


FlyingThingsDataset

            self.root, "data", pairname + "_img2.ppm"
        )
        self.pairname_to_flowname = lambda pairname: osp.join(
            self.root, "data", pairname + "_flow.flo"
        )
        self.pairname_to_str = lambda pairname: pairname
        self.load_flow = _read_flo_file

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

class FlyingThingsDataset(FlowDataset):

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

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

class MPISintelDataset(FlowDataset):

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


MPISintelDataset

        stride = alltest_nsamples // test1024_nsamples
        remove = alltest_nsamples % test1024_nsamples
        for cam in ["left", "right"]:
            camstr = "" if cam == "left" else f"_{cam}cam"
            for pass_ in ["final", "clean"]:
                tosave[f"test1024_{pass_}pass{camstr}"] = sorted(
                    tosave[f"test_{pass_}pass{camstr}"]
                )[:-remove][
                    ::stride
                ]  # warning, it was not sorted before
            assert (
                len(tosave["test1024_cleanpass"]) == 1024
            ), "incorrect parsing of pairs in Flying Things"
            tosave[f"test1024_allpass{camstr}"] = (
                tosave[f"test1024_cleanpass{camstr}"]
                + tosave[f"test1024_finalpass{camstr}"]
            )
        return tosave

class MPISintelDataset(FlowDataset):

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

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

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

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

class SpringDataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Spring"
        self._set_root()
        assert self.split in ["train", "test", "subtrain", "subval"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root,
            pairname[0],
            pairname[1],
            "frame_" + pairname[3],
            "frame_{:s}_{:04d}.png".format(pairname[3], pairname[4]),
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root,
            pairname[0],
            pairname[1],
            "frame_" + pairname[3],


SpringDataset

        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowFile(prediction, outfile)

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

class SpringDataset(FlowDataset):

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

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

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

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

class Kitti12Dataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Kitti12"
        self._set_root()
        assert self.split in ["train", "test"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname + "_11.png"
        )
        self.pairname_to_flowname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/colored_0/", "/flow_occ/") + "_10.png"
            )


            f"flow_{pairname[2]}_{pairname[3]}",
            f"flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}.flo5",
        )
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlo5File(prediction, outfile)

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

class Kitti12Dataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Kitti12"
        self._set_root()
        assert self.split in ["train", "test"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname + "_11.png"
        )
        self.pairname_to_flowname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/colored_0/", "/flow_occ/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/colored_0/", "/")
        self.load_flow = _read_kitti_flow

    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 == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

class Kitti15Dataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Kitti15"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval", "test"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname + "_11.png"
        )
        self.pairname_to_flowname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/image_2/", "/flow_occ/") + "_10.png"
            )


            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 == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

class Kitti15Dataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Kitti15"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval", "test"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname + "_11.png"
        )
        self.pairname_to_flowname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/image_2/", "/flow_occ/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/image_2/", "/")
        self.load_flow = _read_kitti_flow

    def _build_cache(self):
        trainseqs = ["training/image_2/%06d" % (i) for i in range(200)]
        subtrainseqs = trainseqs[:-10]
        subvalseqs = trainseqs[-10:]
        testseqs = ["testing/image_2/%06d" % (i) for i in range(200)]
        assert (
            len(trainseqs) == 200
            and len(subtrainseqs) == 190
            and len(subvalseqs) == 10
            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 == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, "flow", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

import cv2

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

def _read_pfm_flow(filename):
    f, _ = _read_pfm(filename)
    assert np.all(f[:, :, 2] == 0.0)
    return np.ascontiguousarray(f[:, :, :2])

TAG_FLOAT = 202021.25  # tag to check the sanity of the file
TAG_STRING = "PIEH"  # string containing the tag
MIN_WIDTH = 1
MAX_WIDTH = 99999
MIN_HEIGHT = 1


        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, "flow", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

import cv2

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

def _read_pfm_flow(filename):
    f, _ = _read_pfm(filename)
    assert np.all(f[:, :, 2] == 0.0)
    return np.ascontiguousarray(f[:, :, :2])

TAG_FLOAT = 202021.25  # tag to check the sanity of the file
TAG_STRING = "PIEH"  # string containing the tag
MIN_WIDTH = 1
MAX_WIDTH = 99999
MIN_HEIGHT = 1
MAX_HEIGHT = 99999

def readFlowFile(filename):
    """
    readFlowFile(<FILENAME>) reads a flow file <FILENAME> into a 2-band np.array.
    if <FILENAME> does not exist, an IOError is raised.


        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, "flow", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

import cv2

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

def _read_pfm_flow(filename):
    f, _ = _read_pfm(filename)
    assert np.all(f[:, :, 2] == 0.0)
    return np.ascontiguousarray(f[:, :, :2])

TAG_FLOAT = 202021.25  # tag to check the sanity of the file
TAG_STRING = "PIEH"  # string containing the tag
MIN_WIDTH = 1
MAX_WIDTH = 99999
MIN_HEIGHT = 1
MAX_HEIGHT = 99999

def readFlowFile(filename):
    """
    readFlowFile(<FILENAME>) reads a flow file <FILENAME> into a 2-band np.array.
    if <FILENAME> does not exist, an IOError is raised.
    if <FILENAME> does not finish by '.flo' or the tag, the width, the height or the file's size is illegal, an Expcetion is raised.
    ---- PARAMETERS ----
        filename: string containg the name of the file to read a flow
    ---- OUTPUTS ----
        a np.array of dimension (height x width x 2) containing the flow of type 'float32'
    """


readFlowFile

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

def _read_pfm_flow(filename):
    f, _ = _read_pfm(filename)
    assert np.all(f[:, :, 2] == 0.0)
    return np.ascontiguousarray(f[:, :, :2])

TAG_FLOAT = 202021.25  # tag to check the sanity of the file
TAG_STRING = "PIEH"  # string containing the tag
MIN_WIDTH = 1
MAX_WIDTH = 99999
MIN_HEIGHT = 1
MAX_HEIGHT = 99999

def readFlowFile(filename):
    """
    readFlowFile(<FILENAME>) reads a flow file <FILENAME> into a 2-band np.array.
    if <FILENAME> does not exist, an IOError is raised.
    if <FILENAME> does not finish by '.flo' or the tag, the width, the height or the file's size is illegal, an Expcetion is raised.
    ---- PARAMETERS ----
        filename: string containg the name of the file to read a flow
    ---- OUTPUTS ----
        a np.array of dimension (height x width x 2) containing the flow of type 'float32'
    """

    # check filename
    if not filename.endswith(".flo"):
        raise Exception(
            "readFlowFile({:s}): filename must finish with '.flo'".format(filename)
        )

    # open the file and read it
    with open(filename, "rb") as f:
        # check tag
        tag = struct.unpack("f", f.read(4))[0]
        if tag != TAG_FLOAT:
            raise Exception("flow_utils.readFlowFile({:s}): wrong tag".format(filename))
        # read dimension
        w, h = struct.unpack("ii", f.read(8))
        if w < MIN_WIDTH or w > MAX_WIDTH:
            raise Exception(
                "flow_utils.readFlowFile({:s}: illegal width {:d}".format(filename, w)
            )
        if h < MIN_HEIGHT or h > MAX_HEIGHT:
            raise Exception(
                "flow_utils.readFlowFile({:s}: illegal height {:d}".format(filename, h)
            )
        flow = np.fromfile(f, "float32")
        if not flow.shape == (h * w * 2,):
            raise Exception(
                "flow_utils.readFlowFile({:s}: illegal size of the file".format(
                    filename
                )
            )
        flow.shape = (h, w, 2)
        return flow

def writeFlowFile(flow, filename):
    """
    writeFlowFile(flow,<FILENAME>) write flow to the file <FILENAME>.
    if <FILENAME> does not exist, an IOError is raised.
    if <FILENAME> does not finish with '.flo' or the flow has not 2 bands, an Exception is raised.
    ---- PARAMETERS ----
        flow: np.array of dimension (height x width x 2) containing the flow to write
        filename: string containg the name of the file to write a flow
    """

    # check filename
    if not filename.endswith(".flo"):
        raise Exception(
            "flow_utils.writeFlowFile(<flow>,{:s}): filename must finish with '.flo'".format(
                filename
            )
        )


writeFlowFile

        w, h = struct.unpack("ii", f.read(8))
        if w < MIN_WIDTH or w > MAX_WIDTH:
            raise Exception(
                "flow_utils.readFlowFile({:s}: illegal width {:d}".format(filename, w)
            )
        if h < MIN_HEIGHT or h > MAX_HEIGHT:
            raise Exception(
                "flow_utils.readFlowFile({:s}: illegal height {:d}".format(filename, h)
            )
        flow = np.fromfile(f, "float32")
        if not flow.shape == (h * w * 2,):
            raise Exception(
                "flow_utils.readFlowFile({:s}: illegal size of the file".format(
                    filename
                )
            )
        flow.shape = (h, w, 2)
        return flow

def writeFlowFile(flow, filename):
    """
    writeFlowFile(flow,<FILENAME>) write flow to the file <FILENAME>.
    if <FILENAME> does not exist, an IOError is raised.
    if <FILENAME> does not finish with '.flo' or the flow has not 2 bands, an Exception is raised.
    ---- PARAMETERS ----
        flow: np.array of dimension (height x width x 2) containing the flow to write
        filename: string containg the name of the file to write a flow
    """

    # check filename
    if not filename.endswith(".flo"):
        raise Exception(
            "flow_utils.writeFlowFile(<flow>,{:s}): filename must finish with '.flo'".format(
                filename
            )
        )

    if not flow.shape[2:] == (2,):
        raise Exception(
            "flow_utils.writeFlowFile(<flow>,{:s}): <flow> must have 2 bands".format(
                filename
            )
        )

    # open the file and write it
    with open(filename, "wb") as f:
        # write TAG
        f.write(TAG_STRING.encode("utf-8"))
        # write dimension
        f.write(struct.pack("ii", flow.shape[1], flow.shape[0]))
        # write the flow

        flow.astype(np.float32).tofile(f)

_read_flo_file = readFlowFile

def _read_kitti_flow(filename):
    flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR)
    flow = flow[:, :, ::-1].astype(np.float32)
    valid = flow[:, :, 2] > 0
    flow = flow[:, :, :2]
    flow = (flow - 2**15) / 64.0
    flow[~valid, 0] = np.inf
    flow[~valid, 1] = np.inf
    return flow

_read_hd1k_flow = _read_kitti_flow

def writeFlowKitti(filename, uv):


        raise Exception(
            "flow_utils.writeFlowFile(<flow>,{:s}): <flow> must have 2 bands".format(
                filename
            )
        )

    # open the file and write it
    with open(filename, "wb") as f:
        # write TAG
        f.write(TAG_STRING.encode("utf-8"))
        # write dimension
        f.write(struct.pack("ii", flow.shape[1], flow.shape[0]))
        # write the flow

        flow.astype(np.float32).tofile(f)

_read_flo_file = readFlowFile

def _read_kitti_flow(filename):
    flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR)
    flow = flow[:, :, ::-1].astype(np.float32)
    valid = flow[:, :, 2] > 0
    flow = flow[:, :, :2]
    flow = (flow - 2**15) / 64.0
    flow[~valid, 0] = np.inf
    flow[~valid, 1] = np.inf
    return flow

_read_hd1k_flow = _read_kitti_flow

def writeFlowKitti(filename, uv):
    uv = 64.0 * uv + 2**15
    valid = np.ones([uv.shape[0], uv.shape[1], 1])
    uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
    cv2.imwrite(filename, uv[..., ::-1])

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

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


writeFlowKitti

        flow.astype(np.float32).tofile(f)

_read_flo_file = readFlowFile

def _read_kitti_flow(filename):
    flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR)
    flow = flow[:, :, ::-1].astype(np.float32)
    valid = flow[:, :, 2] > 0
    flow = flow[:, :, :2]
    flow = (flow - 2**15) / 64.0
    flow[~valid, 0] = np.inf
    flow[~valid, 1] = np.inf
    return flow

_read_hd1k_flow = _read_kitti_flow

def writeFlowKitti(filename, uv):
    uv = 64.0 * uv + 2**15
    valid = np.ones([uv.shape[0], uv.shape[1], 1])
    uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
    cv2.imwrite(filename, uv[..., ::-1])

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

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

# flow visualization
RY = 15
YG = 6
GC = 4
CB = 11
BM = 13
MR = 6


    flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH | cv2.IMREAD_COLOR)
    flow = flow[:, :, ::-1].astype(np.float32)
    valid = flow[:, :, 2] > 0
    flow = flow[:, :, :2]
    flow = (flow - 2**15) / 64.0
    flow[~valid, 0] = np.inf
    flow[~valid, 1] = np.inf
    return flow

_read_hd1k_flow = _read_kitti_flow

def writeFlowKitti(filename, uv):
    uv = 64.0 * uv + 2**15
    valid = np.ones([uv.shape[0], uv.shape[1], 1])
    uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
    cv2.imwrite(filename, uv[..., ::-1])

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

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

# flow visualization
RY = 15
YG = 6
GC = 4
CB = 11
BM = 13
MR = 6
UNKNOWN_THRESH = 1e9

def colorTest():
    """


    flow[~valid, 0] = np.inf
    flow[~valid, 1] = np.inf
    return flow

_read_hd1k_flow = _read_kitti_flow

def writeFlowKitti(filename, uv):
    uv = 64.0 * uv + 2**15
    valid = np.ones([uv.shape[0], uv.shape[1], 1])
    uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
    cv2.imwrite(filename, uv[..., ::-1])

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

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

# flow visualization
RY = 15
YG = 6
GC = 4
CB = 11
BM = 13
MR = 6
UNKNOWN_THRESH = 1e9

def colorTest():
    """
    flow_utils.colorTest(): display an example of image showing the color encoding scheme
    """
    import matplotlib.pylab as plt

    truerange = 1
    h, w = 151, 151


colorTest

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

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

# flow visualization
RY = 15
YG = 6
GC = 4
CB = 11
BM = 13
MR = 6
UNKNOWN_THRESH = 1e9

def colorTest():
    """
    flow_utils.colorTest(): display an example of image showing the color encoding scheme
    """
    import matplotlib.pylab as plt

    truerange = 1
    h, w = 151, 151
    trange = truerange * 1.04
    s2 = round(h / 2)
    x, y = np.meshgrid(range(w), range(h))
    u = x * trange / s2 - trange
    v = y * trange / s2 - trange
    img = _computeColor(
        np.concatenate((u[:, :, np.newaxis], v[:, :, np.newaxis]), 2)
        / trange
        / np.sqrt(2)
    )
    plt.imshow(img)
    plt.axis("off")
    plt.axhline(round(h / 2), color="k")
    plt.axvline(round(w / 2), color="k")

def flowToColor(flow, maxflow=None, maxmaxflow=None, saturate=False):
    """
    flow_utils.flowToColor(flow): return a color code flow field, normalized based on the maximum l2-norm of the flow
    flow_utils.flowToColor(flow,maxflow): return a color code flow field, normalized by maxflow
    ---- PARAMETERS ----
        flow: flow to display of shape (height x width x 2)
        maxflow (default:None): if given, normalize the flow by its value, otherwise by the flow norm
        maxmaxflow (default:None): if given, normalize the flow by the max of its value and the flow norm
    ---- OUTPUT ----
        an np.array of shape (height x width x 3) of type uint8 containing a color code of the flow
    """
    h, w, n = flow.shape
    # check size of flow
    assert n == 2, "flow_utils.flowToColor(flow): flow must have 2 bands"
    # fix unknown flow
    unknown_idx = np.max(np.abs(flow), 2) > UNKNOWN_THRESH
    flow[unknown_idx] = 0.0
    # compute max flow if needed


flowToColor

    import matplotlib.pylab as plt

    truerange = 1
    h, w = 151, 151
    trange = truerange * 1.04
    s2 = round(h / 2)
    x, y = np.meshgrid(range(w), range(h))
    u = x * trange / s2 - trange
    v = y * trange / s2 - trange
    img = _computeColor(
        np.concatenate((u[:, :, np.newaxis], v[:, :, np.newaxis]), 2)
        / trange
        / np.sqrt(2)
    )
    plt.imshow(img)
    plt.axis("off")
    plt.axhline(round(h / 2), color="k")
    plt.axvline(round(w / 2), color="k")

def flowToColor(flow, maxflow=None, maxmaxflow=None, saturate=False):
    """
    flow_utils.flowToColor(flow): return a color code flow field, normalized based on the maximum l2-norm of the flow
    flow_utils.flowToColor(flow,maxflow): return a color code flow field, normalized by maxflow
    ---- PARAMETERS ----
        flow: flow to display of shape (height x width x 2)
        maxflow (default:None): if given, normalize the flow by its value, otherwise by the flow norm
        maxmaxflow (default:None): if given, normalize the flow by the max of its value and the flow norm
    ---- OUTPUT ----
        an np.array of shape (height x width x 3) of type uint8 containing a color code of the flow
    """
    h, w, n = flow.shape
    # check size of flow
    assert n == 2, "flow_utils.flowToColor(flow): flow must have 2 bands"
    # fix unknown flow
    unknown_idx = np.max(np.abs(flow), 2) > UNKNOWN_THRESH
    flow[unknown_idx] = 0.0
    # compute max flow if needed
    if maxflow is None:
        maxflow = flowMaxNorm(flow)
    if maxmaxflow is not None:
        maxflow = min(maxmaxflow, maxflow)
    # normalize flow
    eps = np.spacing(1)  # minimum positive float value to avoid division by 0
    # compute the flow
    img = _computeColor(flow / (maxflow + eps), saturate=saturate)
    # put black pixels in unknown location
    img[np.tile(unknown_idx[:, :, np.newaxis], [1, 1, 3])] = 0.0
    return img

def flowMaxNorm(flow):
    """
    flow_utils.flowMaxNorm(flow): return the maximum of the l2-norm of the given flow
    ---- PARAMETERS ----
        flow: the flow

    ---- OUTPUT ----
        a float containing the maximum of the l2-norm of the flow
    """
    return np.max(np.sqrt(np.sum(np.square(flow), 2)))

def _computeColor(flow, saturate=True):
    """
    flow_utils._computeColor(flow): compute color codes for the flow field flow

    ---- PARAMETERS ----
        flow: np.array of dimension (height x width x 2) containing the flow to display


flowMaxNorm

    h, w, n = flow.shape
    # check size of flow
    assert n == 2, "flow_utils.flowToColor(flow): flow must have 2 bands"
    # fix unknown flow
    unknown_idx = np.max(np.abs(flow), 2) > UNKNOWN_THRESH
    flow[unknown_idx] = 0.0
    # compute max flow if needed
    if maxflow is None:
        maxflow = flowMaxNorm(flow)
    if maxmaxflow is not None:
        maxflow = min(maxmaxflow, maxflow)
    # normalize flow
    eps = np.spacing(1)  # minimum positive float value to avoid division by 0
    # compute the flow
    img = _computeColor(flow / (maxflow + eps), saturate=saturate)
    # put black pixels in unknown location
    img[np.tile(unknown_idx[:, :, np.newaxis], [1, 1, 3])] = 0.0
    return img

def flowMaxNorm(flow):
    """
    flow_utils.flowMaxNorm(flow): return the maximum of the l2-norm of the given flow
    ---- PARAMETERS ----
        flow: the flow

    ---- OUTPUT ----
        a float containing the maximum of the l2-norm of the flow
    """
    return np.max(np.sqrt(np.sum(np.square(flow), 2)))

def _computeColor(flow, saturate=True):
    """
    flow_utils._computeColor(flow): compute color codes for the flow field flow

    ---- PARAMETERS ----
        flow: np.array of dimension (height x width x 2) containing the flow to display
    ---- OUTPUTS ----
        an np.array of dimension (height x width x 3) containing the color conversion of the flow
    """
    # set nan to 0
    nanidx = np.isnan(flow[:, :, 0])
    flow[nanidx] = 0.0

    # colorwheel
    ncols = RY + YG + GC + CB + BM + MR
    nchans = 3
    colorwheel = np.zeros((ncols, nchans), "uint8")
    col = 0


    eps = np.spacing(1)  # minimum positive float value to avoid division by 0
    # compute the flow
    img = _computeColor(flow / (maxflow + eps), saturate=saturate)
    # put black pixels in unknown location
    img[np.tile(unknown_idx[:, :, np.newaxis], [1, 1, 3])] = 0.0
    return img

def flowMaxNorm(flow):
    """
    flow_utils.flowMaxNorm(flow): return the maximum of the l2-norm of the given flow
    ---- PARAMETERS ----
        flow: the flow

    ---- OUTPUT ----
        a float containing the maximum of the l2-norm of the flow
    """
    return np.max(np.sqrt(np.sum(np.square(flow), 2)))

def _computeColor(flow, saturate=True):
    """
    flow_utils._computeColor(flow): compute color codes for the flow field flow

    ---- PARAMETERS ----
        flow: np.array of dimension (height x width x 2) containing the flow to display
    ---- OUTPUTS ----
        an np.array of dimension (height x width x 3) containing the color conversion of the flow
    """
    # set nan to 0
    nanidx = np.isnan(flow[:, :, 0])
    flow[nanidx] = 0.0

    # colorwheel
    ncols = RY + YG + GC + CB + BM + MR
    nchans = 3
    colorwheel = np.zeros((ncols, nchans), "uint8")
    col = 0
    # RY
    colorwheel[:RY, 0] = 255
    colorwheel[:RY, 1] = [(255 * i) // RY for i in range(RY)]
    col += RY
    # YG
    colorwheel[col : col + YG, 0] = [255 - (255 * i) // YG for i in range(YG)]
    colorwheel[col : col + YG, 1] = 255
    col += YG
    # GC
    colorwheel[col : col + GC, 1] = 255
    colorwheel[col : col + GC, 2] = [(255 * i) // GC for i in range(GC)]
    col += GC
    # CB
    colorwheel[col : col + CB, 1] = [255 - (255 * i) // CB for i in range(CB)]
    colorwheel[col : col + CB, 2] = 255
    col += CB
    # BM
    colorwheel[col : col + BM, 0] = [(255 * i) // BM for i in range(BM)]
    colorwheel[col : col + BM, 2] = 255
    col += BM
    # MR
    colorwheel[col : col + MR, 0] = 255
    colorwheel[col : col + MR, 2] = [255 - (255 * i) // MR for i in range(MR)]

    # compute utility variables
    rad = np.sqrt(np.sum(np.square(flow), 2))  # magnitude
    a = np.arctan2(-flow[:, :, 1], -flow[:, :, 0]) / np.pi  # angle
    fk = (a + 1) / 2 * (ncols - 1)  # map [-1,1] to [0,ncols-1]
    k0 = np.floor(fk).astype("int")
    k1 = k0 + 1
    k1[k1 == ncols] = 0
    f = fk - k0

    if not saturate:
        rad = np.minimum(rad, 1)

    # compute the image
    img = np.zeros((flow.shape[0], flow.shape[1], nchans), "uint8")
    for i in range(nchans):
        tmp = colorwheel[:, i].astype("float")
        col0 = tmp[k0] / 255
        col1 = tmp[k1] / 255
        col = (1 - f) * col0 + f * col1
        idx = rad <= 1
        col[idx] = 1 - rad[idx] * (1 - col[idx])  # increase saturation with radius
        col[~idx] *= 0.75  # out of range
        img[:, :, i] = (255 * col * (1 - nanidx.astype("float"))).astype("uint8")

    return img

# flow dataset getter

def get_train_dataset_flow(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_flow(dataset_str):
    dataset_str = dataset_str.replace("(", "Dataset(")
    return [eval(s) for s in dataset_str.split("+")]


        rad = np.minimum(rad, 1)

    # compute the image
    img = np.zeros((flow.shape[0], flow.shape[1], nchans), "uint8")
    for i in range(nchans):
        tmp = colorwheel[:, i].astype("float")
        col0 = tmp[k0] / 255
        col1 = tmp[k1] / 255
        col = (1 - f) * col0 + f * col1
        idx = rad <= 1
        col[idx] = 1 - rad[idx] * (1 - col[idx])  # increase saturation with radius
        col[~idx] *= 0.75  # out of range
        img[:, :, i] = (255 * col * (1 - nanidx.astype("float"))).astype("uint8")

    return img

# flow dataset getter

def get_train_dataset_flow(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_flow(dataset_str):
    dataset_str = dataset_str.replace("(", "Dataset(")
    return [eval(s) for s in dataset_str.split("+")]


        col[~idx] *= 0.75  # out of range
        img[:, :, i] = (255 * col * (1 - nanidx.astype("float"))).astype("uint8")

    return img

# flow dataset getter

def get_train_dataset_flow(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_flow(dataset_str):
    dataset_str = dataset_str.replace("(", "Dataset(")
    return [eval(s) for s in dataset_str.split("+")]


dataset_to_root = deepcopy(dataset_to_root)

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

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

class FlowDataset(data.Dataset):

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

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

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

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

        # get filenames
        img1name = self.pairname_to_img1name(pairname)
        img2name = self.pairname_to_img2name(pairname)
        flowname = (
            self.pairname_to_flowname(pairname)


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

class FlowDataset(data.Dataset):

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

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

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

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

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

        # load images and disparities
        img1 = _read_img(img1name)


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

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

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

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

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

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

        # apply augmentations
        if self.augmentor is not None:


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

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

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

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

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

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

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

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

        return img1, img2, flow, pairname

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

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

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

    def _set_root(self):
        self.root = dataset_to_root[self.name]
        assert os.path.isdir(


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

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

        return img1, img2, flow, pairname

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

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

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

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

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")


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

        return img1, img2, flow, pairname

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

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

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

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

    def _load_or_build_cache(self):
        cache_file = osp.join(cache_dir, self.name + ".pkl")
        if osp.isfile(cache_file):
            with open(cache_file, "rb") as fid:
                self.pairnames = pickle.load(fid)[self.split]


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

        return img1, img2, flow, pairname

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

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

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

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

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


            )  # transform potential tuple to str to be able to batch it

        return img1, img2, flow, pairname

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

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

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

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

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

class TartanAirDataset(FlowDataset):

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


        self.pairnames = v * self.pairnames
        return self

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

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

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

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

class TartanAirDataset(FlowDataset):

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


        tosave = {"train": trainseqs, "test": testseqs}
        return tosave

    def submission_save_pairname(self, pairname, prediction, outdir, time):
        assert prediction.ndim == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

class Kitti15Dataset(FlowDataset):

    def _prepare_data(self):
        self.name = "Kitti15"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval", "test"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname + "_11.png"
        )
        self.pairname_to_flowname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/image_2/", "/flow_occ/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/image_2/", "/")
        self.load_flow = _read_kitti_flow

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


    def _prepare_data(self):
        self.name = "Kitti15"
        self._set_root()
        assert self.split in ["train", "subtrain", "subval", "test"]
        self.pairname_to_img1name = lambda pairname: osp.join(
            self.root, pairname + "_10.png"
        )
        self.pairname_to_img2name = lambda pairname: osp.join(
            self.root, pairname + "_11.png"
        )
        self.pairname_to_flowname = (
            None
            if self.split == "test"
            else lambda pairname: osp.join(
                self.root, pairname.replace("/image_2/", "/flow_occ/") + "_10.png"
            )
        )
        self.pairname_to_str = lambda pairname: pairname.replace("/image_2/", "/")
        self.load_flow = _read_kitti_flow

    def _build_cache(self):
        trainseqs = ["training/image_2/%06d" % (i) for i in range(200)]
        subtrainseqs = trainseqs[:-10]
        subvalseqs = trainseqs[-10:]
        testseqs = ["testing/image_2/%06d" % (i) for i in range(200)]
        assert (
            len(trainseqs) == 200
            and len(subtrainseqs) == 190
            and len(subvalseqs) == 10
            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 == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, "flow", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

import cv2

def _read_numpy_flow(filename):


    def _build_cache(self):
        trainseqs = ["training/image_2/%06d" % (i) for i in range(200)]
        subtrainseqs = trainseqs[:-10]
        subvalseqs = trainseqs[-10:]
        testseqs = ["testing/image_2/%06d" % (i) for i in range(200)]
        assert (
            len(trainseqs) == 200
            and len(subtrainseqs) == 190
            and len(subvalseqs) == 10
            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 == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, "flow", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

import cv2

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

def _read_pfm_flow(filename):
    f, _ = _read_pfm(filename)
    assert np.all(f[:, :, 2] == 0.0)
    return np.ascontiguousarray(f[:, :, :2])


            len(trainseqs) == 200
            and len(subtrainseqs) == 190
            and len(subvalseqs) == 10
            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 == 3
        assert prediction.shape[2] == 2
        outfile = os.path.join(outdir, "flow", pairname.split("/")[-1] + "_10.png")
        os.makedirs(os.path.dirname(outfile), exist_ok=True)
        writeFlowKitti(outfile, prediction)

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

import cv2

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

def _read_pfm_flow(filename):
    f, _ = _read_pfm(filename)
    assert np.all(f[:, :, 2] == 0.0)
    return np.ascontiguousarray(f[:, :, :2])

TAG_FLOAT = 202021.25  # tag to check the sanity of the file
TAG_STRING = "PIEH"  # string containing the tag
MIN_WIDTH = 1
MAX_WIDTH = 99999
MIN_HEIGHT = 1


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

import os
from torch.utils.data import Dataset
from PIL import Image

from datasets.transforms import get_pair_transforms

def load_image(impath):
    return Image.open(impath)

def load_pairs_from_cache_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l.split()[0]), os.path.join(root, l.split()[1]))
        for l in lines
    ]
    return pairs

def load_pairs_from_list_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l + "_1.jpg"), os.path.join(root, l + "_2.jpg"))
        for l in lines
        if not l.startswith("#")
    ]
    return pairs

def write_cache_file(fname, pairs, root=""):
    if len(root) > 0:
        if not root.endswith("/"):
            root += "/"
        assert os.path.isdir(root)
    s = ""
    for im1, im2 in pairs:
        if len(root) > 0:
            assert im1.startswith(root), im1
            assert im2.startswith(root), im2
        s += "{:s} {:s}\n".format(im1[len(root) :], im2[len(root) :])
    with open(fname, "w") as fid:
        fid.write(s[:-1])

def parse_and_cache_all_pairs(dname, data_dir="./data/"):
    if dname == "habitat_release":
        dirname = os.path.join(data_dir, "habitat_release")
        assert os.path.isdir(dirname), (
            "cannot find folder for habitat_release pairs: " + dirname
        )
        cache_file = os.path.join(dirname, "pairs.txt")
        assert not os.path.isfile(cache_file), (
            "cache file already exists: " + cache_file
        )

        print("Parsing pairs for dataset: " + dname)
        pairs = []
        for root, dirs, files in os.walk(dirname):
            if "val" in root:
                continue
            dirs.sort()
            pairs += [
                (
                    os.path.join(root, f),
                    os.path.join(root, f[: -len("_1.jpeg")] + "_2.jpeg"),
                )
                for f in sorted(files)
                if f.endswith("_1.jpeg")
            ]
        print("Found {:,} pairs".format(len(pairs)))
        print("Writing cache to: " + cache_file)
        write_cache_file(cache_file, pairs, root=dirname)

    else:
        raise NotImplementedError("Unknown dataset: " + dname)

def dnames_to_image_pairs(dnames, data_dir="./data/"):
    """
    dnames: list of datasets with image pairs, separated by +
    """
    all_pairs = []
    for dname in dnames.split("+"):
        if dname == "habitat_release":
            dirname = os.path.join(data_dir, "habitat_release")
            assert os.path.isdir(dirname), (
                "cannot find folder for habitat_release pairs: " + dirname
            )
            cache_file = os.path.join(dirname, "pairs.txt")
            assert os.path.isfile(cache_file), (
                "cannot find cache file for habitat_release pairs, please first create the cache file, see instructions. "
                + cache_file
            )
            pairs = load_pairs_from_cache_file(cache_file, root=dirname)
        elif dname in ["ARKitScenes", "MegaDepth", "3DStreetView", "IndoorVL"]:
            dirname = os.path.join(data_dir, dname + "_crops")
            assert os.path.isdir(
                dirname
            ), "cannot find folder for {:s} pairs: {:s}".format(dname, dirname)
            list_file = os.path.join(dirname, "listing.txt")
            assert os.path.isfile(
                list_file
            ), "cannot find list file for {:s} pairs, see instructions. {:s}".format(
                dname, list_file
            )
            pairs = load_pairs_from_list_file(list_file, root=dirname)
        print("  {:s}: {:,} pairs".format(dname, len(pairs)))
        all_pairs += pairs
    if "+" in dnames:
        print(" Total: {:,} pairs".format(len(all_pairs)))
    return all_pairs

class PairsDataset(Dataset):

    def __init__(
        self, dnames, trfs="", totensor=True, normalize=True, data_dir="./data/"
    ):
        super().__init__()
        self.image_pairs = dnames_to_image_pairs(dnames, data_dir=data_dir)
        self.transforms = get_pair_transforms(
            transform_str=trfs, totensor=totensor, normalize=normalize
        )

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, index):
        im1path, im2path = self.image_pairs[index]
        im1 = load_image(im1path)
        im2 = load_image(im2path)
        if self.transforms is not None:
            im1, im2 = self.transforms(im1, im2)
        return im1, im2

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        prog="Computing and caching list of pairs for a given dataset"
    )
    parser.add_argument(
        "--data_dir", default="./data/", type=str, help="path where data are stored"
    )
    parser.add_argument(
        "--dataset", default="habitat_release", type=str, help="name of the dataset"
    )
    args = parser.parse_args()
    parse_and_cache_all_pairs(dname=args.dataset, data_dir=args.data_dir)


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

import os
from torch.utils.data import Dataset
from PIL import Image

from datasets.transforms import get_pair_transforms

def load_image(impath):
    return Image.open(impath)

def load_pairs_from_cache_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l.split()[0]), os.path.join(root, l.split()[1]))
        for l in lines
    ]
    return pairs

def load_pairs_from_list_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:


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

import os
from torch.utils.data import Dataset
from PIL import Image

from datasets.transforms import get_pair_transforms

def load_image(impath):
    return Image.open(impath)

def load_pairs_from_cache_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l.split()[0]), os.path.join(root, l.split()[1]))
        for l in lines
    ]
    return pairs

def load_pairs_from_list_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l + "_1.jpg"), os.path.join(root, l + "_2.jpg"))
        for l in lines
        if not l.startswith("#")
    ]
    return pairs

def write_cache_file(fname, pairs, root=""):
    if len(root) > 0:
        if not root.endswith("/"):
            root += "/"


from datasets.transforms import get_pair_transforms

def load_image(impath):
    return Image.open(impath)

def load_pairs_from_cache_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l.split()[0]), os.path.join(root, l.split()[1]))
        for l in lines
    ]
    return pairs

def load_pairs_from_list_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l + "_1.jpg"), os.path.join(root, l + "_2.jpg"))
        for l in lines
        if not l.startswith("#")
    ]
    return pairs

def write_cache_file(fname, pairs, root=""):
    if len(root) > 0:
        if not root.endswith("/"):
            root += "/"
        assert os.path.isdir(root)
    s = ""
    for im1, im2 in pairs:
        if len(root) > 0:
            assert im1.startswith(root), im1
            assert im2.startswith(root), im2
        s += "{:s} {:s}\n".format(im1[len(root) :], im2[len(root) :])
    with open(fname, "w") as fid:
        fid.write(s[:-1])

def parse_and_cache_all_pairs(dname, data_dir="./data/"):
    if dname == "habitat_release":
        dirname = os.path.join(data_dir, "habitat_release")


        (os.path.join(root, l.split()[0]), os.path.join(root, l.split()[1]))
        for l in lines
    ]
    return pairs

def load_pairs_from_list_file(fname, root=""):
    assert os.path.isfile(
        fname
    ), "cannot parse pairs from {:s}, file does not exist".format(fname)
    with open(fname, "r") as fid:
        lines = fid.read().strip().splitlines()
    pairs = [
        (os.path.join(root, l + "_1.jpg"), os.path.join(root, l + "_2.jpg"))
        for l in lines
        if not l.startswith("#")
    ]
    return pairs

def write_cache_file(fname, pairs, root=""):
    if len(root) > 0:
        if not root.endswith("/"):
            root += "/"
        assert os.path.isdir(root)
    s = ""
    for im1, im2 in pairs:
        if len(root) > 0:
            assert im1.startswith(root), im1
            assert im2.startswith(root), im2
        s += "{:s} {:s}\n".format(im1[len(root) :], im2[len(root) :])
    with open(fname, "w") as fid:
        fid.write(s[:-1])

def parse_and_cache_all_pairs(dname, data_dir="./data/"):
    if dname == "habitat_release":
        dirname = os.path.join(data_dir, "habitat_release")
        assert os.path.isdir(dirname), (
            "cannot find folder for habitat_release pairs: " + dirname
        )
        cache_file = os.path.join(dirname, "pairs.txt")
        assert not os.path.isfile(cache_file), (
            "cache file already exists: " + cache_file
        )

        print("Parsing pairs for dataset: " + dname)
        pairs = []
        for root, dirs, files in os.walk(dirname):
            if "val" in root:
                continue
            dirs.sort()
            pairs += [


        if not l.startswith("#")
    ]
    return pairs

def write_cache_file(fname, pairs, root=""):
    if len(root) > 0:
        if not root.endswith("/"):
            root += "/"
        assert os.path.isdir(root)
    s = ""
    for im1, im2 in pairs:
        if len(root) > 0:
            assert im1.startswith(root), im1
            assert im2.startswith(root), im2
        s += "{:s} {:s}\n".format(im1[len(root) :], im2[len(root) :])
    with open(fname, "w") as fid:
        fid.write(s[:-1])

def parse_and_cache_all_pairs(dname, data_dir="./data/"):
    if dname == "habitat_release":
        dirname = os.path.join(data_dir, "habitat_release")
        assert os.path.isdir(dirname), (
            "cannot find folder for habitat_release pairs: " + dirname
        )
        cache_file = os.path.join(dirname, "pairs.txt")
        assert not os.path.isfile(cache_file), (
            "cache file already exists: " + cache_file
        )

        print("Parsing pairs for dataset: " + dname)
        pairs = []
        for root, dirs, files in os.walk(dirname):
            if "val" in root:
                continue
            dirs.sort()
            pairs += [
                (
                    os.path.join(root, f),
                    os.path.join(root, f[: -len("_1.jpeg")] + "_2.jpeg"),
                )
                for f in sorted(files)
                if f.endswith("_1.jpeg")
            ]
        print("Found {:,} pairs".format(len(pairs)))
        print("Writing cache to: " + cache_file)
        write_cache_file(cache_file, pairs, root=dirname)

    else:
        raise NotImplementedError("Unknown dataset: " + dname)

def dnames_to_image_pairs(dnames, data_dir="./data/"):
    """
    dnames: list of datasets with image pairs, separated by +
    """
    all_pairs = []
    for dname in dnames.split("+"):
        if dname == "habitat_release":
            dirname = os.path.join(data_dir, "habitat_release")
            assert os.path.isdir(dirname), (
                "cannot find folder for habitat_release pairs: " + dirname
            )
            cache_file = os.path.join(dirname, "pairs.txt")
            assert os.path.isfile(cache_file), (
                "cannot find cache file for habitat_release pairs, please first create the cache file, see instructions. "
                + cache_file
            )
            pairs = load_pairs_from_cache_file(cache_file, root=dirname)
        elif dname in ["ARKitScenes", "MegaDepth", "3DStreetView", "IndoorVL"]:


        for root, dirs, files in os.walk(dirname):
            if "val" in root:
                continue
            dirs.sort()
            pairs += [
                (
                    os.path.join(root, f),
                    os.path.join(root, f[: -len("_1.jpeg")] + "_2.jpeg"),
                )
                for f in sorted(files)
                if f.endswith("_1.jpeg")
            ]
        print("Found {:,} pairs".format(len(pairs)))
        print("Writing cache to: " + cache_file)
        write_cache_file(cache_file, pairs, root=dirname)

    else:
        raise NotImplementedError("Unknown dataset: " + dname)

def dnames_to_image_pairs(dnames, data_dir="./data/"):
    """
    dnames: list of datasets with image pairs, separated by +
    """
    all_pairs = []
    for dname in dnames.split("+"):
        if dname == "habitat_release":
            dirname = os.path.join(data_dir, "habitat_release")
            assert os.path.isdir(dirname), (
                "cannot find folder for habitat_release pairs: " + dirname
            )
            cache_file = os.path.join(dirname, "pairs.txt")
            assert os.path.isfile(cache_file), (
                "cannot find cache file for habitat_release pairs, please first create the cache file, see instructions. "
                + cache_file
            )
            pairs = load_pairs_from_cache_file(cache_file, root=dirname)
        elif dname in ["ARKitScenes", "MegaDepth", "3DStreetView", "IndoorVL"]:
            dirname = os.path.join(data_dir, dname + "_crops")
            assert os.path.isdir(
                dirname
            ), "cannot find folder for {:s} pairs: {:s}".format(dname, dirname)
            list_file = os.path.join(dirname, "listing.txt")
            assert os.path.isfile(
                list_file
            ), "cannot find list file for {:s} pairs, see instructions. {:s}".format(
                dname, list_file
            )
            pairs = load_pairs_from_list_file(list_file, root=dirname)
        print("  {:s}: {:,} pairs".format(dname, len(pairs)))
        all_pairs += pairs
    if "+" in dnames:
        print(" Total: {:,} pairs".format(len(all_pairs)))
    return all_pairs

class PairsDataset(Dataset):

    def __init__(
        self, dnames, trfs="", totensor=True, normalize=True, data_dir="./data/"
    ):
        super().__init__()
        self.image_pairs = dnames_to_image_pairs(dnames, data_dir=data_dir)
        self.transforms = get_pair_transforms(
            transform_str=trfs, totensor=totensor, normalize=normalize
        )

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, index):
        im1path, im2path = self.image_pairs[index]
        im1 = load_image(im1path)
        im2 = load_image(im2path)


PairsDataset

            pairs = load_pairs_from_cache_file(cache_file, root=dirname)
        elif dname in ["ARKitScenes", "MegaDepth", "3DStreetView", "IndoorVL"]:
            dirname = os.path.join(data_dir, dname + "_crops")
            assert os.path.isdir(
                dirname
            ), "cannot find folder for {:s} pairs: {:s}".format(dname, dirname)
            list_file = os.path.join(dirname, "listing.txt")
            assert os.path.isfile(
                list_file
            ), "cannot find list file for {:s} pairs, see instructions. {:s}".format(
                dname, list_file
            )
            pairs = load_pairs_from_list_file(list_file, root=dirname)
        print("  {:s}: {:,} pairs".format(dname, len(pairs)))
        all_pairs += pairs
    if "+" in dnames:
        print(" Total: {:,} pairs".format(len(all_pairs)))
    return all_pairs

class PairsDataset(Dataset):

    def __init__(
        self, dnames, trfs="", totensor=True, normalize=True, data_dir="./data/"
    ):
        super().__init__()
        self.image_pairs = dnames_to_image_pairs(dnames, data_dir=data_dir)
        self.transforms = get_pair_transforms(
            transform_str=trfs, totensor=totensor, normalize=normalize
        )

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, index):
        im1path, im2path = self.image_pairs[index]
        im1 = load_image(im1path)
        im2 = load_image(im2path)
        if self.transforms is not None:
            im1, im2 = self.transforms(im1, im2)
        return im1, im2

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        prog="Computing and caching list of pairs for a given dataset"
    )
    parser.add_argument(
        "--data_dir", default="./data/", type=str, help="path where data are stored"
    )
    parser.add_argument(
        "--dataset", default="habitat_release", type=str, help="name of the dataset"
    )
    args = parser.parse_args()
    parse_and_cache_all_pairs(dname=args.dataset, data_dir=args.data_dir)


            dirname = os.path.join(data_dir, dname + "_crops")
            assert os.path.isdir(
                dirname
            ), "cannot find folder for {:s} pairs: {:s}".format(dname, dirname)
            list_file = os.path.join(dirname, "listing.txt")
            assert os.path.isfile(
                list_file
            ), "cannot find list file for {:s} pairs, see instructions. {:s}".format(
                dname, list_file
            )
            pairs = load_pairs_from_list_file(list_file, root=dirname)
        print("  {:s}: {:,} pairs".format(dname, len(pairs)))
        all_pairs += pairs
    if "+" in dnames:
        print(" Total: {:,} pairs".format(len(all_pairs)))
    return all_pairs

class PairsDataset(Dataset):

    def __init__(
        self, dnames, trfs="", totensor=True, normalize=True, data_dir="./data/"
    ):
        super().__init__()
        self.image_pairs = dnames_to_image_pairs(dnames, data_dir=data_dir)
        self.transforms = get_pair_transforms(
            transform_str=trfs, totensor=totensor, normalize=normalize
        )

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, index):
        im1path, im2path = self.image_pairs[index]
        im1 = load_image(im1path)
        im2 = load_image(im2path)
        if self.transforms is not None:
            im1, im2 = self.transforms(im1, im2)
        return im1, im2

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        prog="Computing and caching list of pairs for a given dataset"
    )
    parser.add_argument(


            )
            pairs = load_pairs_from_list_file(list_file, root=dirname)
        print("  {:s}: {:,} pairs".format(dname, len(pairs)))
        all_pairs += pairs
    if "+" in dnames:
        print(" Total: {:,} pairs".format(len(all_pairs)))
    return all_pairs

class PairsDataset(Dataset):

    def __init__(
        self, dnames, trfs="", totensor=True, normalize=True, data_dir="./data/"
    ):
        super().__init__()
        self.image_pairs = dnames_to_image_pairs(dnames, data_dir=data_dir)
        self.transforms = get_pair_transforms(
            transform_str=trfs, totensor=totensor, normalize=normalize
        )

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, index):
        im1path, im2path = self.image_pairs[index]
        im1 = load_image(im1path)
        im2 = load_image(im2path)
        if self.transforms is not None:
            im1, im2 = self.transforms(im1, im2)
        return im1, im2

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        prog="Computing and caching list of pairs for a given dataset"
    )
    parser.add_argument(
        "--data_dir", default="./data/", type=str, help="path where data are stored"
    )
    parser.add_argument(


        all_pairs += pairs
    if "+" in dnames:
        print(" Total: {:,} pairs".format(len(all_pairs)))
    return all_pairs

class PairsDataset(Dataset):

    def __init__(
        self, dnames, trfs="", totensor=True, normalize=True, data_dir="./data/"
    ):
        super().__init__()
        self.image_pairs = dnames_to_image_pairs(dnames, data_dir=data_dir)
        self.transforms = get_pair_transforms(
            transform_str=trfs, totensor=totensor, normalize=normalize
        )

    def __len__(self):
        return len(self.image_pairs)

    def __getitem__(self, index):
        im1path, im2path = self.image_pairs[index]
        im1 = load_image(im1path)
        im2 = load_image(im2path)
        if self.transforms is not None:
            im1, im2 = self.transforms(im1, im2)
        return im1, im2

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(
        prog="Computing and caching list of pairs for a given dataset"
    )
    parser.add_argument(
        "--data_dir", default="./data/", type=str, help="path where data are stored"
    )
    parser.add_argument(
        "--dataset", default="habitat_release", type=str, help="name of the dataset"
    )
    args = parser.parse_args()
    parse_and_cache_all_pairs(dname=args.dataset, data_dir=args.data_dir)