import copy import itertools import logging import os from pathlib import Path import numpy as np import torch.utils.data from detectron2.config import CfgNode as CN import utils from datasets import FlowPairDetectron, FlowEvalDetectron logger = logging.getLogger('gwm') def scan_train_flow(folders, res, pairs, basepath): pair_list = [p for p in itertools.combinations(pairs, 2)] flow_dir = {} for pair in pair_list: p1, p2 = pair flowpairs = [] for f in folders: path1 = basepath / f'Flows_gap{p1}' / res / f path2 = basepath / f'Flows_gap{p2}' / res / f flows1 = [p.name for p in path1.glob('*.flo')] flows2 = [p.name for p in path2.glob('*.flo')] flows1 = sorted(flows1) flows2 = sorted(flows2) intersect = list(set(flows1).intersection(flows2)) intersect.sort() flowpair = np.array([[path1 / i, path2 / i] for i in intersect]) flowpairs += [flowpair] flow_dir['gap_{}_{}'.format(p1, p2)] = flowpairs # flow_dir is a dictionary, with keys indicating the flow gap, and each value is a list of sequence names, # each item then is an array with Nx2, N indicates the number of available pairs. return flow_dir def setup_dataset(cfg=None, multi_val=False): dataset_str = cfg.GWM.DATASET if '+' in dataset_str: datasets = dataset_str.split('+') logger.info(f'Multiple datasets detected: {datasets}') train_datasets = [] val_datasets = [] for ds in datasets: proxy_cfg = copy.deepcopy(cfg) proxy_cfg.merge_from_list(['GWM.DATASET', ds]), train_ds, val_ds = setup_dataset(proxy_cfg, multi_val=multi_val) train_datasets.append(train_ds) val_datasets.append(val_ds) logger.info(f'Multiple datasets detected: {datasets}') logger.info(f'Validation is still : {datasets[0]}') return torch.utils.data.ConcatDataset(train_datasets), val_datasets[0] resolution = cfg.GWM.RESOLUTION # h,w res = "" with_gt = True pairs = [1, 2, -1, -2] trainval_data_dir = None if cfg.GWM.DATASET == 'DAVIS': basepath = '/DAVIS2016' img_dir = '/DAVIS2016/JPEGImages/480p' gt_dir = '/DAVIS2016/Annotations/480p' val_flow_dir = '/DAVIS2016/Flows_gap1/1080p' val_seq = ['dog', 'cows', 'goat', 'camel', 'libby', 'parkour', 'soapbox', 'blackswan', 'bmx-trees', 'kite-surf', 'car-shadow', 'breakdance', 'dance-twirl', 'scooter-black', 'drift-chicane', 'motocross-jump', 'horsejump-high', 'drift-straight', 'car-roundabout', 'paragliding-launch'] val_data_dir = [val_flow_dir, img_dir, gt_dir] res = "1080p" elif cfg.GWM.DATASET in ['FBMS']: basepath = '/FBMS_clean' img_dir = '/FBMS_clean/JPEGImages/' gt_dir = '/FBMS_clean/Annotations/' val_flow_dir = '/FBMS_val/Flows_gap1/' val_seq = ['camel01', 'cars1', 'cars10', 'cars4', 'cars5', 'cats01', 'cats03', 'cats06', 'dogs01', 'dogs02', 'farm01', 'giraffes01', 'goats01', 'horses02', 'horses04', 'horses05', 'lion01', 'marple12', 'marple2', 'marple4', 'marple6', 'marple7', 'marple9', 'people03', 'people1', 'people2', 'rabbits02', 'rabbits03', 'rabbits04', 'tennis'] val_img_dir = '/FBMS_val/JPEGImages/' val_gt_dir = '/FBMS_val/Annotations/' val_data_dir = [val_flow_dir, val_img_dir, val_gt_dir] with_gt = False pairs = [3, 6, -3, -6] elif cfg.GWM.DATASET in ['STv2']: basepath = '/SegTrackv2' img_dir = '/SegTrackv2/JPEGImages' gt_dir = '/SegTrackv2/Annotations' val_flow_dir = '/SegTrackv2/Flows_gap1/' val_seq = ['drift', 'birdfall', 'girl', 'cheetah', 'worm', 'parachute', 'monkeydog', 'hummingbird', 'soldier', 'bmx', 'frog', 'penguin', 'monkey', 'bird_of_paradise'] val_data_dir = [val_flow_dir, img_dir, gt_dir] else: raise ValueError('Unknown Setting/Dataset.') # Switching this section to pathlib, which should prevent double // errors in paths and dict keys root_path_str = cfg.GWM.DATA_ROOT logger.info(f"Found DATA_ROOT in config: {root_path_str}") root_path_str = '../data' if root_path_str.startswith('/'): root_path = Path(f"/{root_path_str.lstrip('/').rstrip('/')}") else: root_path = Path(f"{root_path_str.lstrip('/').rstrip('/')}") logger.info(f"Loading dataset from: {root_path}") basepath = root_path / basepath.lstrip('/').rstrip('/') img_dir = root_path / img_dir.lstrip('/').rstrip('/') gt_dir = root_path / gt_dir.lstrip('/').rstrip('/') val_data_dir = [root_path / path.lstrip('/').rstrip('/') for path in val_data_dir] folders = [p.name for p in (basepath / f'Flows_gap{pairs[0]}' / res).iterdir() if p.is_dir()] folders = sorted(folders) # flow_dir is a dictionary, with keys indicating the flow gap, and each value is a list of sequence names, # each item then is an array with Nx2, N indicates the number of available pairs. flow_dir = scan_train_flow(folders, res, pairs, basepath) data_dir = [flow_dir, img_dir, gt_dir] force1080p = ('DAVIS' not in cfg.GWM.DATASET) and 'RGB_BIG' in cfg.GWM.SAMPLE_KEYS enable_photometric_augmentations = cfg.FLAGS.INF_TPS train_dataset = FlowPairDetectron(data_dir=data_dir, resolution=resolution, to_rgb=cfg.GWM.FLOW2RGB, size_divisibility=cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY if not cfg.FLAGS.IGNORE_SIZE_DIV else -1, enable_photo_aug=enable_photometric_augmentations, flow_clip=cfg.GWM.FLOW_CLIP, norm=cfg.GWM.FLOW_NORM, force1080p=force1080p, flow_res=cfg.GWM.FLOW_RES, ) if multi_val: print(f"Using multiple validation datasets from {val_data_dir}") val_dataset = [FlowEvalDetectron(data_dir=val_data_dir, resolution=resolution, pair_list=pairs, val_seq=[vs], to_rgb=cfg.GWM.FLOW2RGB, with_rgb=False, size_divisibility=cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY if not cfg.FLAGS.IGNORE_SIZE_DIV else -1, flow_clip=cfg.GWM.FLOW_CLIP, norm=cfg.GWM.FLOW_NORM, force1080p=force1080p) for vs in val_seq] for vs, vds in zip(val_seq, val_dataset): print(f"Validation dataset for {vs}: {len(vds)}") if len(vds) == 0: raise ValueError(f"Empty validation dataset for {vs}") if cfg.GWM.TTA_AS_TRAIN: if trainval_data_dir is None: trainval_data_dir = val_data_dir else: trainval_data_dir = [root_path / path.lstrip('/').rstrip('/') for path in trainval_data_dir] trainval_dataset = [] tvd_basepath = root_path / str(trainval_data_dir[0].relative_to(root_path)).split('/')[0] print("TVD BASE DIR", tvd_basepath) for vs in val_seq: tvd_data_dir = [scan_train_flow([vs], res, pairs, tvd_basepath), *trainval_data_dir[1:]] tvd = FlowPairDetectron(data_dir=tvd_data_dir, resolution=resolution, to_rgb=cfg.GWM.FLOW2RGB, size_divisibility=cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY if not cfg.FLAGS.IGNORE_SIZE_DIV else -1, enable_photo_aug=cfg.GWM.LOSS_MULT.EQV is not None, flow_clip=cfg.GWM.FLOW_CLIP, norm=cfg.GWM.FLOW_NORM, force1080p=force1080p, flow_res=cfg.GWM.FLOW_RES, ) trainval_dataset.append(tvd) print(f'Seq {trainval_data_dir[0]}/{vs} dataset: {len(tvd)}') else: if trainval_data_dir is None: trainval_dataset = val_dataset else: trainval_data_dir = [root_path / path.lstrip('/').rstrip('/') for path in trainval_data_dir] trainval_dataset = [] for vs in val_seq: tvd = FlowEvalDetectron(data_dir=trainval_data_dir, resolution=resolution, pair_list=pairs, val_seq=[vs], to_rgb=cfg.GWM.FLOW2RGB, with_rgb=False, size_divisibility=cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY if not cfg.FLAGS.IGNORE_SIZE_DIV else -1, flow_clip=cfg.GWM.FLOW_CLIP, norm=cfg.GWM.FLOW_NORM, force1080p=force1080p) trainval_dataset.append(tvd) print(f'Seq {trainval_data_dir[0]}/{vs} dataset: {len(tvd)}') return train_dataset, val_dataset, trainval_dataset val_dataset = FlowEvalDetectron(data_dir=val_data_dir, resolution=resolution, pair_list=pairs, val_seq=val_seq, to_rgb=cfg.GWM.FLOW2RGB, with_rgb=False, size_divisibility=cfg.MODEL.MASK_FORMER.SIZE_DIVISIBILITY if not cfg.FLAGS.IGNORE_SIZE_DIV else -1, flow_clip=cfg.GWM.FLOW_CLIP, norm=cfg.GWM.FLOW_NORM, force1080p=force1080p) return train_dataset, val_dataset def loaders(cfg): train_dataset, val_dataset = setup_dataset(cfg) logger.info(f"Sourcing data from {val_dataset.data_dir[0]}") if cfg.FLAGS.DEV_DATA: subset = cfg.SOLVER.IMS_PER_BATCH * 3 train_dataset = torch.utils.data.Subset(train_dataset, list(range(subset))) val_dataset = torch.utils.data.Subset(val_dataset, list(range(subset))) g = torch.Generator() data_generator_seed = int(torch.randint(int(1e6), (1,)).item()) logger.info(f"Dataloaders generator seed {data_generator_seed}") g.manual_seed(data_generator_seed) train_loader = torch.utils.data.DataLoader(train_dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_size=cfg.SOLVER.IMS_PER_BATCH, collate_fn=lambda x: x, shuffle=True, pin_memory=True, drop_last=True, persistent_workers=cfg.DATALOADER.NUM_WORKERS > 0, worker_init_fn=utils.random_state.worker_init_function, generator=g ) val_loader = torch.utils.data.DataLoader(val_dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_size=1, shuffle=False, pin_memory=True, collate_fn=lambda x: x, drop_last=False, persistent_workers=cfg.DATALOADER.NUM_WORKERS > 0, worker_init_fn=utils.random_state.worker_init_function, generator=g) return train_loader, val_loader def multi_loaders(cfg): train_dataset, val_datasets, train_val_datasets = setup_dataset(cfg, multi_val=True) logger.info(f"Sourcing multiple loaders from {len(val_datasets)}") logger.info(f"Sourcing data from {val_datasets[0].data_dir[0]}") g = torch.Generator() data_generator_seed = int(torch.randint(int(1e6), (1,)).item()) logger.info(f"Dataloaders generator seed {data_generator_seed}") g.manual_seed(data_generator_seed) train_loader = torch.utils.data.DataLoader(train_dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_size=cfg.SOLVER.IMS_PER_BATCH, collate_fn=lambda x: x, shuffle=True, pin_memory=True, drop_last=True, persistent_workers=cfg.DATALOADER.NUM_WORKERS > 0, worker_init_fn=utils.random_state.worker_init_function, generator=g ) val_loaders = [(torch.utils.data.DataLoader(val_dataset, num_workers=0, batch_size=1, shuffle=False, pin_memory=True, collate_fn=lambda x: x, drop_last=False, persistent_workers=False, worker_init_fn=utils.random_state.worker_init_function, generator=g), torch.utils.data.DataLoader(tv_dataset, num_workers=0, batch_size=cfg.SOLVER.IMS_PER_BATCH, shuffle=True, pin_memory=False, collate_fn=lambda x: x, drop_last=False, persistent_workers=False, worker_init_fn=utils.random_state.worker_init_function, generator=g)) for val_dataset, tv_dataset in zip(val_datasets, train_val_datasets)] return train_loader, val_loaders def add_gwm_config(cfg): cfg.GWM = CN() cfg.GWM.MODEL = "MASKFORMER" cfg.GWM.RESOLUTION = (128, 224) cfg.GWM.FLOW_RES = (480, 854) cfg.GWM.SAMPLE_KEYS = ["rgb"] cfg.GWM.ADD_POS_EMB = False cfg.GWM.CRITERION = "L2" cfg.GWM.L1_OPTIMIZE = False cfg.GWM.HOMOGRAPHY = 'quad' # False cfg.GWM.HOMOGRAPHY_SUBSAMPLE = 8 cfg.GWM.HOMOGRAPHY_SKIP = 0.4 cfg.GWM.DATASET = 'DAVIS' cfg.GWM.DATA_ROOT = None cfg.GWM.FLOW2RGB = False cfg.GWM.SIMPLE_REC = False cfg.GWM.DAVIS_SINGLE_VID = None cfg.GWM.USE_MULT_FLOW = False cfg.GWM.FLOW_COLORSPACE_REC = None cfg.GWM.FLOW_CLIP_U_LOW = float('-inf') cfg.GWM.FLOW_CLIP_U_HIGH = float('inf') cfg.GWM.FLOW_CLIP_V_LOW = float('-inf') cfg.GWM.FLOW_CLIP_V_HIGH = float('inf') cfg.GWM.FLOW_CLIP = float('inf') cfg.GWM.FLOW_NORM = False cfg.GWM.LOSS_MULT = CN() cfg.GWM.LOSS_MULT.REC = 0.03 cfg.GWM.LOSS_MULT.HEIR_W = [0.1, 0.3, 0.6] cfg.GWM.TTA = 100 # Test-time-adaptation cfg.GWM.TTA_AS_TRAIN = False # Use train-like data logic for test-time-adaptation cfg.GWM.LOSS = 'OG' cfg.FLAGS = CN() cfg.FLAGS.MAKE_VIS_VIDEOS = False # Making videos is kinda slow cfg.FLAGS.EXTENDED_FLOW_RECON_VIS = False # Does not cost much cfg.FLAGS.COMP_NLL_FOR_GT = False # Should we log loss against ground truth? cfg.FLAGS.DEV_DATA = False cfg.FLAGS.KEEP_ALL = True # Keep all checkoints cfg.FLAGS.ORACLE_CHECK = False # Use oracle check to estimate max performance when grouping multiple components cfg.FLAGS.INF_TPS = False # cfg.FLAGS.UNFREEZE_AT = [(1, 10000), (0, 20000), (-1, 30000)] cfg.FLAGS.UNFREEZE_AT = [(4, 0), (2, 500), (1, 1000), (-1, 10000)] cfg.FLAGS.IGNORE_SIZE_DIV = False cfg.FLAGS.IGNORE_TMP = True cfg.WANDB = CN() cfg.WANDB.ENABLE = False cfg.WANDB.BASEDIR = '../' cfg.DEBUG = False cfg.LOG_ID = 'exp' cfg.LOG_FREQ = 250 cfg.OUTPUT_BASEDIR = '../outputs' cfg.SLURM = False cfg.SKIP_TB = False cfg.TOTAL_ITER = 20000 cfg.CONFIG_FILE = None if os.environ.get('SLURM_JOB_ID', None): cfg.LOG_ID = os.environ.get('SLURM_JOB_NAME', cfg.LOG_ID) logger.info(f"Setting name {cfg.LOG_ID} based on SLURM job name")