vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/duster/croco/stereoflow/datasets_flow.py

@@ -0,0 +1,630 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+# --------------------------------------------------------
+# Dataset structure for flow
+# --------------------------------------------------------
+
+import os
+import os.path as osp
+import pickle
+import numpy as np
+import struct
+from PIL import Image 
+import json
+import h5py
+import torch
+from torch.utils import data
+
+from .augmentor import FlowAugmentor
+from .datasets_stereo import _read_img, img_to_tensor, dataset_to_root, _read_pfm
+from copy import deepcopy
+dataset_to_root = deepcopy(dataset_to_root)
+
+dataset_to_root.update(**{
+    'TartanAir': './data/stereoflow/TartanAir',
+    'FlyingChairs': './data/stereoflow/FlyingChairs/',
+    'FlyingThings': osp.join(dataset_to_root['SceneFlow'],'FlyingThings')+'/',
+    'MPISintel': './data/stereoflow//MPI-Sintel/'+'/',
+})
+cache_dir = "./data/stereoflow/datasets_flow_cache/"
+
+
+def flow_to_tensor(disp):
+    return torch.from_numpy(disp).float().permute(2, 0, 1)
+
+class FlowDataset(data.Dataset):
+    
+    def __init__(self, split, augmentor=False, crop_size=None, totensor=True):
+        self.split = split
+        if not augmentor: assert crop_size is None 
+        if crop_size is not None: assert augmentor
+        self.crop_size = crop_size
+        self.augmentor_str = augmentor
+        self.augmentor = FlowAugmentor(crop_size) if augmentor else None
+        self.totensor = totensor
+        self.rmul = 1 # keep track of rmul
+        self.has_constant_resolution = True # whether the dataset has constant resolution or not (=> don't use batch_size>1 at test time)
+        self._prepare_data()
+        self._load_or_build_cache()
+        
+    def prepare_data(self):
+        """
+        to be defined for each dataset 
+        """
+        raise NotImplementedError 
+        
+    def __len__(self):
+        return len(self.pairnames) # each pairname is typically of the form (str, int1, int2) 
+        
+    def __getitem__(self, index):
+        pairname = self.pairnames[index]
+        
+        # get filenames 
+        img1name = self.pairname_to_img1name(pairname)
+        img2name = self.pairname_to_img2name(pairname)
+        flowname = self.pairname_to_flowname(pairname) if self.pairname_to_flowname is not None else None
+        
+        # load images and disparities
+        img1 = _read_img(img1name)
+        img2 = _read_img(img2name)
+        flow = self.load_flow(flowname) if flowname is not None else None
+
+        # apply augmentations
+        if self.augmentor is not None:
+            img1, img2, flow = self.augmentor(img1, img2, flow, self.name)
+        
+        if self.totensor:
+            img1 = img_to_tensor(img1)
+            img2 = img_to_tensor(img2)
+            if flow is not None: 
+                flow = flow_to_tensor(flow)
+            else:
+                flow = torch.tensor([]) # to allow dataloader batching with default collate_gn
+            pairname = str(pairname) # transform potential tuple to str to be able to batch it
+
+        return img1, img2, flow, pairname
+        
+    def __rmul__(self, v):
+        self.rmul *= v
+        self.pairnames = v * self.pairnames
+        return self
+        
+    def __str__(self):
+        return f'{self.__class__.__name__}_{self.split}'
+        
+    def __repr__(self):
+        s = f'{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})'
+        if self.rmul==1:
+            s+=f'\n\tnum pairs: {len(self.pairnames)}'
+        else:
+            s+=f'\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})'
+        return s
+
+    def _set_root(self):
+        self.root = dataset_to_root[self.name]
+        assert os.path.isdir(self.root), f"could not find root directory for dataset {self.name}: {self.root}"       
+
+    def _load_or_build_cache(self):
+        cache_file = osp.join(cache_dir, self.name+'.pkl')
+        if osp.isfile(cache_file):
+            with open(cache_file, 'rb') as fid:
+                self.pairnames = pickle.load(fid)[self.split]
+        else:
+            tosave = self._build_cache()
+            os.makedirs(cache_dir, exist_ok=True)
+            with open(cache_file, 'wb') as fid:
+                pickle.dump(tosave, fid)
+            self.pairnames = tosave[self.split]
+
+class TartanAirDataset(FlowDataset):
+
+    def _prepare_data(self):
+        self.name = "TartanAir"
+        self._set_root()
+        assert self.split in ['train']
+        self.pairname_to_img1name = lambda pairname: osp.join(self.root, pairname[0], 'image_left/{:06d}_left.png'.format(pairname[1]))
+        self.pairname_to_img2name = lambda pairname: osp.join(self.root, pairname[0], 'image_left/{:06d}_left.png'.format(pairname[2]))
+        self.pairname_to_flowname = lambda pairname: osp.join(self.root, pairname[0], 'flow/{:06d}_{:06d}_flow.npy'.format(pairname[1],pairname[2]))
+        self.pairname_to_str = lambda pairname: os.path.join(pairname[0][pairname[0].find('/')+1:], '{:06d}_{:06d}'.format(pairname[1], pairname[2]))
+        self.load_flow = _read_numpy_flow
+        
+    def _build_cache(self):
+        seqs = sorted(os.listdir(self.root))
+        pairs = [(osp.join(s,s,difficulty,Pxxx),int(a[:6]),int(a[:6])+1) for s in seqs for difficulty in ['Easy','Hard'] for Pxxx in sorted(os.listdir(osp.join(self.root,s,s,difficulty))) for a in sorted(os.listdir(osp.join(self.root,s,s,difficulty,Pxxx,'image_left/')))[:-1]]
+        assert len(pairs)==306268, "incorrect parsing of pairs in TartanAir"
+        tosave = {'train': pairs}
+        return tosave
+        
+class FlyingChairsDataset(FlowDataset):
+
+    def _prepare_data(self):
+        self.name = "FlyingChairs"
+        self._set_root()
+        assert self.split in ['train','val']
+        self.pairname_to_img1name = lambda pairname: osp.join(self.root, 'data', pairname+'_img1.ppm')
+        self.pairname_to_img2name = lambda pairname: osp.join(self.root, 'data', pairname+'_img2.ppm')
+        self.pairname_to_flowname = lambda pairname: osp.join(self.root, 'data', pairname+'_flow.flo')
+        self.pairname_to_str = lambda pairname: pairname
+        self.load_flow = _read_flo_file
+        
+    def _build_cache(self):
+        split_file = osp.join(self.root, 'chairs_split.txt')
+        split_list = np.loadtxt(split_file, dtype=np.int32)
+        trainpairs = ['{:05d}'.format(i) for i in np.where(split_list==1)[0]+1]
+        valpairs = ['{:05d}'.format(i) for i in np.where(split_list==2)[0]+1]
+        assert len(trainpairs)==22232 and len(valpairs)==640, "incorrect parsing of pairs in MPI-Sintel"
+        tosave = {'train': trainpairs, 'val': valpairs}
+        return tosave
+        
+class FlyingThingsDataset(FlowDataset):
+    
+    def _prepare_data(self):
+        self.name = "FlyingThings"
+        self._set_root()
+        assert self.split in [f'{set_}_{pass_}pass{camstr}' for set_ in ['train','test','test1024'] for camstr in ['','_rightcam'] for pass_ in ['clean','final','all']]
+        self.pairname_to_img1name = lambda pairname: osp.join(self.root, f'frames_{pairname[3]}pass', pairname[0].replace('into_future','').replace('into_past',''), '{:04d}.png'.format(pairname[1]))
+        self.pairname_to_img2name = lambda pairname: osp.join(self.root, f'frames_{pairname[3]}pass', pairname[0].replace('into_future','').replace('into_past',''), '{:04d}.png'.format(pairname[2]))
+        self.pairname_to_flowname = lambda pairname: osp.join(self.root, 'optical_flow', pairname[0], 'OpticalFlowInto{f:s}_{i:04d}_{c:s}.pfm'.format(f='Future' if 'future' in pairname[0] else 'Past', i=pairname[1], c='L' if 'left' in pairname[0] else 'R' ))
+        self.pairname_to_str = lambda pairname: os.path.join(pairname[3]+'pass', pairname[0], 'Into{f:s}_{i:04d}_{c:s}'.format(f='Future' if 'future' in pairname[0] else 'Past',  i=pairname[1], c='L' if 'left' in pairname[0] else 'R' ))
+        self.load_flow = _read_pfm_flow
+        
+    def _build_cache(self):
+        tosave = {}
+        # train and test splits for the different passes 
+        for set_ in ['train', 'test']:
+            sroot = osp.join(self.root, 'optical_flow', set_.upper())
+            fname_to_i = lambda f: int(f[len('OpticalFlowIntoFuture_'):-len('_L.pfm')])
+            pp = [(osp.join(set_.upper(), d, s, 'into_future/left'),fname_to_i(fname)) for d in sorted(os.listdir(sroot)) for s in sorted(os.listdir(osp.join(sroot,d))) for fname in sorted(os.listdir(osp.join(sroot,d, s, 'into_future/left')))[:-1]]
+            pairs  = [(a,i,i+1) for a,i in pp]
+            pairs += [(a.replace('into_future','into_past'),i+1,i) for a,i in pp]
+            assert len(pairs)=={'train': 40302, 'test': 7866}[set_], "incorrect parsing of pairs Flying Things"
+            for cam in ['left','right']:
+                camstr = '' if cam=='left' else f'_{cam}cam'
+                for pass_ in ['final', 'clean']:
+                    tosave[f'{set_}_{pass_}pass{camstr}'] = [(a.replace('left',cam),i,j,pass_) for a,i,j in pairs]
+                tosave[f'{set_}_allpass{camstr}'] = tosave[f'{set_}_cleanpass{camstr}'] + tosave[f'{set_}_finalpass{camstr}']
+        # test1024: this is the same split as unimatch 'validation' split
+        # see https://github.com/autonomousvision/unimatch/blob/master/dataloader/flow/datasets.py#L229
+        test1024_nsamples = 1024
+        alltest_nsamples = len(tosave['test_cleanpass'])  # 7866
+        stride = alltest_nsamples // test1024_nsamples
+        remove = alltest_nsamples % test1024_nsamples
+        for cam in ['left','right']:
+            camstr = '' if cam=='left' else f'_{cam}cam'
+            for pass_ in ['final','clean']:
+                tosave[f'test1024_{pass_}pass{camstr}'] = sorted(tosave[f'test_{pass_}pass{camstr}'])[:-remove][::stride] # warning, it was not sorted before
+            assert len(tosave['test1024_cleanpass'])==1024, "incorrect parsing of pairs in Flying Things"
+            tosave[f'test1024_allpass{camstr}'] = tosave[f'test1024_cleanpass{camstr}'] + tosave[f'test1024_finalpass{camstr}']
+        return tosave
+        
+      
+class MPISintelDataset(FlowDataset):
+    
+    def _prepare_data(self):
+        self.name = "MPISintel"
+        self._set_root()
+        assert self.split in [s+'_'+p for s in ['train','test','subval','subtrain'] for p in ['cleanpass','finalpass','allpass']]
+        self.pairname_to_img1name = lambda pairname: osp.join(self.root, pairname[0], 'frame_{:04d}.png'.format(pairname[1]))
+        self.pairname_to_img2name = lambda pairname: osp.join(self.root, pairname[0], 'frame_{:04d}.png'.format(pairname[1]+1))
+        self.pairname_to_flowname = lambda pairname: None if pairname[0].startswith('test/') else osp.join(self.root, pairname[0].replace('/clean/','/flow/').replace('/final/','/flow/'), 'frame_{:04d}.flo'.format(pairname[1]))
+        self.pairname_to_str = lambda pairname: osp.join(pairname[0], 'frame_{:04d}'.format(pairname[1]))
+        self.load_flow = _read_flo_file
+        
+    def _build_cache(self):
+        trainseqs = sorted(os.listdir(self.root+'training/clean'))
+        trainpairs = [ (osp.join('training/clean', s),i) for s in trainseqs for i in range(1, len(os.listdir(self.root+'training/clean/'+s)))]
+        subvalseqs = ['temple_2','temple_3']
+        subtrainseqs = [s for s in trainseqs if s not in subvalseqs]
+        subvalpairs = [ (p,i) for p,i in trainpairs if any(s in p for s in subvalseqs)]
+        subtrainpairs = [ (p,i) for p,i in trainpairs if any(s in p for s in subtrainseqs)]
+        testseqs = sorted(os.listdir(self.root+'test/clean'))
+        testpairs = [ (osp.join('test/clean', s),i) for s in testseqs for i in range(1, len(os.listdir(self.root+'test/clean/'+s)))]
+        assert len(trainpairs)==1041 and len(testpairs)==552 and len(subvalpairs)==98 and len(subtrainpairs)==943, "incorrect parsing of pairs in MPI-Sintel"
+        tosave = {}
+        tosave['train_cleanpass'] = trainpairs
+        tosave['test_cleanpass'] = testpairs
+        tosave['subval_cleanpass'] = subvalpairs
+        tosave['subtrain_cleanpass'] = subtrainpairs         
+        for t in ['train','test','subval','subtrain']: 
+            tosave[t+'_finalpass'] = [(p.replace('/clean/','/final/'),i) for p,i in tosave[t+'_cleanpass']]
+            tosave[t+'_allpass'] = tosave[t+'_cleanpass'] + tosave[t+'_finalpass'] 
+        return tosave
+        
+    def submission_save_pairname(self, pairname, prediction, outdir, _time):
+        assert prediction.shape[2]==2
+        outfile = os.path.join(outdir, 'submission', self.pairname_to_str(pairname)+'.flo')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        writeFlowFile(prediction, outfile)
+        
+    def finalize_submission(self, outdir):
+        assert self.split == 'test_allpass'
+        bundle_exe = "/nfs/data/ffs-3d/datasets/StereoFlow/MPI-Sintel/bundler/linux-x64/bundler" # eg <bundle_exe> <path_to_results_for_clean> <path_to_results_for_final> <output/bundled.lzma>
+        if os.path.isfile(bundle_exe):
+            cmd = f'{bundle_exe} "{outdir}/submission/test/clean/" "{outdir}/submission/test/final" "{outdir}/submission/bundled.lzma"'
+            print(cmd)
+            os.system(cmd)
+            print(f'Done. Submission file at: "{outdir}/submission/bundled.lzma"')
+        else:
+            print('Could not find bundler executable for submission.')
+            print('Please download it and run:')
+            print(f'<bundle_exe> "{outdir}/submission/test/clean/" "{outdir}/submission/test/final" "{outdir}/submission/bundled.lzma"')
+        
+class SpringDataset(FlowDataset):
+
+    def _prepare_data(self):
+        self.name = "Spring"
+        self._set_root()
+        assert self.split in ['train','test','subtrain','subval']
+        self.pairname_to_img1name = lambda pairname: osp.join(self.root, pairname[0], pairname[1], 'frame_'+pairname[3], 'frame_{:s}_{:04d}.png'.format(pairname[3], pairname[4]))
+        self.pairname_to_img2name = lambda pairname: osp.join(self.root, pairname[0], pairname[1], 'frame_'+pairname[3], 'frame_{:s}_{:04d}.png'.format(pairname[3], pairname[4]+(1 if pairname[2]=='FW' else -1)))
+        self.pairname_to_flowname = lambda pairname: None if pairname[0]=='test' else osp.join(self.root, pairname[0], pairname[1], f'flow_{pairname[2]}_{pairname[3]}', f'flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}.flo5')
+        self.pairname_to_str = lambda pairname: osp.join(pairname[0], pairname[1], f'flow_{pairname[2]}_{pairname[3]}', f'flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}')
+        self.load_flow = _read_hdf5_flow
+
+    def _build_cache(self):
+        # train 
+        trainseqs = sorted(os.listdir( osp.join(self.root,'train')))
+        trainpairs = []
+        for leftright in ['left','right']:
+            for fwbw in ['FW','BW']:
+                trainpairs += [('train',s,fwbw,leftright,int(f[len(f'flow_{fwbw}_{leftright}_'):-len('.flo5')])) for s in trainseqs for f in sorted(os.listdir(osp.join(self.root,'train',s,f'flow_{fwbw}_{leftright}')))]
+        # test 
+        testseqs = sorted(os.listdir( osp.join(self.root,'test')))
+        testpairs = []
+        for leftright in ['left','right']:
+            testpairs += [('test',s,'FW',leftright,int(f[len(f'frame_{leftright}_'):-len('.png')])) for s in testseqs for f in sorted(os.listdir(osp.join(self.root,'test',s,f'frame_{leftright}')))[:-1]]
+            testpairs += [('test',s,'BW',leftright,int(f[len(f'frame_{leftright}_'):-len('.png')])+1) for s in testseqs for f in sorted(os.listdir(osp.join(self.root,'test',s,f'frame_{leftright}')))[:-1]]
+        # subtrain / subval
+        subtrainpairs = [p for p in trainpairs if p[1]!='0041']
+        subvalpairs = [p for p in trainpairs if p[1]=='0041']
+        assert len(trainpairs)==19852 and len(testpairs)==3960 and len(subtrainpairs)==19472 and len(subvalpairs)==380, "incorrect parsing of pairs in Spring"
+        tosave = {'train': trainpairs, 'test': testpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs}
+        return tosave
+        
+    def submission_save_pairname(self, pairname, prediction, outdir, time):
+        assert prediction.ndim==3
+        assert prediction.shape[2]==2
+        assert prediction.dtype==np.float32
+        outfile = osp.join(outdir, pairname[0], pairname[1], f'flow_{pairname[2]}_{pairname[3]}', f'flow_{pairname[2]}_{pairname[3]}_{pairname[4]:04d}.flo5')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        writeFlo5File(prediction, outfile)
+        
+    def finalize_submission(self, outdir):
+        assert self.split=='test'
+        exe = "{self.root}/flow_subsampling"
+        if os.path.isfile(exe):
+            cmd = f'cd "{outdir}/test"; {exe} .'
+            print(cmd)
+            os.system(cmd)
+            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('+')]