config.py

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")