Spaces:
Running
Running
| # 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('+')] |