Skip to content
Permalink
Browse files

fix(preprocessor): modify ffmpeg video preprocessor

  • Loading branch information...
Larryjianfeng committed Aug 9, 2019
1 parent 877a313 commit d1cfa539ed0f60b0dd1d4c09dab9ef9b4a626ef1
Showing with 15 additions and 2 deletions.
  1. +6 −0 gnes/preprocessor/helper.py
  2. +9 −2 gnes/preprocessor/video/ffmpeg.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':
@@ -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)))

0 comments on commit d1cfa53

Please sign in to comment.
You can’t perform that action at this time.