fix(preprocessor): modify ffmpeg video preprocessor

gnes/preprocessor/helper.py

@@ -186,6 +186,12 @@ def get_video_frames(buffer_data: bytes, image_format: str = 'cv2',
     return frames
 
 
+def split_video_frames(buffer_data: bytes,
+                       splitter: str = '__split__'):
+    chunks = buffer_data.split(splitter.encode())
+    return [np.array(Image.open(io.BytesIO(chunk))) for chunk in chunks]
+
+
 def block_descriptor(image: 'np.ndarray',
                      descriptor_fn: Callable,
                      num_blocks: int = 3) -> 'np.ndarray':

gnes/preprocessor/video/ffmpeg.py

@@ -18,7 +18,7 @@
 import numpy as np
 import random
 from .base import BaseVideoPreprocessor
-from ..helper import get_video_frames, phash_descriptor
+from ..helper import get_video_frames, split_video_frames, phash_descriptor
 from ...proto import gnes_pb2, array2blob
 
 
@@ -108,6 +108,8 @@ def __init__(self,
                  segment_interval: int = -1,
                  segment_num: int = 3,
                  max_frames_per_doc: int = -1,
+                 use_image_input: bool = False,
+                 splitter: str = '__split__',
                  audio_interval: int = 30,
                  sample_rate: int = 16000,
                  *args,
@@ -119,12 +121,17 @@ def __init__(self,
         self.max_frames_per_doc = max_frames_per_doc
         self.audio_interval = audio_interval
         self.sample_rate = sample_rate
+        self.splitter = splitter
+        self.use_image_input = use_image_input
         self._ffmpeg_kwargs = kwargs
 
     def apply(self, doc: 'gnes_pb2.Document') -> None:
         super().apply(doc)
         if doc.raw_bytes:
-            frames = get_video_frames(doc.raw_bytes, **self._ffmpeg_kwargs)
+            if self.use_image_input:
+                frames = split_video_frames(doc.raw_bytes, self.splitter)
+            else:
+                frames = get_video_frames(doc.raw_bytes, **self._ffmpeg_kwargs)
             if self.max_frames_per_doc > 0:
                 random_id = random.sample(range(len(frames)),
                                           k=min(self.max_frames_per_doc, len(frames)))