Source code for imaginaire.datasets.base

# 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 perform_individual_video_frame(self, data, augment_ops): r"""Perform data augmentation on images only. Args: data (dict): Keys are from data types. Values can be numpy.ndarray or list of numpy.ndarray (image or list of images). augment_ops (list): The augmentation operations for individual frames. Returns: (tuple): - data (dict): Augmented data, with same keys as input data. - is_flipped (bool): Flag which tells if images have been left-right flipped. """ if augment_ops: all_data = dict() for ix, key in enumerate(data.keys()): if ix == 0: num = len(data[key]) for j in range(num): all_data['%d' % j] = dict() for j in range(num): all_data['%d' % j][key] = data[key][j:(j+1)] for j in range(num): all_data['%d' % j], _ = self.perform_augmentation( all_data['%d' % j], paired=True, augment_ops=augment_ops) for key in data.keys(): tmp = [] for j in range(num): tmp += all_data['%d' % j][key] data[key] = tmp return data
[docs] def perform_augmentation(self, data, paired, augment_ops=None): r"""Perform data augmentation on images only. Args: data (dict): Keys are from data types. Values can be numpy.ndarray or list of numpy.ndarray (image or list of images). paired (bool): Apply same augmentation to all input keys? augment_ops (list): The augmentation operations. Returns: (tuple): - data (dict): Augmented data, with same keys as input data. - is_flipped (bool): Flag which tells if images have been left-right flipped. """ aug_inputs = {} for data_type in self.augmentable_types: aug_inputs[data_type] = data[data_type] augmented, is_flipped = self.augmentor.perform_augmentation( aug_inputs, paired=paired, augment_ops=augment_ops) for data_type in self.augmentable_types: data[data_type] = augmented[data_type] return data, is_flipped
[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