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