Spaces:
Sleeping
Sleeping
# Data loading based on https://github.com/NVIDIA/flownet2-pytorch | |
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 model.raft.core.utils import frame_utils | |
from model.raft.core.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) | |
# 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() | |
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) ] # scene and frame_id | |
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 | |