vismatch 1.1.1__py3-none-any.whl

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

@@ -0,0 +1,674 @@
+# Copyright (C) 2022-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+
+# --------------------------------------------------------
+# Dataset structure for stereo
+# --------------------------------------------------------
+
+import sys, os
+import os.path as osp
+import pickle
+import numpy as np
+from PIL import Image
+import json
+import h5py
+from glob import glob
+import cv2
+
+import torch
+from torch.utils import data
+
+from .augmentor import StereoAugmentor
+
+
+
+dataset_to_root = {
+    'CREStereo': './data/stereoflow//crenet_stereo_trainset/stereo_trainset/crestereo/',
+    'SceneFlow': './data/stereoflow//SceneFlow/',
+    'ETH3DLowRes': './data/stereoflow/eth3d_lowres/',
+    'Booster': './data/stereoflow/booster_gt/',
+    'Middlebury2021': './data/stereoflow/middlebury/2021/data/',
+    'Middlebury2014': './data/stereoflow/middlebury/2014/',
+    'Middlebury2006': './data/stereoflow/middlebury/2006/',
+    'Middlebury2005': './data/stereoflow/middlebury/2005/train/',
+    'MiddleburyEval3':  './data/stereoflow/middlebury/MiddEval3/',
+    'Spring': './data/stereoflow/spring/', 
+    'Kitti15': './data/stereoflow/kitti-stereo-2015/',
+    'Kitti12': './data/stereoflow/kitti-stereo-2012/',
+}
+cache_dir = "./data/stereoflow/datasets_stereo_cache/"
+
+
+in1k_mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)
+in1k_std =  torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)
+def img_to_tensor(img):
+    img = torch.from_numpy(img).permute(2, 0, 1).float() / 255.
+    img = (img-in1k_mean)/in1k_std
+    return img
+def disp_to_tensor(disp):
+    return torch.from_numpy(disp)[None,:,:]
+
+class StereoDataset(data.Dataset):
+    
+    def __init__(self, split, augmentor=False, crop_size=None, totensor=True):
+        self.split = split
+        if not augmentor: assert crop_size is None 
+        if crop_size: assert augmentor
+        self.crop_size = crop_size
+        self.augmentor_str = augmentor
+        self.augmentor = StereoAugmentor(crop_size) if augmentor else None
+        self.totensor = totensor
+        self.rmul = 1 # keep track of rmul
+        self.has_constant_resolution = True # whether the dataset has constant resolution or not (=> don't use batch_size>1 at test time)
+        self._prepare_data()
+        self._load_or_build_cache()
+        
+    def prepare_data(self):
+        """
+        to be defined for each dataset 
+        """
+        raise NotImplementedError 
+        
+    def __len__(self):
+        return len(self.pairnames)
+        
+    def __getitem__(self, index):
+        pairname = self.pairnames[index]
+        
+        # get filenames 
+        Limgname = self.pairname_to_Limgname(pairname)
+        Rimgname = self.pairname_to_Rimgname(pairname)
+        Ldispname = self.pairname_to_Ldispname(pairname) if self.pairname_to_Ldispname is not None else None
+        
+        # load images and disparities
+        Limg = _read_img(Limgname)
+        Rimg = _read_img(Rimgname)
+        disp = self.load_disparity(Ldispname) if Ldispname is not None else None
+        
+        # sanity check
+        if disp is not None: assert np.all(disp>0) or self.name=="Spring", (self.name, pairname, Ldispname)
+        
+        # apply augmentations
+        if self.augmentor is not None:
+            Limg, Rimg, disp = self.augmentor(Limg, Rimg, disp, self.name)
+        
+        if self.totensor:
+            Limg = img_to_tensor(Limg)
+            Rimg = img_to_tensor(Rimg)
+            if disp is None:
+                disp = torch.tensor([]) # to allow dataloader batching with default collate_gn
+            else:
+                disp = disp_to_tensor(disp)
+        
+        return Limg, Rimg, disp, str(pairname)
+        
+    def __rmul__(self, v):
+        self.rmul *= v
+        self.pairnames = v * self.pairnames
+        return self
+        
+    def __str__(self):
+        return f'{self.__class__.__name__}_{self.split}'
+        
+    def __repr__(self):
+        s = f'{self.__class__.__name__}(split={self.split}, augmentor={self.augmentor_str}, crop_size={str(self.crop_size)}, totensor={self.totensor})'
+        if self.rmul==1:
+            s+=f'\n\tnum pairs: {len(self.pairnames)}'
+        else:
+            s+=f'\n\tnum pairs: {len(self.pairnames)} ({len(self.pairnames)//self.rmul}x{self.rmul})'
+        return s
+
+    def _set_root(self):
+        self.root = dataset_to_root[self.name]
+        assert os.path.isdir(self.root), f"could not find root directory for dataset {self.name}: {self.root}"       
+
+    def _load_or_build_cache(self):
+        cache_file = osp.join(cache_dir, self.name+'.pkl')
+        if osp.isfile(cache_file):
+            with open(cache_file, 'rb') as fid:
+                self.pairnames = pickle.load(fid)[self.split]
+        else:
+            tosave = self._build_cache()
+            os.makedirs(cache_dir, exist_ok=True)
+            with open(cache_file, 'wb') as fid:
+                pickle.dump(tosave, fid)
+            self.pairnames = tosave[self.split]
+        
+class CREStereoDataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = 'CREStereo'
+        self._set_root()
+        assert self.split in ['train']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname+'_left.jpg')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname+'_right.jpg')
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, pairname+'_left.disp.png')
+        self.pairname_to_str = lambda pairname: pairname
+        self.load_disparity = _read_crestereo_disp
+        
+    
+    def _build_cache(self):
+        allpairs = [s+'/'+f[:-len('_left.jpg')] for s in sorted(os.listdir(self.root)) for f in sorted(os.listdir(self.root+'/'+s)) if f.endswith('_left.jpg')]
+        assert len(allpairs)==200000, "incorrect parsing of pairs in CreStereo"
+        tosave = {'train': allpairs}
+        return tosave
+        
+class SceneFlowDataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "SceneFlow"
+        self._set_root()
+        assert self.split in ['train_finalpass','train_cleanpass','train_allpass','test_finalpass','test_cleanpass','test_allpass','test1of100_cleanpass','test1of100_finalpass']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname).replace('/left/','/right/')
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, pairname).replace('/frames_finalpass/','/disparity/').replace('/frames_cleanpass/','/disparity/')[:-4]+'.pfm'
+        self.pairname_to_str = lambda pairname: pairname[:-4]
+        self.load_disparity = _read_sceneflow_disp
+        
+    def _build_cache(self):
+        trainpairs = []
+        # driving
+        pairs = sorted(glob(self.root+'Driving/frames_finalpass/*/*/*/left/*.png'))
+        pairs = list(map(lambda x: x[len(self.root):], pairs))
+        assert len(pairs) == 4400, "incorrect parsing of pairs in SceneFlow"
+        trainpairs += pairs
+        # monkaa
+        pairs = sorted(glob(self.root+'Monkaa/frames_finalpass/*/left/*.png'))
+        pairs = list(map(lambda x: x[len(self.root):], pairs))
+        assert len(pairs) == 8664, "incorrect parsing of pairs in SceneFlow"
+        trainpairs += pairs
+        # flyingthings
+        pairs = sorted(glob(self.root+'FlyingThings/frames_finalpass/TRAIN/*/*/left/*.png'))
+        pairs = list(map(lambda x: x[len(self.root):], pairs))
+        assert len(pairs) == 22390, "incorrect parsing of pairs in SceneFlow"
+        trainpairs += pairs
+        assert len(trainpairs) == 35454, "incorrect parsing of pairs in SceneFlow"
+        testpairs = sorted(glob(self.root+'FlyingThings/frames_finalpass/TEST/*/*/left/*.png'))
+        testpairs = list(map(lambda x: x[len(self.root):], testpairs))
+        assert len(testpairs) == 4370, "incorrect parsing of pairs in SceneFlow"
+        test1of100pairs = testpairs[::100]
+        assert len(test1of100pairs) == 44, "incorrect parsing of pairs in SceneFlow"
+        # all 
+        tosave = {'train_finalpass': trainpairs,
+                  'train_cleanpass': list(map(lambda x: x.replace('frames_finalpass','frames_cleanpass'), trainpairs)),
+                  'test_finalpass': testpairs,
+                  'test_cleanpass': list(map(lambda x: x.replace('frames_finalpass','frames_cleanpass'), testpairs)),
+                  'test1of100_finalpass': test1of100pairs,
+                  'test1of100_cleanpass': list(map(lambda x: x.replace('frames_finalpass','frames_cleanpass'), test1of100pairs)),
+                 }
+        tosave['train_allpass'] = tosave['train_finalpass']+tosave['train_cleanpass']
+        tosave['test_allpass'] = tosave['test_finalpass']+tosave['test_cleanpass']
+        return tosave
+   
+class Md21Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Middlebury2021"
+        self._set_root()
+        assert self.split in ['train','subtrain','subval']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname.replace('/im0','/im1'))
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, pairname.split('/')[0], 'disp0.pfm')
+        self.pairname_to_str = lambda pairname: pairname[:-4]
+        self.load_disparity = _read_middlebury_disp
+        
+    def _build_cache(self):
+        seqs = sorted(os.listdir(self.root))
+        trainpairs = []
+        for s in seqs:
+            #trainpairs += [s+'/im0.png'] # we should remove it, it is included as such in other lightings
+            trainpairs += [s+'/ambient/'+b+'/'+a for b in sorted(os.listdir(osp.join(self.root,s,'ambient'))) for a in sorted(os.listdir(osp.join(self.root,s,'ambient',b))) if a.startswith('im0')]
+        assert len(trainpairs)==355
+        subtrainpairs = [p for p in trainpairs if any(p.startswith(s+'/') for s in seqs[:-2])]
+        subvalpairs = [p for p in trainpairs if any(p.startswith(s+'/') for s in seqs[-2:])]
+        assert len(subtrainpairs)==335 and len(subvalpairs)==20, "incorrect parsing of pairs in Middlebury 2021"
+        tosave = {'train': trainpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs}
+        return tosave 
+
+class Md14Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Middlebury2014"
+        self._set_root()
+        assert self.split in ['train','subtrain','subval']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, osp.dirname(pairname), 'im0.png')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname)
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, osp.dirname(pairname), 'disp0.pfm')
+        self.pairname_to_str = lambda pairname: pairname[:-4]
+        self.load_disparity = _read_middlebury_disp
+        self.has_constant_resolution = False
+        
+    def _build_cache(self):
+        seqs = sorted(os.listdir(self.root))
+        trainpairs = []
+        for s in seqs:
+            trainpairs += [s+'/im1.png',s+'/im1E.png',s+'/im1L.png']
+        assert len(trainpairs)==138
+        valseqs = ['Umbrella-imperfect','Vintage-perfect']
+        assert all(s in seqs for s in valseqs)
+        subtrainpairs = [p for p in trainpairs if not any(p.startswith(s+'/') for s in valseqs)]
+        subvalpairs = [p for p in trainpairs if any(p.startswith(s+'/') for s in valseqs)]
+        assert len(subtrainpairs)==132 and len(subvalpairs)==6, "incorrect parsing of pairs in Middlebury 2014"
+        tosave = {'train': trainpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs}
+        return tosave 
+
+class Md06Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Middlebury2006"
+        self._set_root()
+        assert self.split in ['train','subtrain','subval']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, osp.dirname(pairname), 'view5.png')
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, pairname.split('/')[0], 'disp1.png')
+        self.load_disparity = _read_middlebury20052006_disp
+        self.has_constant_resolution = False
+        
+    def _build_cache(self):
+        seqs = sorted(os.listdir(self.root))
+        trainpairs = []
+        for s in seqs:
+            for i in ['Illum1','Illum2','Illum3']:
+                for e in ['Exp0','Exp1','Exp2']:
+                    trainpairs.append(osp.join(s,i,e,'view1.png'))
+        assert len(trainpairs)==189
+        valseqs = ['Rocks1','Wood2']
+        assert all(s in seqs for s in valseqs)
+        subtrainpairs = [p for p in trainpairs if not any(p.startswith(s+'/') for s in valseqs)]
+        subvalpairs = [p for p in trainpairs if any(p.startswith(s+'/') for s in valseqs)]
+        assert len(subtrainpairs)==171 and len(subvalpairs)==18, "incorrect parsing of pairs in Middlebury 2006"
+        tosave = {'train': trainpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs}
+        return tosave
+
+class Md05Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Middlebury2005"
+        self._set_root()
+        assert self.split in ['train','subtrain','subval']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, osp.dirname(pairname), 'view5.png')
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, pairname.split('/')[0], 'disp1.png')
+        self.pairname_to_str = lambda pairname: pairname[:-4]
+        self.load_disparity = _read_middlebury20052006_disp
+        
+    def _build_cache(self):
+        seqs = sorted(os.listdir(self.root))
+        trainpairs = []
+        for s in seqs:
+            for i in ['Illum1','Illum2','Illum3']:
+                for e in ['Exp0','Exp1','Exp2']:
+                    trainpairs.append(osp.join(s,i,e,'view1.png'))
+        assert len(trainpairs)==54, "incorrect parsing of pairs in Middlebury 2005"
+        valseqs = ['Reindeer']
+        assert all(s in seqs for s in valseqs)
+        subtrainpairs = [p for p in trainpairs if not any(p.startswith(s+'/') for s in valseqs)]
+        subvalpairs = [p for p in trainpairs if any(p.startswith(s+'/') for s in valseqs)]
+        assert len(subtrainpairs)==45 and len(subvalpairs)==9, "incorrect parsing of pairs in Middlebury 2005"
+        tosave = {'train': trainpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs}
+        return tosave
+        
+class MdEval3Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "MiddleburyEval3"
+        self._set_root()
+        assert self.split in [s+'_'+r for s in ['train','subtrain','subval','test','all'] for r in ['full','half','quarter']]
+        if self.split.endswith('_full'):
+            self.root = self.root.replace('/MiddEval3','/MiddEval3_F')
+        elif self.split.endswith('_half'):        
+            self.root = self.root.replace('/MiddEval3','/MiddEval3_H')
+        else:
+            assert self.split.endswith('_quarter')
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname, 'im0.png')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname, 'im1.png')
+        self.pairname_to_Ldispname = lambda pairname: None if pairname.startswith('test') else osp.join(self.root, pairname, 'disp0GT.pfm')
+        self.pairname_to_str = lambda pairname: pairname
+        self.load_disparity = _read_middlebury_disp
+        # for submission only
+        self.submission_methodname = "CroCo-Stereo"
+        self.submission_sresolution = 'F' if self.split.endswith('_full') else ('H' if self.split.endswith('_half') else 'Q')
+        
+    def _build_cache(self):
+        trainpairs = ['train/'+s for s in sorted(os.listdir(self.root+'train/'))]
+        testpairs = ['test/'+s for s in sorted(os.listdir(self.root+'test/'))]
+        subvalpairs = trainpairs[-1:]
+        subtrainpairs = trainpairs[:-1]
+        allpairs = trainpairs+testpairs
+        assert len(trainpairs)==15 and len(testpairs)==15 and len(subvalpairs)==1 and len(subtrainpairs)==14 and len(allpairs)==30, "incorrect parsing of pairs in Middlebury Eval v3"
+        tosave = {}
+        for r in ['full','half','quarter']:
+            tosave.update(**{'train_'+r: trainpairs, 'subtrain_'+r: subtrainpairs, 'subval_'+r: subvalpairs, 'test_'+r: testpairs, 'all_'+r: allpairs})
+        return tosave
+        
+    def submission_save_pairname(self, pairname, prediction, outdir, time):
+        assert prediction.ndim==2
+        assert prediction.dtype==np.float32
+        outfile = os.path.join(outdir, pairname.split('/')[0].replace('train','training')+self.submission_sresolution, pairname.split('/')[1], 'disp0'+self.submission_methodname+'.pfm')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        writePFM(outfile, prediction)
+        timefile = os.path.join( os.path.dirname(outfile), "time"+self.submission_methodname+'.txt')
+        with open(timefile, 'w') as fid:
+            fid.write(str(time))
+
+    def finalize_submission(self, outdir):
+        cmd = f'cd {outdir}/; zip -r "{self.submission_methodname}.zip" .'
+        print(cmd)
+        os.system(cmd)
+        print(f'Done. Submission file at {outdir}/{self.submission_methodname}.zip')
+
+class ETH3DLowResDataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "ETH3DLowRes"
+        self._set_root()
+        assert self.split in ['train','test','subtrain','subval','all']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname, 'im0.png')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname, 'im1.png')
+        self.pairname_to_Ldispname = None if self.split=='test' else lambda pairname: None if pairname.startswith('test/') else osp.join(self.root, pairname.replace('train/','train_gt/'), 'disp0GT.pfm')
+        self.pairname_to_str = lambda pairname: pairname
+        self.load_disparity = _read_eth3d_disp
+        self.has_constant_resolution = False
+        
+    def _build_cache(self):
+        trainpairs = ['train/' + s for s in sorted(os.listdir(self.root+'train/'))]
+        testpairs = ['test/' + s for s in sorted(os.listdir(self.root+'test/'))]
+        assert len(trainpairs) == 27 and len(testpairs) == 20, "incorrect parsing of pairs in ETH3D Low Res"
+        subvalpairs = ['train/delivery_area_3s','train/electro_3l','train/playground_3l']
+        assert all(p in trainpairs for p in subvalpairs)
+        subtrainpairs = [p for p in trainpairs if not p in subvalpairs]
+        assert len(subvalpairs)==3 and len(subtrainpairs)==24, "incorrect parsing of pairs in ETH3D Low Res"
+        tosave = {'train': trainpairs, 'test': testpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs, 'all': trainpairs+testpairs}
+        return tosave
+
+    def submission_save_pairname(self, pairname, prediction, outdir, time):
+        assert prediction.ndim==2
+        assert prediction.dtype==np.float32
+        outfile = os.path.join(outdir, 'low_res_two_view', pairname.split('/')[1]+'.pfm')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        writePFM(outfile, prediction)
+        timefile = outfile[:-4]+'.txt'
+        with open(timefile, 'w') as fid:
+            fid.write('runtime '+str(time))
+
+    def finalize_submission(self, outdir):
+        cmd = f'cd {outdir}/; zip -r "eth3d_low_res_two_view_results.zip" low_res_two_view'
+        print(cmd)
+        os.system(cmd)
+        print(f'Done. Submission file at {outdir}/eth3d_low_res_two_view_results.zip')
+
+class BoosterDataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Booster"
+        self._set_root()
+        assert self.split in ['train_balanced','test_balanced','subtrain_balanced','subval_balanced'] # we use only the balanced version
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname)
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname).replace('/camera_00/','/camera_02/')
+        self.pairname_to_Ldispname = lambda pairname: osp.join(self.root, osp.dirname(pairname), '../disp_00.npy') # same images with different colors, same gt per sequence
+        self.pairname_to_str = lambda pairname: pairname[:-4].replace('/camera_00/','/')
+        self.load_disparity = _read_booster_disp
+        
+        
+    def _build_cache(self):
+        trainseqs = sorted(os.listdir(self.root+'train/balanced'))
+        trainpairs = ['train/balanced/'+s+'/camera_00/'+imname for s in trainseqs for imname in sorted(os.listdir(self.root+'train/balanced/'+s+'/camera_00/'))]
+        testpairs = ['test/balanced/'+s+'/camera_00/'+imname for s in sorted(os.listdir(self.root+'test/balanced')) for imname in sorted(os.listdir(self.root+'test/balanced/'+s+'/camera_00/'))]
+        assert len(trainpairs) == 228 and len(testpairs) == 191
+        subtrainpairs = [p for p in trainpairs if any(s in p for s in trainseqs[:-2])]
+        subvalpairs = [p for p in trainpairs if any(s in p for s in trainseqs[-2:])]
+        # warning: if we do validation split, we should split scenes!!!
+        tosave = {'train_balanced': trainpairs, 'test_balanced': testpairs, 'subtrain_balanced': subtrainpairs, 'subval_balanced': subvalpairs,}
+        return tosave
+        
+class SpringDataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Spring"
+        self._set_root()
+        assert self.split in ['train', 'test', 'subtrain', 'subval']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname+'.png')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname+'.png').replace('frame_right','<frame_right>').replace('frame_left','frame_right').replace('<frame_right>','frame_left')
+        self.pairname_to_Ldispname = lambda pairname: None if pairname.startswith('test') else osp.join(self.root, pairname+'.dsp5').replace('frame_left','disp1_left').replace('frame_right','disp1_right')
+        self.pairname_to_str = lambda pairname: pairname
+        self.load_disparity = _read_hdf5_disp        
+        
+    def _build_cache(self):
+        trainseqs = sorted(os.listdir( osp.join(self.root,'train')))
+        trainpairs = [osp.join('train',s,'frame_left',f[:-4]) for s in trainseqs for f in sorted(os.listdir(osp.join(self.root,'train',s,'frame_left')))]
+        testseqs = sorted(os.listdir( osp.join(self.root,'test')))
+        testpairs = [osp.join('test',s,'frame_left',f[:-4]) for s in testseqs for f in sorted(os.listdir(osp.join(self.root,'test',s,'frame_left')))]
+        testpairs += [p.replace('frame_left','frame_right') for p in testpairs]
+        """maxnorm = {'0001': 32.88, '0002': 228.5, '0004': 298.2, '0005': 142.5, '0006': 113.6, '0007': 27.3, '0008': 554.5, '0009': 155.6, '0010': 126.1, '0011': 87.6, '0012': 303.2, '0013': 24.14, '0014': 82.56, '0015': 98.44, '0016': 156.9, '0017': 28.17, '0018': 21.03, '0020': 178.0, '0021': 58.06, '0022': 354.2, '0023': 8.79, '0024': 97.06, '0025': 55.16, '0026': 91.9, '0027': 156.6, '0030': 200.4, '0032': 58.66, '0033': 373.5, '0036': 149.4, '0037': 5.625, '0038': 37.0, '0039': 12.2, '0041': 453.5, '0043': 457.0, '0044': 379.5, '0045': 161.8, '0047': 105.44} # => let'use 0041"""
+        subtrainpairs = [p for p in trainpairs if p.split('/')[1]!='0041']
+        subvalpairs = [p for p in trainpairs if p.split('/')[1]=='0041']
+        assert len(trainpairs)==5000 and len(testpairs)==2000 and len(subtrainpairs)==4904 and len(subvalpairs)==96, "incorrect parsing of pairs in Spring"
+        tosave = {'train': trainpairs, 'test': testpairs, 'subtrain': subtrainpairs, 'subval': subvalpairs}
+        return tosave
+        
+    def submission_save_pairname(self, pairname, prediction, outdir, time):
+        assert prediction.ndim==2
+        assert prediction.dtype==np.float32
+        outfile = os.path.join(outdir, pairname+'.dsp5').replace('frame_left','disp1_left').replace('frame_right','disp1_right')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        writeDsp5File(prediction, outfile)
+        
+    def finalize_submission(self, outdir):
+        assert self.split=='test'
+        exe = "{self.root}/disp1_subsampling"
+        if os.path.isfile(exe):
+            cmd = f'cd "{outdir}/test"; {exe} .'
+            print(cmd)
+            os.system(cmd)
+        else:
+            print('Could not find disp1_subsampling executable for submission.')
+            print('Please download it and run:')
+            print(f'cd "{outdir}/test"; <disp1_subsampling_exe> .')
+
+class Kitti12Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Kitti12"
+        self._set_root()
+        assert self.split in ['train','test']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname+'_10.png')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname.replace('/colored_0/','/colored_1/')+'_10.png')
+        self.pairname_to_Ldispname = None if self.split=='test' else lambda pairname: osp.join(self.root, pairname.replace('/colored_0/','/disp_occ/')+'_10.png')
+        self.pairname_to_str = lambda pairname: pairname.replace('/colored_0/','/')
+        self.load_disparity = _read_kitti_disp
+        
+    def _build_cache(self):
+        trainseqs = ["training/colored_0/%06d"%(i) for i in range(194)]
+        testseqs = ["testing/colored_0/%06d"%(i) for i in range(195)]
+        assert len(trainseqs)==194 and len(testseqs)==195, "incorrect parsing of pairs in Kitti12"
+        tosave = {'train': trainseqs, 'test': testseqs}
+        return tosave 
+
+    def submission_save_pairname(self, pairname, prediction, outdir, time):
+        assert prediction.ndim==2
+        assert prediction.dtype==np.float32
+        outfile = os.path.join(outdir, pairname.split('/')[-1]+'_10.png')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        img = (prediction * 256).astype('uint16')
+        Image.fromarray(img).save(outfile)
+
+    def finalize_submission(self, outdir):
+        assert self.split=='test'
+        cmd = f'cd {outdir}/; zip -r "kitti12_results.zip" .'
+        print(cmd)
+        os.system(cmd)
+        print(f'Done. Submission file at {outdir}/kitti12_results.zip')
+
+class Kitti15Dataset(StereoDataset):
+
+    def _prepare_data(self):
+        self.name = "Kitti15"
+        self._set_root()
+        assert self.split in ['train','subtrain','subval','test']
+        self.pairname_to_Limgname = lambda pairname: osp.join(self.root, pairname+'_10.png')
+        self.pairname_to_Rimgname = lambda pairname: osp.join(self.root, pairname.replace('/image_2/','/image_3/')+'_10.png')
+        self.pairname_to_Ldispname = None if self.split=='test' else lambda pairname: osp.join(self.root, pairname.replace('/image_2/','/disp_occ_0/')+'_10.png')
+        self.pairname_to_str = lambda pairname: pairname.replace('/image_2/','/')
+        self.load_disparity = _read_kitti_disp
+        
+    def _build_cache(self):
+        trainseqs = ["training/image_2/%06d"%(i) for i in range(200)]
+        subtrainseqs = trainseqs[:-5]
+        subvalseqs = trainseqs[-5:]
+        testseqs = ["testing/image_2/%06d"%(i) for i in range(200)]
+        assert len(trainseqs)==200 and len(subtrainseqs)==195 and len(subvalseqs)==5 and len(testseqs)==200, "incorrect parsing of pairs in Kitti15"
+        tosave = {'train': trainseqs, 'subtrain': subtrainseqs, 'subval': subvalseqs, 'test': testseqs}
+        return tosave 
+
+    def submission_save_pairname(self, pairname, prediction, outdir, time):
+        assert prediction.ndim==2
+        assert prediction.dtype==np.float32
+        outfile = os.path.join(outdir, 'disp_0', pairname.split('/')[-1]+'_10.png')
+        os.makedirs( os.path.dirname(outfile), exist_ok=True)
+        img = (prediction * 256).astype('uint16')
+        Image.fromarray(img).save(outfile)
+
+    def finalize_submission(self, outdir):
+        assert self.split=='test'
+        cmd = f'cd {outdir}/; zip -r "kitti15_results.zip" disp_0'
+        print(cmd)
+        os.system(cmd)
+        print(f'Done. Submission file at {outdir}/kitti15_results.zip')
+
+
+### auxiliary functions
+
+def _read_img(filename):
+    # convert to RGB for scene flow finalpass data
+    img = np.asarray(Image.open(filename).convert('RGB'))
+    return img
+
+def _read_booster_disp(filename):
+    disp = np.load(filename)
+    disp[disp==0.0] = np.inf
+    return disp
+
+def _read_png_disp(filename, coef=1.0):
+    disp = np.asarray(Image.open(filename))
+    disp = disp.astype(np.float32) / coef
+    disp[disp==0.0] = np.inf
+    return disp 
+
+def _read_pfm_disp(filename):
+    disp = np.ascontiguousarray(_read_pfm(filename)[0])
+    disp[disp<=0] = np.inf # eg /nfs/data/ffs-3d/datasets/middlebury/2014/Shopvac-imperfect/disp0.pfm
+    return disp
+
+def _read_npy_disp(filename):
+    return np.load(filename)
+
+def _read_crestereo_disp(filename): return _read_png_disp(filename, coef=32.0)
+def _read_middlebury20052006_disp(filename): return _read_png_disp(filename, coef=1.0)
+def _read_kitti_disp(filename): return _read_png_disp(filename, coef=256.0)
+_read_sceneflow_disp = _read_pfm_disp
+_read_eth3d_disp = _read_pfm_disp
+_read_middlebury_disp = _read_pfm_disp
+_read_carla_disp = _read_pfm_disp
+_read_tartanair_disp = _read_npy_disp
+    
+def _read_hdf5_disp(filename):
+    disp = np.asarray(h5py.File(filename)['disparity'])
+    disp[np.isnan(disp)] = np.inf # make invalid values as +inf
+    #disp[disp==0.0] = np.inf # make invalid values as +inf
+    return disp.astype(np.float32)
+    
+import re
+def _read_pfm(file):
+    file = open(file, 'rb')
+
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().rstrip()
+    if header.decode("ascii") == 'PF':
+        color = True
+    elif header.decode("ascii") == 'Pf':
+        color = False
+    else:
+        raise Exception('Not a PFM file.')
+
+    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode("ascii"))
+    if dim_match:
+        width, height = list(map(int, dim_match.groups()))
+    else:
+        raise Exception('Malformed PFM header.')
+
+    scale = float(file.readline().decode("ascii").rstrip())
+    if scale < 0:  # little-endian
+        endian = '<'
+        scale = -scale
+    else:
+        endian = '>'  # big-endian
+
+    data = np.fromfile(file, endian + 'f')
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    return data, scale
+
+def writePFM(file, image, scale=1):
+    file = open(file, 'wb')
+
+    color = None
+
+    if image.dtype.name != 'float32':
+        raise Exception('Image dtype must be float32.')
+
+    image = np.flipud(image)
+
+    if len(image.shape) == 3 and image.shape[2] == 3:  # color image
+        color = True
+    elif len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1:  # greyscale
+        color = False
+    else:
+        raise Exception('Image must have H x W x 3, H x W x 1 or H x W dimensions.')
+
+    file.write('PF\n' if color else 'Pf\n'.encode())
+    file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0]))
+
+    endian = image.dtype.byteorder
+
+    if endian == '<' or endian == '=' and sys.byteorder == 'little':
+        scale = -scale
+
+    file.write('%f\n'.encode() % scale)
+
+    image.tofile(file)
+
+def writeDsp5File(disp, filename):
+    with h5py.File(filename, "w") as f:
+        f.create_dataset("disparity", data=disp, compression="gzip", compression_opts=5)
+
+
+# disp visualization
+
+def vis_disparity(disp, m=None, M=None):
+    if m is None: m = disp.min()
+    if M is None: M = disp.max()
+    disp_vis = (disp - m) / (M-m) * 255.0
+    disp_vis = disp_vis.astype("uint8")
+    disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_INFERNO)
+    return disp_vis
+
+# dataset getter 
+    
+def get_train_dataset_stereo(dataset_str, augmentor=True, crop_size=None):
+    dataset_str = dataset_str.replace('(','Dataset(')
+    if augmentor:
+        dataset_str = dataset_str.replace(')',', augmentor=True)')
+    if crop_size is not None:
+        dataset_str = dataset_str.replace(')',', crop_size={:s})'.format(str(crop_size)))
+    return eval(dataset_str)
+    
+def get_test_datasets_stereo(dataset_str):
+    dataset_str = dataset_str.replace('(','Dataset(')
+    return [eval(s) for s in dataset_str.split('+')]