Commit 6249f829 authored by xiongyu's avatar xiongyu
Browse files

add negative indexing support

parent c75ff259
......@@ -241,6 +241,11 @@ class VideoReader(object):
def __getitem__(self, index):
if isinstance(index, slice):
return [self.get_frame(i) for i in range(*index.indices(self.frame_cnt))]
# support negative indexing
if index < 0:
index += self.frame_cnt
if index < 0:
raise IndexError('index out of range')
return self.get_frame(index)
def __iter__(self):
......
......@@ -62,12 +62,18 @@ class TestVideo(object):
assert int(round(img.mean())) == 94
img = v[64]
assert int(round(img.mean())) == 205
img = v[-104]
assert int(round(img.mean())) == 205
img = v[63]
assert int(round(img.mean())) == 94
img = v[-105]
assert int(round(img.mean())) == 94
img = v.read()
assert int(round(img.mean())) == 205
with pytest.raises(IndexError):
v.get_frame(self.num_frames + 1)
with pytest.raises(IndexError):
v[-self.num_frames - 1]
def test_slice(self):
v = mmcv.VideoReader(self.video_path)
......
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