feat(preprocessor): add frame selector

gnes/preprocessor/video/frame_select.py

@@ -13,7 +13,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import random
+import numpy as np
+import math
 
 from gnes.preprocessor.base import BaseVideoPreprocessor
 from gnes.proto import gnes_pb2, array2blob, blob2array
@@ -35,36 +36,16 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
                 images = blob2array(chunk.blob)
                 if len(images) == 0:
                     self.logger.info("this chunk has no frame!")
-                if self.sframes == -1:
-                    self.logger.info("keep all frames!")
-                elif self.sframes > len(images):
-                    self.logger.error("sframes should not be larger than number of frames!")
                 elif self.sframes == 1:
                     idx = [int(len(images) / 2)]
                     chunk.blob.CopyFrom(array2blob(images[idx]))
-                else:
-                    idx = []
-                    rs = ReservoirSample(self.sframes)
-                    for item in range(0, len(images)):
-                        idx = rs.feed(item)
-                    chunk.blob.CopyFrom(array2blob(images[idx]))
+                elif self.sframes > 0 and len(images) > self.sframes:
+                    if len(images) >= 2 * self.sframes:
+                        step = math.ceil(len(images) / self.sframes)
+                        chunk.blob.CopyFrom(array2blob(images[::step]))
+                    else:
+                        idx = np.sort(np.random.choice(len(images), self.sframes, replace=False))
+                        chunk.blob.CopyFrom(array2blob(images[idx]))
         else:
             self.logger.error(
                 'bad document: "doc.chunks" is empty!')
-
-
-class ReservoirSample():
-    def __init__(self, size):
-        self._size = size
-        self._counter = 0
-        self._sample = []
-
-    def feed(self, item):
-        self._counter += 1
-        if len(self._sample) < self._size:
-            self._sample.append(item)
-            return self._sample
-        rand_int = random.randint(1, self._counter)
-        if rand_int <= self._size:
-            self._sample[rand_int - 1] = item
-        return self._sample

tests/test_frame_selector.py

@@ -2,6 +2,7 @@
 from gnes.proto import gnes_pb2, array2blob, blob2array
 from gnes.preprocessor.video.frame_select import FrameSelectPreprocessor
 import numpy as np
+import copy
 
 
 class TestFrameSelector(unittest.TestCase):
@@ -22,22 +23,25 @@ def test_emtpy_document(self):
         frame_selector.apply(gnes_pb2.Document())
 
     def test_big_sframe(self):
+        doc = copy.deepcopy(self.doc)
         frame_selector = FrameSelectPreprocessor(sframes=100)
-        frame_selector.apply(self.doc)
+        frame_selector.apply(doc)
 
     def test_get_frames(self):
+        doc = copy.deepcopy(self.doc)
         frame_selector = FrameSelectPreprocessor(sframes=3)
-        frame_selector.apply(self.doc)
-        for idx, chunk in enumerate(self.doc.chunks):
+        frame_selector.apply(doc)
+        for idx, chunk in enumerate(doc.chunks):
             if idx == 0:
                 self.assertEqual(blob2array(chunk.blob).shape[0], 2)
             else:
                 self.assertEqual(blob2array(chunk.blob).shape[0], 3)
 
     def test_get_one_frame(self):
+        doc = copy.deepcopy(self.doc)
         frame_selector = FrameSelectPreprocessor(sframes=1)
-        frame_selector.apply(self.doc)
-        for chunk in self.doc.chunks:
+        frame_selector.apply(doc)
+        for chunk in doc.chunks:
             self.assertEqual(blob2array(chunk.blob).shape[0], 1)