Commit 30dc736b authored by David GEISSBUHLER's avatar David GEISSBUHLER
Browse files

handle slice properly

parent 6df3b1c3
Pipeline #42646 passed with stage
in 5 minutes and 26 seconds
...@@ -2,6 +2,7 @@ from bob.io.stream import Stream, StreamFile, StreamFilter, stream_filter #, Str ...@@ -2,6 +2,7 @@ from bob.io.stream import Stream, StreamFile, StreamFilter, stream_filter #, Str
import numpy as np import numpy as np
import bob.io.base import bob.io.base
# create data sets # create data sets
num_frame = 10 num_frame = 10
...@@ -40,7 +41,31 @@ assert(stream_b.timestamps == None) ...@@ -40,7 +41,31 @@ assert(stream_b.timestamps == None)
assert(stream_a.camera == None) assert(stream_a.camera == None)
assert(stream_b.camera == None) assert(stream_b.camera == None)
print(stream_a.image_points) # load full datasets
ld_a = stream_a.load()
ld_b = stream_b.load()
assert(ld_a.shape == data_a.shape)
assert(ld_b.shape == data_b.shape)
assert(np.array_equal(ld_a, data_a))
assert(np.array_equal(ld_b, data_b))
# try some indices slices cases and compare data
tests = [ ( None, None, None),
( None, 3, None),
( 5, None, None),
( None, None, 3),
( 1, 10, 3),
( 9, 0, -3),
( -5, -1, None),
( 10, 0, -3)]
for t in tests:
s = slice(t[0], t[1], t[2])
dd = data_a[s]
sd = stream_a.load(s)
assert(np.array_equal(sd, dd))
########### ###########
......
...@@ -128,7 +128,7 @@ class Stream: ...@@ -128,7 +128,7 @@ class Stream:
raise Exception("not yet implemented") raise Exception("not yet implemented")
# load one or several frames # load one or several frames
def load(self, index): def load(self, index=None):
indices = self.get_indices(index) indices = self.get_indices(index)
# return buffered data OR load from file OR process data # return buffered data OR load from file OR process data
if self.__loaded == indices and self.__data is not None: if self.__loaded == indices and self.__data is not None:
...@@ -141,20 +141,50 @@ class Stream: ...@@ -141,20 +141,50 @@ class Stream:
self.__loaded = indices self.__loaded = indices
return self.__data return self.__data
# get list of indices # get list of frame indices
# TODO check all cases
def get_indices(self, index): def get_indices(self, index):
# None index is equivalent to [:] i.e. slice(None, None, None)
if index is None:
index = slice(None, None, None)
# frame index transform to list
if isinstance(index, int): if isinstance(index, int):
indices = [index] indices = [index]
# slice transform to list
elif isinstance(index, slice): elif isinstance(index, slice):
if index.step == None: # start value: handle None and negative
indices = list(range(index.start, index.stop)) if index.start is not None:
if index.start < 0:
start = self.shape[0] + index.start
else:
start = index.start
# boundary case
if start >= self.shape[0]:
start = self.shape[0] - 1
else:
start = 0
# stop value: handle None and negative
if index.stop is not None:
if index.stop < 0:
stop = self.shape[0] + index.stop
else:
stop = index.stop
# boundary case
if stop >= self.shape[0]:
stop = self.shape[0] - 1
else:
stop = self.shape[0]
# step value: handle None
if index.step is not None:
step = index.step
else: else:
indices = list(range(index.start, index.stop, index.step)) step = 1
# generate list
indices = list(range(start, stop, step))
# pass lists thru
elif isinstance(index, list): elif isinstance(index, list):
indices = index indices = index
else: else:
raise Exception("index can only be int, slice, tuple or list") raise Exception("index can only be None, int, slice or list")
return indices return indices
# filters # filters
......
...@@ -51,7 +51,7 @@ class StreamFile: ...@@ -51,7 +51,7 @@ class StreamFile:
else: else:
# return a generic config if no config is present # return a generic config if no config is present
# TODO: make formal # TODO: make formal
data_config = { 'array_format' : None, data_config = { 'array_format' : {},
'rotation' : None, 'rotation' : None,
'camera' : None, 'camera' : None,
'path' : stream_name} 'path' : stream_name}
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment