Commit 9124c770 authored by David GEISSBUHLER's avatar David GEISSBUHLER
Browse files

self.get_frame_indices -> utils.get_index_list

parent 040ab89b
......@@ -5,7 +5,7 @@ from functools import wraps
import numpy as np
from builtins import str
from .utils import StreamArray
from .utils import StreamArray, get_index_list
from .stream_file import StreamFile
......@@ -149,55 +149,9 @@ class Stream:
def __setitem__(self, index, data):
raise Exception("not yet implemented")
# get list of frame indices
def get_frame_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):
indices = [index]
# slice transform to list
elif isinstance(index, slice):
# start value: handle None and negative
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:
step = 1
# generate list
indices = list(range(start, stop, step))
# pass lists thru
elif isinstance(index, list):
indices = index
else:
raise Exception("index can only be None, int, slice or list")
return indices
# load one or several frames
def load(self, index=None):
indices = self.get_frame_indices(index)
indices = get_index_list(index, self.shape[0])
# return buffered data OR load from file OR process data
if self._loaded == indices and self._data is not None:
# print('loaded', self.name)
......@@ -250,9 +204,6 @@ class StreamFilter(Stream):
super().__init__(name=name, parent=parent)
pass
def get_frame_indices(self, index):
return super().get_frame_indices(index)
def process(self, data, indices):
assert isinstance(indices, list)
return np.stack([self.process_frame(data[i], i, indices[i]) for i in range(data.shape[0])])
......@@ -262,7 +213,7 @@ class StreamFilter(Stream):
# load one or several frames
def load(self, index=None):
indices = self.get_frame_indices(index)
indices = get_index_list(index, self.shape[0])
# return buffered data OR process data
if self._loaded == indices and self._data is not None:
# print('loaded', self.name)
......@@ -368,12 +319,8 @@ class StreamView(StreamFilter):
__ndim -= 1
return __ndim
# returns frame indices to be loaded from parent
def get_frame_indices(self, index):
return super().get_frame_indices(index)
def load(self, index=None):
indices = self.get_frame_indices(index)
indices = get_index_list(index, self.shape[0])
return super().load(indices)
def process(self, data, indices):
......@@ -411,16 +358,13 @@ class StreamAdjust(StreamFilter):
def timestamps(self):
return self.adjust_to.timestamps
def get_frame_indices(self, index):
return super().get_frame_indices(index)
def load(self, index):
# TODO load only relevant data if cropped
if isinstance(index, tuple):
index = index[0]
print("WARNING: cropping not yet implemented")
# original stream indices
old_indices = self.get_frame_indices(index)
old_indices = get_index_list(index, self.shape[0])
selected_timestamps = [self.adjust_to.timestamps[i] for i in old_indices]
kdtree = cKDTree(self.parent.timestamps[:, np.newaxis])
......
......@@ -41,6 +41,53 @@ def _crop(video, mask, crop):
mask = mask[:, starty:starty + height, startx:startx + width]
return video, mask
################################################################################
# return a list of indices given an index (int / slice / list / None) and a size
# TODO: check all boundary cases
def get_index_list(index, size):
# 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):
indices = [index]
# slice transform to list
elif isinstance(index, slice):
# start value: handle None and negative
if index.start is not None:
if index.start < 0:
start = size + index.start
else:
start = index.start
# boundary case
if start >= size:
start = size - 1
else:
start = 0
# stop value: handle None and negative
if index.stop is not None:
if index.stop < 0:
stop = size + index.stop
else:
stop = index.stop
# boundary case
if stop >= size:
stop = size - 1
else:
stop = size
# step value: handle None
if index.step is not None:
step = index.step
else:
step = 1
# generate list
indices = list(range(start, stop, step))
# pass lists thru
elif isinstance(index, list):
indices = index
else:
raise Exception("index can only be None, int, slice or list")
return indices
################################################################################
class StreamArray():
......
Markdown is supported
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