refactor(shot_detector): update ffmpeg api

gnes-ai · Aug 26, 2019 · e3bbbd9 · e3bbbd9
1 parent 29abf0b
commit e3bbbd9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/gnes/preprocessor/video/ffmpeg.py b/gnes/preprocessor/video/ffmpeg.py
@@ -48,7 +48,7 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
         # video could't be processed from ndarray!
         # only bytes can be passed into ffmpeg pipeline
         if doc.raw_bytes:
-            frames = video_util.capture_frames(video_data=doc.raw_bytes, scale=self.frame_size, fps=self.frame_rate)
+            frames = video_util.capture_frames(input_data=doc.raw_bytes, scale=self.frame_size, fps=self.frame_rate)
             # remove dupliated key frames by phash value
             if self.duplicate_rm:
                 frames = self.duplicate_rm_hash(frames)
@@ -135,7 +135,7 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
             if self.use_image_input:
                 frames = split_video_frames(doc.raw_bytes, self.splitter)
             else:
-                frames = video_util.capture_frames(video_data=doc.raw_bytes, scale=self.frame_size, fps=self.frame_rate)
+                frames = video_util.capture_frames(input_data=doc.raw_bytes, scale=self.frame_size, fps=self.frame_rate)
             if self.max_frames_per_doc > 0:
                 random_id = random.sample(range(len(frames)),
                                           k=min(self.max_frames_per_doc, len(frames)))

diff --git a/gnes/preprocessor/video/shotdetect.py b/gnes/preprocessor/video/shotdetect.py
@@ -71,7 +71,7 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
 
         if doc.raw_bytes:
             all_frames = video_util.capture_frames(
-                video_data=doc.raw_bytes,
+                input_data=doc.raw_bytes,
                 scale=self.frame_size,
                 fps=self.frame_rate)
             num_frames = len(all_frames)