iicd/RAFT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

added upsampling module

Browse Source
This commit is contained in:
Zach Teed
2020-07-25 17:36:17 -06:00
parent dc1220825d
commit a2408eab78
32 changed files with 23559 additions and 619 deletions
Download Patch File Download Diff File Expand all files Collapse all files

367
core/datasets.py
View File

@@ -6,53 +6,42 @@ import torch.utils.data as data
import torch.nn.functional as F
import os
import cv2
import math
import random
from glob import glob
import os.path as osp
from utils import frame_utils
from utils.augmentor import FlowAugmentor, FlowAugmentorKITTI
from utils.augmentor import FlowAugmentor, SparseFlowAugmentor
class CombinedDataset(data.Dataset):
def __init__(self, datasets):
self.datasets = datasets
def __len__(self):
length = 0
for i in range(len(self.datasets)):
length += len(self.datsaets[i])
return length
def __getitem__(self, index):
i = 0
for j in range(len(self.datasets)):
if i + len(self.datasets[j]) >= index:
yield self.datasets[j][index-i]
break
i += len(self.datasets[j])
def __add__(self, other):
self.datasets.append(other)
return self
class FlowDataset(data.Dataset):
def __init__(self, args, image_size=None, do_augument=False):
self.image_size = image_size
self.do_augument = do_augument
if self.do_augument:
self.augumentor = FlowAugmentor(self.image_size)
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.init_seed = False
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:
@@ -62,133 +51,96 @@ class FlowDataset(data.Dataset):
self.init_seed = True
index = index % len(self.image_list)
flow = frame_utils.read_gen(self.flow_list[index])
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])
img1 = np.array(img1).astype(np.uint8)[..., :3]
img2 = np.array(img2).astype(np.uint8)[..., :3]
flow = np.array(flow).astype(np.float32)
img1 = np.array(img1).astype(np.uint8)
img2 = np.array(img2).astype(np.uint8)
if self.do_augument:
img1, img2, flow = self.augumentor(img1, img2, flow)
# grayscale images
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()
valid = torch.ones_like(flow[0])
return img1, img2, flow, valid
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)
def __add(self, other):
return CombinedDataset([self, other])
class MpiSintelTest(FlowDataset):
def __init__(self, args, root='datasets/Sintel/test', dstype='clean'):
super(MpiSintelTest, self).__init__(args, image_size=None, do_augument=False)
self.root = root
self.dstype = dstype
image_dir = osp.join(self.root, dstype)
all_sequences = os.listdir(image_dir)
self.image_list = []
for sequence in all_sequences:
frames = sorted(glob(osp.join(image_dir, sequence, '*.png')))
for i in range(len(frames)-1):
self.image_list += [[frames[i], frames[i+1], sequence, i]]
def __getitem__(self, index):
img1 = frame_utils.read_gen(self.image_list[index][0])
img2 = frame_utils.read_gen(self.image_list[index][1])
sequence = self.image_list[index][2]
frame = self.image_list[index][3]
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, sequence, frame
class MpiSintel(FlowDataset):
def __init__(self, args, image_size=None, do_augument=True, root='datasets/Sintel/training', dstype='clean'):
super(MpiSintel, self).__init__(args, image_size, do_augument)
if do_augument:
self.augumentor.min_scale = -0.2
self.augumentor.max_scale = 0.7
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)
self.root = root
self.dstype = dstype
if split == 'test':
self.is_test = True
flow_root = osp.join(root, 'flow')
image_root = osp.join(root, dstype)
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) ] # scene and frame_id
file_list = sorted(glob(osp.join(flow_root, '*/*.flo')))
for flo in file_list:
fbase = flo[len(flow_root)+1:]
fprefix = fbase[:-8]
fnum = int(fbase[-8:-4])
img1 = osp.join(image_root, fprefix + "%04d"%(fnum+0) + '.png')
img2 = osp.join(image_root, fprefix + "%04d"%(fnum+1) + '.png')
if not osp.isfile(img1) or not osp.isfile(img2) or not osp.isfile(flo):
continue
self.image_list.append((img1, img2))
self.flow_list.append(flo)
if split != 'test':
self.flow_list += sorted(glob(osp.join(flow_root, scene, '*.flo')))
class FlyingChairs(FlowDataset):
def __init__(self, args, image_size=None, do_augument=True, root='datasets/FlyingChairs_release/data'):
super(FlyingChairs, self).__init__(args, image_size, do_augument)
self.root = root
self.augumentor.min_scale = -0.2
self.augumentor.max_scale = 1.0
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')))
self.flow_list = sorted(glob(osp.join(root, '*.flo')))
assert (len(images)//2 == len(self.flow_list))
flows = sorted(glob(osp.join(root, '*.flo')))
assert (len(images)//2 == len(flows))
self.image_list = []
for i in range(len(self.flow_list)):
im1 = images[2*i]
im2 = images[2*i + 1]
self.image_list.append([im1, im2])
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 SceneFlow(FlowDataset):
def __init__(self, args, image_size, do_augument=True, root='datasets',
dstype='frames_cleanpass', use_flyingthings=True, use_monkaa=False, use_driving=False):
super(SceneFlow, self).__init__(args, image_size, do_augument)
self.root = root
self.dstype = dstype
self.augumentor.min_scale = -0.2
self.augumentor.max_scale = 0.8
if use_flyingthings:
self.add_flyingthings()
if use_monkaa:
self.add_monkaa()
if use_driving:
self.add_driving()
def add_flyingthings(self):
root = osp.join(self.root, 'FlyingThings3D')
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, self.dstype, 'TRAIN/*/*')))
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/*/*')))
@@ -199,114 +151,85 @@ class SceneFlow(FlowDataset):
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]]
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]]
self.image_list += [ [images[i+1], images[i]] ]
self.flow_list += [ flows[i+1] ]
def add_monkaa(self):
pass # we don't use monkaa
def add_driving(self):
pass # we don't use driving
class KITTI(FlowDataset):
def __init__(self, args, image_size=None, do_augument=True, is_test=False, is_val=False, do_pad=False, split=True, root='datasets/KITTI'):
super(KITTI, self).__init__(args, image_size, do_augument)
self.root = root
self.is_test = is_test
self.is_val = is_val
self.do_pad = do_pad
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
if self.do_augument:
self.augumentor = FlowAugmentorKITTI(self.image_size, min_scale=-0.2, max_scale=0.5)
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')))
if self.is_test:
images1 = sorted(glob(os.path.join(root, 'testing', 'image_2/*_10.png')))
images2 = sorted(glob(os.path.join(root, 'testing', 'image_2/*_11.png')))
for i in range(len(images1)):
self.image_list += [[images1[i], images2[i]]]
for img1, img2 in zip(images1, images2):
frame_id = img1.split('/')[-1]
self.extra_info += [ [frame_id] ]
self.image_list += [ [img1, img2] ]
else:
flows = sorted(glob(os.path.join(root, 'training', 'flow_occ/*_10.png')))
images1 = sorted(glob(os.path.join(root, 'training', 'image_2/*_10.png')))
images2 = sorted(glob(os.path.join(root, 'training', 'image_2/*_11.png')))
if split == 'training':
self.flow_list = sorted(glob(osp.join(root, 'flow_occ/*_10.png')))
for i in range(len(flows)):
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 += [[images1[i], images2[i]]]
self.image_list += [ [images[i], images[i+1]] ]
seq_ix += 1
def __getitem__(self, index):
def fetch_dataloader(args, TRAIN_DS='C+T+K+S+H'):
""" Create the data loader for the corresponding trainign set """
if self.is_test:
frame_id = self.image_list[index][0]
frame_id = frame_id.split('/')[-1]
if args.stage == 'chairs':
aug_params = {'crop_size': args.image_size, 'min_scale': -0.2, '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
img1 = frame_utils.read_gen(self.image_list[index][0])
img2 = frame_utils.read_gen(self.image_list[index][1])
elif args.stage == 'sintel':
aug_params = {'crop_size': args.image_size, 'min_scale': -0.3, 'max_scale': 0.7, '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')
img1 = np.array(img1).astype(np.uint8)[..., :3]
img2 = np.array(img2).astype(np.uint8)[..., :3]
if TRAIN_DS == 'C+T+K+S+H':
kitti = KITTI({'crop_size': args.image_size, 'min_scale': -0.3, 'max_scale': 0.7, 'do_flip': True})
hd1k = HD1K({'crop_size': args.image_size, 'min_scale': -0.5, 'max_scale': 0.5, 'do_flip': True})
train_dataset = 100*sintel_clean + 100*sintel_final + 200*kitti + 5*hd1k + things
img1 = torch.from_numpy(img1).permute(2, 0, 1).float()
img2 = torch.from_numpy(img2).permute(2, 0, 1).float()
return img1, img2, frame_id
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(args, image_size=args.image_size, is_val=False)
else:
if not self.init_seed:
worker_info = torch.utils.data.get_worker_info()
if worker_info is not None:
np.random.seed(worker_info.id)
random.seed(worker_info.id)
self.init_seed = True
train_loader = data.DataLoader(train_dataset, batch_size=args.batch_size,
pin_memory=False, shuffle=True, num_workers=4, drop_last=True)
index = index % len(self.image_list)
frame_id = self.image_list[index][0]
frame_id = frame_id.split('/')[-1]
print('Training with %d image pairs' % len(train_dataset))
return train_loader
img1 = frame_utils.read_gen(self.image_list[index][0])
img2 = frame_utils.read_gen(self.image_list[index][1])
flow, valid = frame_utils.readFlowKITTI(self.flow_list[index])
img1 = np.array(img1).astype(np.uint8)[..., :3]
img2 = np.array(img2).astype(np.uint8)[..., :3]
if self.do_augument:
img1, img2, flow, valid = self.augumentor(img1, img2, flow, valid)
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()
valid = torch.from_numpy(valid).float()
if self.do_pad:
ht, wd = img1.shape[1:]
pad_ht = (((ht // 8) + 1) * 8 - ht) % 8
pad_wd = (((wd // 8) + 1) * 8 - wd) % 8
pad_ht1 = [0, pad_ht]
pad_wd1 = [pad_wd//2, pad_wd - pad_wd//2]
pad = pad_wd1 + pad_ht1
img1 = img1.view(1, 3, ht, wd)
img2 = img2.view(1, 3, ht, wd)
flow = flow.view(1, 2, ht, wd)
valid = valid.view(1, 1, ht, wd)
img1 = torch.nn.functional.pad(img1, pad, mode='replicate')
img2 = torch.nn.functional.pad(img2, pad, mode='replicate')
flow = torch.nn.functional.pad(flow, pad, mode='constant', value=0)
valid = torch.nn.functional.pad(valid, pad, mode='replicate', value=0)
img1 = img1.view(3, ht+pad_ht, wd+pad_wd)
img2 = img2.view(3, ht+pad_ht, wd+pad_wd)
flow = flow.view(2, ht+pad_ht, wd+pad_wd)
valid = valid.view(ht+pad_ht, wd+pad_wd)
if self.is_test:
return img1, img2, flow, valid, frame_id
return img1, img2, flow, valid