Skip to content
Permalink
Browse files

fix(preprocessor): add random sampling to ffmpeg

  • Loading branch information...
Larryjianfeng committed Aug 7, 2019
1 parent 82ee826 commit 7031fe20e9583720fb2f4b1b930f029e50f55303
Showing with 7 additions and 1 deletion.
  1. +7 −1 gnes/preprocessor/video/ffmpeg.py
@@ -16,7 +16,7 @@
from typing import List

import numpy as np

import random
from .base import BaseVideoPreprocessor
from ..helper import get_video_frames, phash_descriptor
from ...proto import gnes_pb2, array2blob
@@ -107,18 +107,24 @@ def __init__(self,
segment_method: str = 'cut_by_frame',
segment_interval: int = -1,
segment_num: int = 3,
max_frames_per_doc: int = -1,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.segment_method = segment_method
self.segment_interval = segment_interval
self.segment_num = segment_num
self.max_frames_per_doc = max_frames_per_doc
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.max_frames_per_doc > 0:
random_id = random.sample(range(len(frames)),
k=min(self.max_frames_per_doc, len(frames)))
frames = [frames[i] for i in sorted(random_id)]

sub_videos = []
if len(frames) >= 1:

0 comments on commit 7031fe2

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