added ndim property and preliminaries for stream_slice (not yet fully working)

api_test.py

@@ -42,6 +42,7 @@ assert(stream_a.camera == None)
 assert(stream_b.camera == None)
 
 # load full datasets
+
 ld_a = stream_a.load()
 ld_b = stream_b.load()
 
@@ -67,6 +68,22 @@ for t in tests:
     sd = stream_a.load(s)
     assert(np.array_equal(sd, dd))
 
+# test loading one frame with [] operator
+
+data = stream_a[1]
+
+assert(isinstance(data, np.ndarray))
+assert(data.shape == data_a[1].shape)
+assert(np.array_equal(data, data_a[1]))
+
+# test slicing in first dimension
+
+stream_sl = stream_a[1:4]
+data_sl = data_a[1:4]
+assert(stream_sl.shape == data_sl.shape)
+assert(stream_sl.ndim == data_sl.ndim)
+#assert(np.array_equal(stream_sl.load(), data_sl))
+
 ###########
 
 #fo = StreamFile('api_test.h5', 'w')

bob/io/stream/stream.py

@@ -66,6 +66,11 @@ class Stream:
     def shape(self, value):
         raise Exception("not yet implemented")
 
+    # ndim
+    @property
+    def ndim(self):
+        return len(self.shape)
+
     # timestamps
     @property
     def timestamps(self):
@@ -107,22 +112,35 @@ class Stream:
             return self.parent.image_points
 
     # data access
-    # reformat data for compatibility with python /np arrays
+    # returns either a frame (if index is int) or a StreamSlice (if index is slice or tuple)
+    # provides modest compatibility with python / numpy arrays
     def __getitem__(self, index):
-        # keep only first (time) index
-        if isinstance(index, tuple):
-            index = index[0]
-            raise Exception("ERROR: cropping not yet implemented")
-            # TODO
-        data = self.load(index)
-        # emulate python/np arrays
+        # case 1: index is int -> load and return the frame data
         if isinstance(index, int):
+            data = self.load(index)
             return data[0]
+        # case 2: index is slice -> pack as tuple
         elif isinstance(index, slice):
-            return data
+            slice_indices = (index, )
+        # case 3: index is tuple -> validate
+        elif isinstance(index, tuple):
+            # we do not check dimensionality here, 
+            # it can be varying if the pipeline has no source
+            for i in index:
+                if isinstance(i, int):
+                    pass
+                elif isinstance(i, slice):
+                    pass
+                else:
+                    raise Exception("index can only be int or slice")
+            # so far so good ...
+            slice_indices = index
         else:
-            raise Exception("index can only be int or slice")
-            return None
+            raise Exception("index can only be int, slice or tuple")
+        # case 2 & 3: return stream_slice
+        print('__getitem__', slice_indices)
+        return self.stream_slice(slice_indices=slice_indices)
+
 
     def __setitem__(self, index, data):
         raise Exception("not yet implemented")
@@ -136,7 +154,6 @@ class Stream:
             pass
         else:
             self.__data = self.parent.load_stream_data(self.name, indices)
-
         # buffer and return data
         self.__loaded = indices
         return self.__data
@@ -187,6 +204,19 @@ class Stream:
             raise Exception("index can only be None, int, slice or list")
         return indices
 
+    # get size of a given (sliced) axis 
+    def get_axis_size(self, axis, _slice=None):
+        if _slice is None:
+            return self.shape[axis]
+        else:
+            if isinstance(_slice, int):
+                return 1
+            elif isinstance(_slice, slice):
+                # TODO check this works always
+                return len(range(*_slice.indices(self.shape[axis])))
+            else:
+                raise Exception('_slice can be None, int or slice')
+
     # filters
     def get_available_filters(self):
         return Stream.filters
@@ -202,18 +232,16 @@ def stream_filter(name):
     def wrapper(Filter):
         def filter(self, *args, **kwargs):
             return Filter(name=name, parent=self, *args, **kwargs)
-
         setattr(Stream, name, filter)
         Stream.filters.append(name)
         return Filter
-
     return wrapper
 
 ################################################################################
 # Stream Filters
 
-### default ###
-@stream_filter("nop")
+### default filter ###
+@stream_filter("filter")
 class StreamFilter(Stream):
     def __init__(self, name, parent):
         super().__init__(name=name, parent=parent)
@@ -230,7 +258,7 @@ class StreamFilter(Stream):
         return data
 
     # load one or several frames
-    def load(self, index):
+    def load(self, index=None):
         indices = self.get_frame_indices(index)
         # return buffered data OR load from file OR process data
         if self._Stream__loaded == indices and self._Stream__data is not None:
@@ -242,3 +270,108 @@ class StreamFilter(Stream):
         # buffer and return data
         self.__loaded = indices
         return self.__data
+
+################################################################################
+### astype ###
+# TODO make that work more as numpy...
+@stream_filter("astype")
+class AsType(StreamFilter):
+    def __init__(self, name, parent, dtype=None):
+        super().__init__(name=name, parent=parent)
+        if dtype is not None:
+            self.dtype = dtype
+        else:
+            raise Exception("dtype parameter not set")
+
+    def process(self, data, indices):
+        return data.astype(self.dtype)
+
+################################################################################
+### channel select ###
+# TODO maybe this is no needed, can be done with stream_slice
+@stream_filter("select")
+class Select(StreamFilter):
+    def __init__(self, name, parent, channel=None):
+        super().__init__(name=name, parent=parent)
+        if channel is not None:
+            self.channel = channel
+        else:
+            raise Exception("channel parameter not set")
+
+    @property
+    def shape(self):
+        return (self.parent.shape[0], 1, self.parent.shape[2], self.parent.shape[3])
+
+    def process(self, data, indices):
+        return np.expand_dims(data[:, self.channel, :, :], axis=1)
+
+
+################################################################################
+### stream_slice ###
+@stream_filter("stream_slice")
+class StreamSlice(StreamFilter):
+    def __init__(self, name, parent, slice_indices=None):
+        super().__init__(name=name, parent=parent)
+        self.frame_slice = None
+        self.bulk_slice = None
+        if isinstance(slice_indices, tuple):
+            # separate frame and bulk slices
+            self.frame_slice = slice_indices[0]
+            # TODO should case with int on frame index be allowed? 
+            assert(self.frame_slice is None or isinstance(self.frame_slice, slice))
+            # bulk slices
+            if len(slice_indices) > 1:
+                self.bulk_slice = slice_indices[1:]
+                for e in self.bulk_slice:
+                    if isinstance(e, int) or isinstance(e, slice) or e is None:
+                        pass
+                    else:
+                        raise Exception('slice_indices should be a tuple of int / slice object(s) or None')
+        # None is passthru
+        elif slice_indices is None:
+            pass
+        else:
+            raise Exception('slice_indices should be a tuple of int / slice object(s) or None')
+
+    @property
+    def shape(self):
+        if self.bulk_slice is not None and self.parent.ndim <= len(self.bulk_slice):
+            raise Exception('slice dimension exceed parent dimension')
+        # first dimension ...
+        __shape = [self.parent.get_axis_size(0, self.frame_slice)]
+        # ... and others 
+        for d in range(1, self.parent.ndim):
+            if self.bulk_slice is not None and d+1 < len(self.bulk_slice):
+                __slice_index = self.bulk_slice[d+1]
+                __shape.append(self.parent.get_axis_size(d, __slice_index))
+            else:
+                __shape.append(self.parent.get_axis_size(d))
+        return tuple(__shape)
+
+    @property
+    def ndim(self):
+        __ndim = self.parent.ndim
+        if self.frame_slice is not None and isinstance(self.frame_slice, int):
+            __ndim -= 1
+        if self.bulk_slice is not None:
+            for a in self.bulk_slice:
+                if isinstance(a, int):
+                    __ndim -= 1
+        return __ndim
+
+    # translate frame indices
+    def get_frame_indices(self, index):
+        if self.frame_slice is None:
+            return super().get_frame_indices(index)
+        elif isinstance(self.frame_slice, slice):
+            # TODO perform translation....
+            return super().get_frame_indices(index)
+
+    def load(self, index=None):
+        indices = self.get_frame_indices(index)
+        print('slice load', indices)
+        return super().load(indices)
+
+    def process(self, data, indices):
+        print('process inidices', indices)
+        return data
\ No newline at end of file

bob/io/stream/stream_filters.py

@@ -11,38 +11,6 @@ from bob.ip.stereo import stereo_match, reproject_image, CameraPair
 
 from .stream import stream_filter, Stream, StreamFilter
 
-### channel ###
-@stream_filter("select")
-class Select(StreamFilter):
-    def __init__(self, name, parent, channel=None):
-        super().__init__(name=name, parent=parent)
-        if channel is not None:
-            self.channel = channel
-        else:
-            raise Exception("channel parameter not set")
-
-    @property
-    def shape(self):
-        return (self.parent.shape[0], 1, self.parent.shape[2], self.parent.shape[3])
-
-    def process(self, data, indices):
-        return np.expand_dims(data[:, self.channel, :, :], axis=1)
-
-
-### astype ###
-@stream_filter("astype")
-class AsType(StreamFilter):
-    def __init__(self, name, parent, dtype=None):
-        super().__init__(name=name, parent=parent)
-        if dtype is not None:
-            self.dtype = dtype
-        else:
-            raise Exception("dtype parameter not set")
-
-    def process(self, data, indices):
-        return data.astype(self.dtype)
-
-
 ### to_rgb ###
 @stream_filter("colormap")
 class StreamColorMap(StreamFilter):