| |
|
|
| 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 core.utils import frame_utils |
| from core.utils.augmentor import FlowAugmentor, SparseFlowAugmentor |
|
|
| import sys |
|
|
| 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 = [] |
|
|
| self.occ_list = [] |
| self.line_list = [] |
|
|
| def __getitem__(self, index): |
| |
| |
| index = index % len(self.image_list) |
| if self.is_test: |
| self.augmentor = None |
|
|
| valid = None |
| 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)) |
| else: |
| img1 = img1[..., :3] |
| |
| if len(img2.shape) == 2: |
| img2 = np.tile(img2[...,None], (1, 1, 3)) |
| else: |
| img2 = img2[..., :3] |
|
|
| if self.augmentor is not None: |
| 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) |
|
|
| if self.is_test: |
| occ = torch.from_numpy(np.load(self.occ_list[index])) |
| line = torch.from_numpy(np.load(self.line_list[index])) |
| return img1, img2, flow, occ, line, self.extra_info[index], valid.float() |
| else: |
| 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='/mnt/lustre/syli/AnimeRun/flow/data/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 AnimeRun(FlowDataset): |
| def __init__(self, aug_params=None, split='train', root='/mnt/lustre/syli/AnimeRun/flow/data/AnimeRun_v2', dstype='Frame_Anime/original'): |
| super(AnimeRun, self).__init__(aug_params) |
| |
| flow_root = osp.join(root, split, 'Flow') |
| image_root = osp.join(root, split, dstype) |
| unmatch_root = osp.join(root, split, 'UnmatchedForward') |
| line_root = osp.join(root, split, 'LineArea') |
|
|
| if split == 'test': |
| self.is_test = True |
|
|
| for scene in os.listdir(image_root): |
| for color_pass in os.listdir(osp.join(image_root, scene)): |
| if color_pass != 'original': |
| continue |
| image_list = sorted(glob(osp.join(image_root, scene, color_pass, '*.png'))) |
| for i in range(len(image_list)-1): |
| self.image_list += [ [image_list[i], image_list[i+1]] ] |
| self.extra_info += [ (scene, i) ] |
| self.occ_list += sorted(glob(osp.join(unmatch_root, scene, '*.npy'))) |
| self.flow_list += sorted(glob(osp.join(flow_root, scene, 'forward', '*.flo'))) |
| self.line_list += [ path.replace(unmatch_root, line_root) for path in sorted(glob(osp.join(unmatch_root, scene, '*.npy')))] |
|
|
| print('Len of Flow is ', len(self.flow_list)) |
| print('Len of Anime is ', len(self.image_list)) |
| print('Len of Occlusion is ', len(self.occ_list)) |
| print('Len of Line Area is ', len(self.line_list)) |
|
|
|
|
|
|
| class FlyingThings3D(FlowDataset): |
| def __init__(self, aug_params=None, root='/mnt/lustre/syli/AnimeRun/flow/data/FlyingThings', 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 CreativeFlow(FlowDataset): |
| def __init__(self, aug_params=None, split='train', root='/mnt/lustre/syli/AnimeRun/flow/data/CreativeFlow/decompressed'): |
| super(CreativeFlow, self).__init__(aug_params) |
|
|
| flow_root = osp.join(root, split) |
| image_root = osp.join(root, split) |
| for scene in os.listdir(image_root): |
| if split == 'test': |
| render_folder = os.listdir(osp.join(image_root, scene, 'cam0', 'renders', 'composite'))[0] |
| image_list = sorted(glob(osp.join(image_root, scene, 'cam0', 'renders', 'composite', render_folder, '*.png'))) |
| flow_list = sorted(glob(osp.join(flow_root, scene, 'cam0', 'metadata', 'flow', '*.flo'))) |
| else: |
| render_folder = os.listdir(osp.join(image_root, scene, 'renders', 'composite'))[0] |
| image_list = sorted(glob(osp.join(image_root, scene, 'renders', 'composite', render_folder, '*.png'))) |
| flow_list = sorted(glob(osp.join(flow_root, scene, 'metadata', 'flow', '*.flo'))) |
|
|
| if len(image_list)-1 == len(flow_list): |
| for i in range(len(image_list)-1): |
| self.image_list += [ [image_list[i], image_list[i+1]] ] |
| self.extra_info += [ (scene, i) ] |
| self.flow_list += [ flow_list[i] ] |
| print(len(self.image_list), len(self.extra_info), len(self.flow_list)) |
|
|
|
|
| def fetch_dataloader(args, TRAIN_DS='C+T+K/S'): |
| """ Create the data loader for the corresponding trainign set """ |
|
|
| if args.stage == 'sintel': |
| print('Training Sintel Stage...') |
| sys.stdout.flush() |
| 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') |
| elif args.stage == 'anime': |
| print('Training AnimeRun Stage...') |
| sys.stdout.flush() |
| 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') |
| anime_run = 200 * AnimeRun(aug_params, split='train', dstype='Frame_Anime') |
|
|
| train_dataset = anime_run + things |
|
|
| elif args.stage == 'creative': |
| print('Training Using Creative Flow...') |
| sys.stdout.flush() |
| 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') |
| creative = CreativeFlow(aug_params, split='test') |
|
|
| train_dataset = 20*creative + things |
|
|
| 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 |
|
|
