Merge pull request #9 from ycxioooong/master

add slice support and unit test

mmcv/video/io.py

@@ -163,8 +163,8 @@ class VideoReader(object):
             ndarray or None: Return the frame if successful, otherwise None.
         """
         if frame_id < 0 or frame_id >= self._frame_cnt:
-            raise ValueError('"frame_id" must be between 0 and {}'.format(
-                self._frame_cnt))
+            raise IndexError('"frame_id" must be between 0 and {}'.format(
+                self._frame_cnt - 1))
         if frame_id == self._position:
             return self.read()
         if self._cache:
@@ -240,7 +240,12 @@ class VideoReader(object):
 
     def __getitem__(self, index):
         if isinstance(index, slice):
-            raise RuntimeError('slice has not been supported yet')
+            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):

tests/test_video.py

@@ -62,12 +62,45 @@ 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(ValueError):
+        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)
+        imgs = v[-105:-103]
+        assert int(round(imgs[0].mean())) == 94
+        assert int(round(imgs[1].mean())) == 205
+        assert len(imgs) == 2
+        imgs = v[63:65]
+        assert int(round(imgs[0].mean())) == 94
+        assert int(round(imgs[1].mean())) == 205
+        assert len(imgs) == 2
+        imgs = v[64:62:-1]
+        assert int(round(imgs[0].mean())) == 205
+        assert int(round(imgs[1].mean())) == 94
+        assert len(imgs) == 2
+        imgs = v[:5]
+        assert len(imgs) == 5
+        for img in imgs:
+            assert int(round(img.mean())) == 94
+        imgs = v[165:]
+        assert len(imgs) == 3
+        for img in imgs:
+            assert int(round(img.mean())) == 0
+        imgs = v[-3:]
+        assert len(imgs) == 3
+        for img in imgs:
+            assert int(round(img.mean())) == 0
 
     def test_current_frame(self):
         v = mmcv.VideoReader(self.video_path)