def _build_dir_tree(self):
"""Build a dict of indices for iterating over the dataset."""
self._dir_tree = collections.OrderedDict()
for path in self._allowed_dirs:
vids = get_subdirs(
path,
nonempty=False,
sort_numerical=True,
)
if vids:
self._dir_tree[path] = vids
self._restrict_dataset_size()
def _build_dir_tree(self):
"""Build a dict of indices for iterating over the dataset."""
self._dir_tree = {}
for path in self._allowed_dirs:
vids = get_subdirs(
path,
nonempty=False,
sort=True,
sortfunc=lambda x: int(osp.splitext(osp.basename(x))[0]),
)
if len(vids) > 0: # pylint: disable=g-explicit-length-test
self._dir_tree[path] = vids
self._restrict_dataset_size()
def __init__(
self,
root_dir,
frame_sampler,
augmentor=None,
max_vids_per_class=-1,
seed=None,
):
"""Constructor.
Args:
root_dir: The path to the dataset directory.
frame_sampler: A sampler specifying the frame sampling strategy.
augmentor: An instance of transforms.VideoAugmentor. If provided, will
apply data augmentation to the sampled video data.
max_vids_per_class: The max number of videos to consider per class. The
remaining videos are ignored, effectively reducing the total dataset
size.
seed: The seed for the rng.
Raises:
ValueError: If the root directory is empty.
"""
super().__init__()
self._root_dir = root_dir
self._frame_sampler = frame_sampler
self._seed = seed
self._max_vids_per_class = max_vids_per_class
self._augmentor = augmentor
self._totensor = ToTensor()
# Get list of available dirs and ensure that it is not empty.
self._available_dirs = get_subdirs(
self._root_dir,
nonempty=True,
sort=False,
)
if len(self._available_dirs) == 0: # pylint: disable=g-explicit-length-test
raise ValueError("{} is an empty directory.".format(root_dir))
self._allowed_dirs = self._available_dirs
self.seed_rng()
self._build_dir_tree()
def dataset_from_config(config, downstream, split, debug):
"""Create a video dataset from a config."""
dataset_path = osp.join(config.data.root, split)
image_size = config.data_augmentation.image_size
if isinstance(image_size, int):
image_size = (image_size, image_size)
image_size = tuple(image_size)
# Note(kevin): We used to disable data augmentation on all downstream
# dataloaders. I've decided to keep them for train downstream loaders.
if debug:
# The minimum data augmentation we want to keep is resizing when
# debugging.
aug_names = ["global_resize"]
else:
if split == "train":
aug_names = config.data_augmentation.train_transforms
else:
aug_names = config.data_augmentation.eval_transforms
# Create a list of data augmentation callables.
aug_funcs = []
for name in aug_names:
if "resize" in name or "crop" in name:
aug_funcs.append(create_transform(name, *image_size))
else:
aug_funcs.append(create_transform(name))
augmentor = transforms.VideoAugmentor({SequenceType.FRAMES: aug_funcs})
# Restrict action classes if they have been provided. Else, load all
# from the data directory.
c_action_class = (config.data.downstream_action_class
if downstream else config.data.pretrain_action_class)
if c_action_class:
action_classes = c_action_class
else:
action_classes = get_subdirs(
dataset_path,
basename=True,
nonempty=True,
sort_lexicographical=True,
)
# We need to separate out the dataclasses for each action class when
# creating downstream datasets.
if downstream:
dataset = {}
for action_class in action_classes:
frame_sampler = frame_sampler_from_config(config, downstream=True)
single_class_dataset = VideoDataset(
dataset_path,
frame_sampler,
seed=config.seed,
augmentor=augmentor,
max_vids_per_class=config.data.max_vids_per_class,
)
single_class_dataset.restrict_subdirs(action_class)
dataset[action_class] = single_class_dataset
else:
frame_sampler = frame_sampler_from_config(config, downstream=False)
dataset = VideoDataset(
dataset_path,
frame_sampler,
seed=config.seed,
augmentor=augmentor,
max_vids_per_class=config.data.max_vids_per_class,
)
dataset.restrict_subdirs(action_classes)
return dataset
def dataset_from_config(config, downstream, split, debug):
"""Create a video dataset from a config."""
dataset_path = osp.join(config.DATA.ROOT, split)
image_size = config.DATA_AUGMENTATION.IMAGE_SIZE
if isinstance(image_size, int):
image_size = (image_size, image_size)
image_size = tuple(image_size)
if debug or downstream:
# The minimum data augmentation we want to keep is resizing when
# debugging.
aug_names = ["global_resize"]
else:
if split == "train":
aug_names = config.DATA_AUGMENTATION.TRAIN_TRANSFORMS
else:
aug_names = config.DATA_AUGMENTATION.EVAL_TRANSFORMS
# Create a list of data augmentation callables.
aug_funcs = []
for name in aug_names:
if "resize" in name or "crop" in name:
aug_funcs.append(create_transform(name, *image_size))
else:
aug_funcs.append(create_transform(name))
augmentor = transforms.VideoAugmentor({SequenceType.FRAMES: aug_funcs})
# Restrict action classes if they have been provided. Else, load all
# from the data directory.
c_action_class = (config.DATA.DOWNSTREAM_ACTION_CLASS
if downstream else config.DATA.PRETRAIN_ACTION_CLASS)
if c_action_class:
action_classes = c_action_class
else:
action_classes = get_subdirs(
dataset_path,
basename=True,
nonempty=True,
sort=False,
)
# We need to separate out the dataclasses for each action class when
# creating downstream datasets.
if downstream:
dataset = {}
for action_class in action_classes:
frame_sampler = frame_sampler_from_config(config, downstream=True)
single_class_dataset = VideoDataset(
dataset_path,
frame_sampler,
seed=config.SEED,
augmentor=augmentor,
max_vids_per_class=config.DATA.MAX_VIDS_PER_CLASS,
)
single_class_dataset.restrict_subdirs(action_class)
dataset[action_class] = single_class_dataset
else:
frame_sampler = frame_sampler_from_config(config, downstream=False)
dataset = VideoDataset(
dataset_path,
frame_sampler,
seed=config.SEED,
augmentor=augmentor,
max_vids_per_class=config.DATA.MAX_VIDS_PER_CLASS,
)
dataset.restrict_subdirs(action_classes)
return dataset