# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
"""All datasets are inherited from this class."""
import importlib
import json
import os
import pickle
from collections import OrderedDict
from functools import partial
from inspect import signature
import numpy as np
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from imaginaire.datasets.folder import FolderDataset
from imaginaire.datasets.lmdb import \
IMG_EXTENSIONS, HDR_IMG_EXTENSIONS, LMDBDataset
from imaginaire.datasets.object_store import ObjectStoreDataset
from imaginaire.utils.data import \
(VIDEO_EXTENSIONS, Augmentor,
load_from_folder, load_from_lmdb, load_from_object_store)
from imaginaire.utils.lmdb import create_metadata
DATASET_TYPES = ['lmdb', 'folder', 'object_store']
[docs]class BaseDataset(data.Dataset):
r"""Base class for image/video datasets.
Args:
cfg (Config object): Input config.
is_inference (bool): Training if False, else validation.
is_test (bool): Final test set after training and validation.
"""
def __init__(self, cfg, is_inference, is_test):
super(BaseDataset, self).__init__()
self.cfg = cfg
self.is_inference = is_inference
self.is_test = is_test
if self.is_test:
self.cfgdata = self.cfg.test_data
data_info = self.cfgdata.test
else:
self.cfgdata = self.cfg.data
if self.is_inference:
data_info = self.cfgdata.val
else:
data_info = self.cfgdata.train
self.name = self.cfgdata.name
self.lmdb_roots = data_info.roots
self.dataset_type = getattr(data_info, 'dataset_type', None)
self.cache = getattr(self.cfgdata, 'cache', None)
self.interpolator = getattr(self.cfgdata, 'interpolator', "INTER_LINEAR")
# Get AWS secret keys.
if self.dataset_type == 'object_store':
assert hasattr(cfg, 'aws_credentials_file')
self.aws_credentials_file = cfg.aws_credentials_file
# Legacy lmdb/folder only support.
if self.dataset_type is None:
self.dataset_is_lmdb = getattr(data_info, 'is_lmdb', False)
if self.dataset_is_lmdb:
self.dataset_type = 'lmdb'
else:
self.dataset_type = 'folder'
# Legacy support ends.
assert self.dataset_type in DATASET_TYPES
if self.dataset_type == 'lmdb':
# Add handle to function to load data from LMDB.
self.load_from_dataset = load_from_lmdb
elif self.dataset_type == 'folder':
# For some unpaired experiments, we would like the dataset to be presented in a paired way
if hasattr(self.cfgdata, 'paired') is False:
self.cfgdata.paired = self.paired
# Add handle to function to load data from folder.
self.load_from_dataset = load_from_folder
# Create metadata for folders.
print('Creating metadata')
all_filenames, all_metadata = [], []
if self.is_test:
cfg.data_backup = cfg.data
cfg.data = cfg.test_data
for root in self.lmdb_roots:
filenames, metadata = create_metadata(
data_root=root, cfg=cfg, paired=self.cfgdata['paired'])
all_filenames.append(filenames)
all_metadata.append(metadata)
if self.is_test:
cfg.data = cfg.data_backup
elif self.dataset_type == 'object_store':
# Add handle to function to load data from AWS S3.
self.load_from_dataset = load_from_object_store
# Get the types of data stored in dataset, and their extensions.
self.data_types = [] # Names of data types.
self.dataset_data_types = [] # These data types are in the dataset.
self.image_data_types = [] # These types are images.
self.hdr_image_data_types = [] # These types are HDR images.
self.normalize = {} # Does this data type need normalization?
self.extensions = {} # What is this data type's file extension.
self.is_mask = {} # Whether this data type is discrete masks?
self.num_channels = {} # How many channels does this data type have?
self.pre_aug_ops = {} # Ops on data type before augmentation.
self.post_aug_ops = {} # Ops on data type after augmentation.
# Extract info from data types.
for data_type in self.cfgdata.input_types:
name = list(data_type.keys())
assert len(name) == 1
name = name[0]
info = data_type[name]
if 'ext' not in info:
info['ext'] = None
if 'normalize' not in info:
info['normalize'] = False
if 'is_mask' not in info:
info['is_mask'] = False
if 'pre_aug_ops' not in info:
info['pre_aug_ops'] = 'None'
if 'post_aug_ops' not in info:
info['post_aug_ops'] = 'None'
if 'computed_on_the_fly' not in info:
info['computed_on_the_fly'] = False
if 'num_channels' not in info:
info['num_channels'] = None
self.data_types.append(name)
if not info['computed_on_the_fly']:
self.dataset_data_types.append(name)
self.extensions[name] = info['ext']
self.normalize[name] = info['normalize']
self.num_channels[name] = info['num_channels']
self.pre_aug_ops[name] = [op.strip() for op in
info['pre_aug_ops'].split(',')]
self.post_aug_ops[name] = [op.strip() for op in
info['post_aug_ops'].split(',')]
self.is_mask[name] = info['is_mask']
if info['ext'] is not None and (info['ext'] in IMG_EXTENSIONS or info['ext'] in VIDEO_EXTENSIONS):
self.image_data_types.append(name)
if info['ext'] is not None and info['ext'] in HDR_IMG_EXTENSIONS:
self.hdr_image_data_types.append(name)
# Add some info into cfgdata for legacy support.
self.cfgdata.data_types = self.data_types
self.cfgdata.num_channels = [self.num_channels[name]
for name in self.data_types]
# Augmentations which need full dict.
self.full_data_post_aug_ops, self.full_data_ops = [], []
if hasattr(self.cfgdata, 'full_data_ops'):
ops = self.cfgdata.full_data_ops
self.full_data_ops.extend([op.strip() for op in ops.split(',')])
if hasattr(self.cfgdata, 'full_data_post_aug_ops'):
ops = self.cfgdata.full_data_post_aug_ops
self.full_data_post_aug_ops.extend(
[op.strip() for op in ops.split(',')])
# These are the labels which will be concatenated for generator input.
self.input_labels = []
if hasattr(self.cfgdata, 'input_labels'):
self.input_labels = self.cfgdata.input_labels
# These are the keypoints which also need to be augmented.
self.keypoint_data_types = []
if hasattr(self.cfgdata, 'keypoint_data_types'):
self.keypoint_data_types = self.cfgdata.keypoint_data_types
# Create augmentation operations.
aug_list = data_info.augmentations
individual_video_frame_aug_list = getattr(data_info, 'individual_video_frame_augmentations', dict())
self.augmentor = Augmentor(
aug_list, individual_video_frame_aug_list, self.image_data_types, self.is_mask,
self.keypoint_data_types, self.interpolator)
self.augmentable_types = self.image_data_types + \
self.keypoint_data_types
# Create torch transformations.
self.transform = {}
for data_type in self.image_data_types:
normalize = self.normalize[data_type]
self.transform[data_type] = self._get_transform(
normalize, self.num_channels[data_type])
# Create torch transformations for HDR images.
for data_type in self.hdr_image_data_types:
normalize = self.normalize[data_type]
self.transform[data_type] = self._get_transform(
normalize, self.num_channels[data_type])
# Initialize handles.
self.sequence_lists = [] # List of sequences per dataset root.
self.lmdbs = {} # Dict for list of lmdb handles per data type.
for data_type in self.dataset_data_types:
self.lmdbs[data_type] = []
self.dataset_probability = None
self.additional_lists = []
# Load each dataset.
for idx, root in enumerate(self.lmdb_roots):
if self.dataset_type == 'lmdb':
self._add_dataset(root)
elif self.dataset_type == 'folder':
self._add_dataset(root, filenames=all_filenames[idx],
metadata=all_metadata[idx])
elif self.dataset_type == 'object_store':
self._add_dataset(
root, aws_credentials_file=self.aws_credentials_file)
# Compute dataset statistics and create whatever self.variables required
# for the specific dataloader.
self._compute_dataset_stats()
# Build index of data to sample.
self.mapping, self.epoch_length = self._create_mapping()
def _create_mapping(self):
r"""Creates mapping from data sample idx to actual LMDB keys.
All children need to implement their own.
Returns:
self.mapping (list): List of LMDB keys.
"""
raise NotImplementedError
def _compute_dataset_stats(self):
r"""Computes required statistics about dataset.
All children need to implement their own.
"""
pass
def __getitem__(self, index):
r"""Entry function for dataset."""
raise NotImplementedError
def _get_transform(self, normalize, num_channels):
r"""Convert numpy to torch tensor.
Args:
normalize (bool): Normalize image i.e. (x - 0.5) * 2.
Goes from [0, 1] -> [-1, 1].
Returns:
Composed list of torch transforms.
"""
transform_list = [transforms.ToTensor()]
if normalize:
transform_list.append(
transforms.Normalize((0.5, ) * num_channels,
(0.5, ) * num_channels, inplace=True))
return transforms.Compose(transform_list)
def _add_dataset(self, root, filenames=None, metadata=None,
aws_credentials_file=None):
r"""Adds an LMDB dataset to a list of datasets.
Args:
root (str): Path to LMDB or folder dataset.
filenames: List of filenames for folder dataset.
metadata: Metadata for folder dataset.
aws_credentials_file: Path to file containing AWS credentials.
"""
if aws_credentials_file and self.dataset_type == 'object_store':
object_store_dataset = ObjectStoreDataset(
root, aws_credentials_file, cache=self.cache)
sequence_list = object_store_dataset.sequence_list
else:
# Get sequences associated with this dataset.
if filenames is None:
list_path = 'all_filenames.json'
with open(os.path.join(root, list_path)) as fin:
sequence_list = OrderedDict(json.load(fin))
else:
sequence_list = filenames
additional_path = 'all_indices.json'
if os.path.exists(os.path.join(root, additional_path)):
print('Using additional list for object indices.')
with open(os.path.join(root, additional_path)) as fin:
additional_list = OrderedDict(json.load(fin))
self.additional_lists.append(additional_list)
self.sequence_lists.append(sequence_list)
# Get LMDB dataset handles.
for data_type in self.dataset_data_types:
if self.dataset_type == 'lmdb':
self.lmdbs[data_type].append(
LMDBDataset(os.path.join(root, data_type)))
elif self.dataset_type == 'folder':
self.lmdbs[data_type].append(
FolderDataset(os.path.join(root, data_type), metadata))
elif self.dataset_type == 'object_store':
# All data types use the same handle.
self.lmdbs[data_type].append(object_store_dataset)
[docs] def flip_hdr(self, data, is_flipped=False):
r"""Flip hdr images.
Args:
data (dict): Keys are from data types. Values can be numpy.ndarray
or list of numpy.ndarray (image or list of images).
is_flipped (bool): Applying left-right flip to the hdr images
Returns:
(tuple):
- data (dict): Augmented data, with same keys as input data.
"""
if is_flipped is False:
return data
for data_type in self.hdr_image_data_types:
# print('Length of data: {}'.format(len(data[data_type])))
data[data_type][0] = data[data_type][0][:, ::-1, :].copy()
return data
[docs] def to_tensor(self, data):
r"""Convert all images to tensor.
Args:
data (dict): Dict containing data_type as key, with each value
as a list of numpy.ndarrays.
Returns:
data (dict): Dict containing data_type as key, with each value
as a list of torch.Tensors.
"""
for data_type in self.image_data_types:
for idx in range(len(data[data_type])):
if data[data_type][idx].dtype == np.uint16:
data[data_type][idx] = data[data_type][idx].astype(
np.float32)
data[data_type][idx] = self.transform[data_type](
data[data_type][idx])
for data_type in self.hdr_image_data_types:
for idx in range(len(data[data_type])):
data[data_type][idx] = self.transform[data_type](
data[data_type][idx])
return data
[docs] def apply_ops(self, data, op_dict, full_data=False):
r"""Apply any ops from op_dict to data types.
Args:
data (dict): Dict containing data_type as key, with each value
as a list of numpy.ndarrays.
op_dict (dict): Dict containing data_type as key, with each value
containing string of operations to apply.
full_data (bool): Do these ops require access to the full data?
Returns:
data (dict): Dict containing data_type as key, with each value
modified by the op if any.
"""
if full_data:
# op needs entire data dict.
for op in op_dict:
if op == 'None':
continue
op, op_type = self.get_op(op)
assert op_type == 'full_data'
data = op(data)
else:
# op per data type.
if not op_dict:
return data
for data_type in data:
for op in op_dict[data_type]:
if op == 'None':
continue
op, op_type = self.get_op(op)
data[data_type] = op(data[data_type])
if op_type == 'vis':
# We have converted this data type to an image. Enter it
# in self.image_data_types and give it a torch
# transform.
if data_type not in self.image_data_types:
self.image_data_types.append(data_type)
normalize = self.normalize[data_type]
num_channels = self.num_channels[data_type]
self.transform[data_type] = \
self._get_transform(normalize, num_channels)
elif op_type == 'convert':
continue
elif op_type is None:
continue
else:
raise NotImplementedError
return data
[docs] def get_op(self, op):
r"""Get function to apply for specific op.
Args:
op (str): Name of the op.
Returns:
function handle.
"""
def list_to_tensor(data):
r"""Convert list of numeric values to tensor."""
assert isinstance(data, list)
return torch.from_numpy(np.array(data, dtype=np.float32))
def decode_json_list(data):
r"""Decode list of strings in json to objects."""
assert isinstance(data, list)
return [json.loads(item) for item in data]
def decode_pkl_list(data):
r"""Decode list of pickled strings to objects."""
assert isinstance(data, list)
return [pickle.loads(item) for item in data]
def list_to_numpy(data):
r"""Convert list of numeric values to numpy array."""
assert isinstance(data, list)
return np.array(data)
def l2_normalize(data):
r"""L2 normalization."""
assert isinstance(data, torch.Tensor)
import torch.nn.functional as F
return F.normalize(data, dim=1)
if op == 'to_tensor':
return list_to_tensor, None
elif op == 'decode_json':
return decode_json_list, None
elif op == 'decode_pkl':
return decode_pkl_list, None
elif op == 'to_numpy':
return list_to_numpy, None
elif op == 'l2_norm':
return l2_normalize, None
elif '::' in op:
parts = op.split('::')
if len(parts) == 2:
module, function = parts
module = importlib.import_module(module)
function = getattr(module, function)
sig = signature(function)
num_params = len(sig.parameters)
assert num_params in [3, 4], \
'Full data functions take in (cfgdata, is_inference, ' \
'full_data) or (cfgdata, is_inference, self, full_data) ' \
'as input.'
if num_params == 3:
function = partial(
function, self.cfgdata, self.is_inference)
elif num_params == 4:
function = partial(
function, self.cfgdata, self.is_inference, self)
function_type = 'full_data'
elif len(parts) == 3:
function_type, module, function = parts
module = importlib.import_module(module)
# Get function inputs, if provided.
partial_fn = False
if '(' in function and ')' in function:
partial_fn = True
function, params = self._get_fn_params(function)
function = getattr(module, function)
# Create partial function.
if partial_fn:
function = partial(function, **params)
# Get function signature.
sig = signature(function)
num_params = 0
for param in sig.parameters.values():
if param.kind == param.POSITIONAL_OR_KEYWORD:
num_params += 1
if function_type == 'vis':
if num_params != 9:
raise ValueError(
'vis function type needs to take ' +
'(resize_h, resize_w, crop_h, crop_w, ' +
'original_h, original_w, is_flipped, cfgdata, ' +
'data) as input.')
function = partial(function,
self.augmentor.resize_h,
self.augmentor.resize_w,
self.augmentor.crop_h,
self.augmentor.crop_w,
self.augmentor.original_h,
self.augmentor.original_w,
self.augmentor.is_flipped,
self.cfgdata)
elif function_type == 'convert':
if num_params != 1:
raise ValueError(
'convert function type needs to take ' +
'(data) as input.')
else:
raise ValueError('Unknown op: %s' % (op))
else:
raise ValueError('Unknown op: %s' % (op))
return function, function_type
else:
raise ValueError('Unknown op: %s' % (op))
def _get_fn_params(self, function_string):
r"""Find key-value inputs to function from string definition.
Args:
function_string (str): String with function name and args. e.g.
my_function(a=10, b=20).
Returns:
function (str): Name of function.
params (dict): Key-value params for function.
"""
start = function_string.find('(')
end = function_string.find(')')
function = function_string[:start]
params_str = function_string[start+1:end]
params = {}
for item in params_str.split(':'):
key, value = item.split('=')
try:
params[key] = float(value)
except: # noqa
params[key] = value
return function, params
def __len__(self):
return self.epoch_length