def add_getter(self, keys, getter):
"""Register a getter function
Args:
keys (string or tuple of strings): The name(s) of data
that the getter function returns.
getter (callable): A getter function that takes an index and
returns data of the corresponding example.
"""
self._getters.append(getter)
if _is_iterable(keys):
for key_index, key in enumerate(keys):
self._keys.append((key, len(self._getters) - 1, key_index))
else:
self._keys.append((keys, len(self._getters) - 1, None))
def __init__(self, dataset, keys, transform=None):
if transform is None:
keys, transform = None, keys
super(TransformDataset, self).__init__()
self._dataset = dataset
self._transform = transform
if keys is None:
sample = self._get(0)
if isinstance(sample, tuple):
keys = (None, ) * len(sample)
else:
keys = None
self.add_getter(keys, self._get)
if not _is_iterable(keys):
self.keys = 0
def keys(self, keys):
self._keys = [
self._keys[key_index]
for key_index in _as_key_indices(keys, self.keys)
]
self._return_tuple = _is_iterable(keys)