third_party/dust3r/croco/stereoflow/datasets_flow.py

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

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

import json
import os
import os.path as osp
import pickle
import struct
from copy import deepcopy

import h5py
import numpy as np
import torch
from PIL import Image
from torch.utils import data

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

dataset_to_root = deepcopy(dataset_to_root)

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


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


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

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

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

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

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

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

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

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

        return img1, img2, flow, pairname

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

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

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

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

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


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

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


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

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


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

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


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

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

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

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


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

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

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

    def finalize_submission(self, outdir):
        assert self.split == "test"
        exe = "{self.root}/flow_subsampling"
        if os.path.isfile(exe):
            cmd = f'cd "{outdir}/test"; {exe} .'
            print(cmd)
            os.system(cmd)
            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"
            )
        )
        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
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
            )
        )

    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):
    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
    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
    ---- 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("+")]