Skip to content
Permalink
Browse files

fix(encoder): fix mixture

  • Loading branch information...
Larryjianfeng committed Jul 31, 2019
1 parent 6799153 commit 3fdf1c06e302a5c3e32d28de431e076802ff5c9e
Showing with 52 additions and 1 deletion.
  1. +3 −1 gnes/encoder/__init__.py
  2. +6 −0 gnes/encoder/base.py
  3. +1 −0 gnes/preprocessor/__init__.py
  4. +42 −0 gnes/preprocessor/video/ffmpeg.py
@@ -34,13 +34,15 @@
'BaseEncoder': 'base',
'BaseBinaryEncoder': 'base',
'BaseTextEncoder': 'base',
'BaseVideoEncoder': 'base',
'BaseNumericEncoder': 'base',
'CompositionalEncoder': 'base',
'PipelineEncoder': 'base',
'HashEncoder': 'numeric.hash',
'BasePytorchEncoder': 'image.base',
'TFInceptionEncoder': 'image.inception',
'CVAEEncoder': 'image.cvae'
'CVAEEncoder': 'image.cvae',
'IncepMixtureEncoder': 'video.incep_mixture'
}

register_all_class(_cls2file_map, 'encoder')
@@ -38,6 +38,12 @@ def encode(self, img: List['np.ndarray'], *args, **kwargs) -> np.ndarray:
pass


class BaseVideoEncoder(BaseEncoder):

def encode(self, img: List['np.ndarray'], *args, **kwargs) -> np.ndarray:
pass


class BaseTextEncoder(BaseEncoder):

def encode(self, text: List[str], *args, **kwargs) -> np.ndarray:
@@ -30,6 +30,7 @@
'BaseUnaryPreprocessor': 'base',
'BaseVideoPreprocessor': 'video.base',
'FFmpegPreprocessor': 'video.ffmpeg',
'FFmpegVideoSegmentor': 'video.ffmpeg',
'ShotDetectPreprocessor': 'video.shotdetect',
}

@@ -106,3 +106,45 @@ def duplicate_rm_hash(self,
ret.append((i, h))

return [images[_[0]] for _ in ret]


class FFmpegVideoSegmentor(BaseVideoPreprocessor):
def __init__(self,
frame_size: str = "192*168",
segment_method: str = 'uniform',
segment_interval: int = -1,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.frame_size = frame_size
self.segment_method = segment_method
self.segment_interval = segment_interval
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,
s=self.frame_size,
vsync=self._ffmpeg_kwargs.get("vsync", "vfr"),
vf=self._ffmpeg_kwargs.get("vf", "select=eq(pict_type\\,I)"))

sub_videos = []
if len(frames) >= 1:
if self.segment_method == 'uniform':
if self.segment_interval == -1:
sub_videos = [frames]
else:
sub_videos = [frames[_: _+self.segment_interval]
for _ in range(0, len(frames), self.segment_interval)]
for ci, chunk in enumerate(sub_videos):
c = doc.chunks.add()
c.doc_id = doc.doc_id
c.blob.CopyFrom(array2blob(np.array(chunk, dtype=np.uint8)))
c.offset_1d = ci
c.weight = 1 / len(sub_videos)
else:
self.logger.info('bad document: no key frames extracted')
else:
self.logger.error('bad document: "raw_bytes" is empty!')

0 comments on commit 3fdf1c0

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