Add SliceableDataset

chainercv/__init__.py

@@ -1,5 +1,6 @@
 import pkg_resources
 
+from chainercv import chainer_experimental  # NOQA
 from chainercv import datasets  # NOQA
 from chainercv import evaluations  # NOQA
 from chainercv import extensions  # NOQA

chainercv/chainer_experimental/__init__.py

@@ -0,0 +1 @@
+from chainercv.chainer_experimental import datasets  # NOQA

chainercv/chainer_experimental/datasets/__init__.py

@@ -0,0 +1 @@
+from chainercv.chainer_experimental.datasets import sliceable  # NOQA

chainercv/chainer_experimental/datasets/sliceable/__init__.py

@@ -0,0 +1,6 @@
+from chainercv.chainer_experimental.datasets.sliceable.sliceable_dataset import SliceableDataset  # NOQA
+
+from chainercv.chainer_experimental.datasets.sliceable.concatenated_dataset import ConcatenatedDataset  # NOQA
+from chainercv.chainer_experimental.datasets.sliceable.getter_dataset import GetterDataset  # NOQA
+from chainercv.chainer_experimental.datasets.sliceable.transform_dataset import TransformDataset  # NOQA
+from chainercv.chainer_experimental.datasets.sliceable.tuple_dataset import TupleDataset  # NOQA

chainercv/chainer_experimental/datasets/sliceable/concatenated_dataset.py

@@ -0,0 +1,45 @@
+from chainercv.chainer_experimental.datasets.sliceable import SliceableDataset
+
+
+class ConcatenatedDataset(SliceableDataset):
+    """A sliceable version of :class:`chainer.datasets.ConcatenatedDataset`.
+
+    Hew is an example.
+
+    >>> dataset_a = TupleDataset([0, 1, 2], [0, 1, 4])
+    >>> dataset_b = TupleDataset([3, 4, 5], [9, 16, 25])
+    >>>
+    >>> dataset = ConcatenatedDataset(dataset_a, dataset_b)
+    >>> dataset.slice[:, 0][:]  # [0, 1, 2, 3, 4, 5]
+
+    Args:
+        datasets: The underlying datasets.
+            Each dataset should inherit
+            :class:~chainer.datasets.sliceable.Sliceabledataset`.
+            and should have the same keys.
+    """
+
+    def __init__(self, *datasets):
+        if len(datasets) == 0:
+            raise ValueError('At least one dataset is required')
+        self._datasets = datasets
+        self._keys = datasets[0].keys
+        for dataset in datasets[1:]:
+            if not dataset.keys == self._keys:
+                raise ValueError('All datasets should have the same keys')
+
+    def __len__(self):
+        return sum(len(dataset) for dataset in self._datasets)
+
+    @property
+    def keys(self):
+        return self._keys
+
+    def get_example_by_keys(self, index, key_indices):
+        if index < 0:
+            raise IndexError
+        for dataset in self._datasets:
+            if index < len(dataset):
+                return dataset.get_example_by_keys(index, key_indices)
+            index -= len(dataset)
+        raise IndexError

chainercv/chainer_experimental/datasets/sliceable/getter_dataset.py

@@ -0,0 +1,91 @@
+from chainercv.chainer_experimental.datasets.sliceable import SliceableDataset
+
+
+def _as_tuple(t):
+    if isinstance(t, tuple):
+        return t
+    else:
+        return t,
+
+
+class GetterDataset(SliceableDataset):
+    """A sliceable dataset class that defined by getters.
+
+    This ia a dataset class with getters.
+
+    Here is an example.
+
+    >>> class SliceableLabeledImageDataset(GetterDataset):
+    >>>     def __init__(self, pairs, root='.'):
+    >>>         super().__init__()
+    >>>         with open(pairs) as f:
+    >>>             self._pairs = [l.split() for l in f]
+    >>>         self._root = root
+    >>>
+    >>>         self.add_getter('image', self.get_image)
+    >>>         self.add_getter('label', self.get_label)
+    >>>
+    >>>     def __len__(self):
+    >>>         return len(self._pairs)
+    >>>
+    >>>     def get_image(self, i):
+    >>>         path, _ = self._pairs[i]
+    >>>         return read_image(os.path.join(self._root, path))
+    >>>
+    >>>     def get_label(self, i):
+    >>>         _, label = self._pairs[i]
+    >>>         return np.int32(label)
+    >>
+    >>> dataset = SliceableLabeledImageDataset('list.txt')
+    >>>
+    >>> # get a subset with label = 0, 1, 2
+    >>> # no images are loaded
+    >>> indices = [i for i, label in
+    >>> enumerate(dataset.slice[:, 'label']) if label in {0, 1, 2}]
+    >>> dataset_012 = dataset.slice[indices]
+    """
+
+    def __init__(self):
+        self._keys = []
+        self._getters = []
+
+    def __len__(self):
+        raise NotImplementedError
+
+    @property
+    def keys(self):
+        return tuple(key for key, _, _ in self._keys)
+
+    def add_getter(self, keys, getter):
+        """Register a getter function
+
+        Args:
+            keys (int or string or tuple of strings): The number or 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 isinstance(keys, int):
+            if keys == 1:
+                keys = None
+            else:
+                keys = (None,) * keys
+        if isinstance(keys, tuple):
+            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 get_example_by_keys(self, index, key_indices):
+        example = []
+        cache = {}
+        for key_index in key_indices:
+            _, getter_index, key_index = self._keys[key_index]
+            if getter_index not in cache:
+                cache[getter_index] = self._getters[getter_index](index)
+            if key_index is None:
+                example.append(cache[getter_index])
+            else:
+                example.append(cache[getter_index][key_index])
+        return tuple(example)

chainercv/chainer_experimental/datasets/sliceable/sliceable_dataset.py

@@ -0,0 +1,137 @@
+import six
+
+import chainer
+
+
+def _as_tuple(t):
+    if isinstance(t, tuple):
+        return t
+    else:
+        return t,
+
+
+class SliceableDataset(chainer.dataset.DatasetMixin):
+    """An abstract dataset class that supports slicing.
+
+    This ia a dataset class that supports slicing.
+    A dataset class inheriting this class should implement
+    three methods: :meth:`__len__`, :meth:`keys`, and
+    :meth:`get_example_by_keys`.
+    """
+
+    def __len__(self):
+        raise NotImplementedError
+
+    @property
+    def keys(self):
+        """Return names of all keys
+
+        Returns:
+            string or tuple of strings
+        """
+        raise NotImplementedError
+
+    def get_example_by_keys(self, index, key_indices):
+        """Return data of an example by keys
+
+        Args:
+            index (int): An index of an example.
+            key_indices (tuple of ints): A tuple of indices of requested keys.
+
+        Returns:
+            tuple of data
+        """
+        raise NotImplementedError
+
+    def get_example(self, index):
+        if isinstance(self.keys, tuple):
+            return self.get_example_by_keys(
+                index, tuple(range(len(self.keys))))
+        else:
+            return self.get_example_by_keys(index, (0,))[0]
+
+    @property
+    def slice(self):
+        return SliceHelper(self)
+
+    def __iter__(self):
+        return (self.get_example(i) for i in six.moves.range(len(self)))
+
+
+class SliceHelper(object):
+    """A helper class for :class:`SliceableDataset`."""
+
+    def __init__(self, dataset):
+        self._dataset = dataset
+
+    def __getitem__(self, args):
+        if isinstance(args, tuple):
+            index, keys = args
+        else:
+            index = args
+            keys = self._dataset.keys
+
+        if isinstance(keys, (list, tuple)):
+            return_tuple = True
+        else:
+            keys, return_tuple = (keys,), False
+
+        # convert name to index
+        key_indices = []
+        for key in keys:
+            if isinstance(key, int):
+                key_index = key
+                if key_index >= len(self._dataset.keys):
+                    raise IndexError('Invalid index of key')
+                if key_index < 0:
+                    key_index += len(self._dataset.keys)
+            else:
+                try:
+                    key_index = _as_tuple(self._dataset.keys).index(key)
+                except ValueError:
+                    raise KeyError('{} does not exists'.format(key))
+            key_indices.append(key_index)
+
+        return SlicedDataset(
+            self._dataset, index,
+            tuple(key_indices) if return_tuple else key_indices[0])
+
+
+class SlicedDataset(SliceableDataset):
+    """A sliced view for :class:`SliceableDataset`."""
+
+    def __init__(self, dataset, index, key_indices):
+        self._dataset = dataset
+        self._index = index
+        self._key_indices = key_indices
+
+    def __len__(self):
+        if isinstance(self._index, slice):
+            start, end, step = self._index.indices(len(self._dataset))
+            return len(range(start, end, step))
+        else:
+            return len(self._index)
+
+    @property
+    def keys(self):
+        keys = _as_tuple(self._dataset.keys)
+        if isinstance(self._key_indices, tuple):
+            return tuple(keys[key_index] for key_index in self._key_indices)
+        else:
+            return keys[self._key_indices]
+
+    def get_example_by_keys(self, index, key_indices):
+        if isinstance(key_indices, tuple):
+            key_indices = tuple(
+                _as_tuple(self._key_indices)[key_index]
+                for key_index in key_indices)
+        else:
+            key_indices = _as_tuple(self._key_indices)[key_indices]
+
+        if isinstance(self._index, slice):
+            start, _, step = self._index.indices(len(self._dataset))
+            return self._dataset.get_example_by_keys(
+                start + index * step, key_indices)
+        else:
+            return self._dataset.get_example_by_keys(
+                self._index[index], key_indices)

chainercv/chainer_experimental/datasets/sliceable/transform_dataset.py

@@ -0,0 +1,40 @@
+from chainercv.chainer_experimental.datasets.sliceable import GetterDataset
+
+
+class TransformDataset(GetterDataset):
+    """A sliceable version of :class:`chainer.datasets.TransformDataset`.
+
+    Note that it reuqires :obj:`keys` to determine the names of returned
+    values.
+
+    Hew is an example.
+
+    >>> def transfrom(in_data):
+    >>>     img, bbox, label = in_data
+    >>>     ...
+    >>>     return new_img, new_label
+    >>>
+    >>> dataset = TramsformDataset(dataset, ('img', 'label'), transform)
+    >>> dataset.keys  # ('img', 'label')
+
+    Args:
+        dataset: The underlying dataset.
+            This dataset should have :meth:`__len__` and :meth:`__getitem__`.
+        keys (int or string or tuple of strings): The number or name(s) of
+            data that the transform function returns.
+        transform (callable): A function that is called to transform values
+            returned by the underlying dataset's :meth:`__getitem__`.
+    """
+
+    def __init__(self, dataset, keys, transform):
+        super(TransformDataset, self).__init__()
+        self._dataset = dataset
+        if isinstance(keys, int):
+            if keys == 1:
+                keys = None
+            else:
+                keys = (None,) * keys
+        self.add_getter(keys, lambda index: transform(dataset[index]))
+
+    def __len__(self):
+        return len(self._dataset)