|
|
|
|
| import numpy as np
|
| import torch
|
| import torch.utils.data as data
|
| import torch.nn.functional as F
|
|
|
| import os
|
| import math
|
| import random
|
| from glob import glob
|
| import os.path as osp
|
|
|
| from .utils import frame_utils
|
| from .utils.augmentor import FlowAugmentor, SparseFlowAugmentor
|
|
|
|
|
| class FlowDataset(data.Dataset):
|
| def __init__(self, aug_params=None, sparse=False):
|
| self.augmentor = None
|
| self.sparse = sparse
|
| if aug_params is not None:
|
| if sparse:
|
| self.augmentor = SparseFlowAugmentor(**aug_params)
|
| else:
|
| self.augmentor = FlowAugmentor(**aug_params)
|
|
|
| self.is_test = False
|
| self.init_seed = False
|
| self.flow_list = []
|
| self.image_list = []
|
| self.extra_info = []
|
|
|
| def __getitem__(self, index):
|
|
|
| if self.is_test:
|
| img1 = frame_utils.read_gen(self.image_list[index][0])
|
| img2 = frame_utils.read_gen(self.image_list[index][1])
|
| img1 = np.array(img1).astype(np.uint8)[..., :3]
|
| img2 = np.array(img2).astype(np.uint8)[..., :3]
|
| img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
|
| img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
|
| return img1, img2, self.extra_info[index]
|
|
|
| if not self.init_seed:
|
| worker_info = torch.utils.data.get_worker_info()
|
| if worker_info is not None:
|
| torch.manual_seed(worker_info.id)
|
| np.random.seed(worker_info.id)
|
| random.seed(worker_info.id)
|
| self.init_seed = True
|
|
|
| index = index % len(self.image_list)
|
| valid = None
|
| if self.sparse:
|
| flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])
|
| else:
|
| flow = frame_utils.read_gen(self.flow_list[index])
|
|
|
| img1 = frame_utils.read_gen(self.image_list[index][0])
|
| img2 = frame_utils.read_gen(self.image_list[index][1])
|
|
|
| flow = np.array(flow).astype(np.float32)
|
| img1 = np.array(img1).astype(np.uint8)
|
| img2 = np.array(img2).astype(np.uint8)
|
|
|
|
|
| if len(img1.shape) == 2:
|
| img1 = np.tile(img1[...,None], (1, 1, 3))
|
| img2 = np.tile(img2[...,None], (1, 1, 3))
|
| else:
|
| img1 = img1[..., :3]
|
| img2 = img2[..., :3]
|
|
|
| if self.augmentor is not None:
|
| if self.sparse:
|
| img1, img2, flow, valid = self.augmentor(img1, img2, flow, valid)
|
| else:
|
| img1, img2, flow = self.augmentor(img1, img2, flow)
|
|
|
| img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
|
| img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
|
| flow = torch.from_numpy(flow).permute(2, 0, 1).float()
|
|
|
| if valid is not None:
|
| valid = torch.from_numpy(valid)
|
| else:
|
| valid = (flow[0].abs() < 1000) & (flow[1].abs() < 1000)
|
|
|
| return img1, img2, flow, valid.float()
|
|
|
|
|
| def __rmul__(self, v):
|
| self.flow_list = v * self.flow_list
|
| self.image_list = v * self.image_list
|
| return self
|
|
|
| def __len__(self):
|
| return len(self.image_list)
|
|
|
|
|
| class MpiSintel(FlowDataset):
|
| def __init__(self, aug_params=None, split='training', root='datasets/Sintel', dstype='clean'):
|
| super(MpiSintel, self).__init__(aug_params)
|
| flow_root = osp.join(root, split, 'flow')
|
| image_root = osp.join(root, split, dstype)
|
|
|
| if split == 'test':
|
| self.is_test = True
|
|
|
| for scene in os.listdir(image_root):
|
| image_list = sorted(glob(osp.join(image_root, scene, '*.png')))
|
| for i in range(len(image_list)-1):
|
| self.image_list += [ [image_list[i], image_list[i+1]] ]
|
| self.extra_info += [ (scene, i) ]
|
|
|
| if split != 'test':
|
| self.flow_list += sorted(glob(osp.join(flow_root, scene, '*.flo')))
|
|
|
|
|
| class FlyingChairs(FlowDataset):
|
| def __init__(self, aug_params=None, split='train', root='datasets/FlyingChairs_release/data'):
|
| super(FlyingChairs, self).__init__(aug_params)
|
|
|
| images = sorted(glob(osp.join(root, '*.ppm')))
|
| flows = sorted(glob(osp.join(root, '*.flo')))
|
| assert (len(images)//2 == len(flows))
|
|
|
| split_list = np.loadtxt('chairs_split.txt', dtype=np.int32)
|
| for i in range(len(flows)):
|
| xid = split_list[i]
|
| if (split=='training' and xid==1) or (split=='validation' and xid==2):
|
| self.flow_list += [ flows[i] ]
|
| self.image_list += [ [images[2*i], images[2*i+1]] ]
|
|
|
|
|
| class FlyingThings3D(FlowDataset):
|
| def __init__(self, aug_params=None, root='datasets/FlyingThings3D', dstype='frames_cleanpass'):
|
| super(FlyingThings3D, self).__init__(aug_params)
|
|
|
| for cam in ['left']:
|
| for direction in ['into_future', 'into_past']:
|
| image_dirs = sorted(glob(osp.join(root, dstype, 'TRAIN/*/*')))
|
| image_dirs = sorted([osp.join(f, cam) for f in image_dirs])
|
|
|
| flow_dirs = sorted(glob(osp.join(root, 'optical_flow/TRAIN/*/*')))
|
| flow_dirs = sorted([osp.join(f, direction, cam) for f in flow_dirs])
|
|
|
| for idir, fdir in zip(image_dirs, flow_dirs):
|
| images = sorted(glob(osp.join(idir, '*.png')) )
|
| flows = sorted(glob(osp.join(fdir, '*.pfm')) )
|
| for i in range(len(flows)-1):
|
| if direction == 'into_future':
|
| self.image_list += [ [images[i], images[i+1]] ]
|
| self.flow_list += [ flows[i] ]
|
| elif direction == 'into_past':
|
| self.image_list += [ [images[i+1], images[i]] ]
|
| self.flow_list += [ flows[i+1] ]
|
|
|
|
|
| class KITTI(FlowDataset):
|
| def __init__(self, aug_params=None, split='training', root='datasets/KITTI'):
|
| super(KITTI, self).__init__(aug_params, sparse=True)
|
| if split == 'testing':
|
| self.is_test = True
|
|
|
| root = osp.join(root, split)
|
| images1 = sorted(glob(osp.join(root, 'image_2/*_10.png')))
|
| images2 = sorted(glob(osp.join(root, 'image_2/*_11.png')))
|
|
|
| for img1, img2 in zip(images1, images2):
|
| frame_id = img1.split('/')[-1]
|
| self.extra_info += [ [frame_id] ]
|
| self.image_list += [ [img1, img2] ]
|
|
|
| if split == 'training':
|
| self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png')))
|
|
|
|
|
| class HD1K(FlowDataset):
|
| def __init__(self, aug_params=None, root='datasets/HD1k'):
|
| super(HD1K, self).__init__(aug_params, sparse=True)
|
|
|
| seq_ix = 0
|
| while 1:
|
| flows = sorted(glob(os.path.join(root, 'hd1k_flow_gt', 'flow_occ/%06d_*.png' % seq_ix)))
|
| images = sorted(glob(os.path.join(root, 'hd1k_input', 'image_2/%06d_*.png' % seq_ix)))
|
|
|
| if len(flows) == 0:
|
| break
|
|
|
| for i in range(len(flows)-1):
|
| self.flow_list += [flows[i]]
|
| self.image_list += [ [images[i], images[i+1]] ]
|
|
|
| seq_ix += 1
|
|
|
|
|
| def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'):
|
| """ Create the data loader for the corresponding trainign set """
|
|
|
| if args.stage == 'chairs':
|
| aug_params = {'crop_size': args.image_size, 'min_scale': -0.1, 'max_scale': 1.0, 'do_flip': True}
|
| train_dataset = FlyingChairs(aug_params, split='training')
|
|
|
| elif args.stage == 'things':
|
| aug_params = {'crop_size': args.image_size, 'min_scale': -0.4, 'max_scale': 0.8, 'do_flip': True}
|
| clean_dataset = FlyingThings3D(aug_params, dstype='frames_cleanpass')
|
| final_dataset = FlyingThings3D(aug_params, dstype='frames_finalpass')
|
| train_dataset = clean_dataset + final_dataset
|
|
|
| elif args.stage == 'sintel':
|
| aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.6, 'do_flip': True}
|
| things = FlyingThings3D(aug_params, dstype='frames_cleanpass')
|
| sintel_clean = MpiSintel(aug_params, split='training', dstype='clean')
|
| sintel_final = MpiSintel(aug_params, split='training', dstype='final')
|
|
|
| if TRAIN_DS == 'C+T+K+S+H':
|
| kitti = KITTI({'crop_size': args.image_size, 'min_scale': -0.3, 'max_scale': 0.5, 'do_flip': True})
|
| hd1k = HD1K({'crop_size': args.image_size, 'min_scale': -0.5, 'max_scale': 0.2, 'do_flip': True})
|
| train_dataset = 100*sintel_clean + 100*sintel_final + 200*kitti + 5*hd1k + things
|
|
|
| elif TRAIN_DS == 'C+T+K/S':
|
| train_dataset = 100*sintel_clean + 100*sintel_final + things
|
|
|
| elif args.stage == 'kitti':
|
| aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, 'max_scale': 0.4, 'do_flip': False}
|
| train_dataset = KITTI(aug_params, split='training')
|
|
|
| train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size,
|
| pin_memory=False, shuffle=True, num_workers=4, drop_last=True)
|
|
|
| print('Training with %d image pairs' % len(train_dataset))
|
| return train_loader
|
|
|
|
|
