Skip to content
Permalink
Browse files

fix(preprocessor): add vframe(frame_num) for video and gif frames cap…

…ture
  • Loading branch information...
raccoonliukai committed Sep 10, 2019
1 parent a88fe11 commit a5999828297494a1bb0644718de373bb915eeb14
@@ -25,7 +25,8 @@
def capture_frames(input_fn: str = 'pipe:',
input_data: bytes = None,
fps: int = None,
pix_fmt: str = 'rgb24') -> 'np.ndarray':
pix_fmt: str = 'rgb24',
vframes: int = -1) -> 'np.ndarray':
_check_input(input_fn, input_data)


@@ -41,6 +42,8 @@ def capture_frames(input_fn: str = 'pipe:',
video_filters += ['fps=%d' % fps]

output_kwargs = {'format': 'rawvideo', 'pix_fmt': pix_fmt}
if vframes > 0:
output_kwargs['vframes'] = vframes

cmd_args = compile_args(
input_fn=input_fn,
@@ -139,6 +139,7 @@ def capture_frames(input_fn: str = 'pipe:',
scale: str = None,
start_time: float = None,
end_time: float = None,
vframes: int = -1,
**kwargs) -> List['np.ndarray']:
_check_input(input_fn, input_data)

@@ -199,6 +200,8 @@ def capture_frames(input_fn: str = 'pipe:',
'vcodec': 'rawvideo',
'movflags': 'faststart',
}
if vframes > 0:
output_kwargs['vframes'] = vframes

cmd_args = compile_args(
input_fn=input_fn,
@@ -29,6 +29,7 @@ class FFmpegPreprocessor(BaseVideoPreprocessor):
def __init__(self,
frame_size: str = '192:168',
frame_rate: int = 10,
frame_num: int = -1,
duplicate_rm: bool = True,
use_phash_weight: bool = False,
phash_thresh: int = 5,
@@ -37,6 +38,7 @@ def __init__(self,
super().__init__(*args, **kwargs)
self.frame_size = frame_size
self.frame_rate = frame_rate
self.frame_num = frame_num
self.phash_thresh = phash_thresh
self.duplicate_rm = duplicate_rm
self.use_phash_weight = use_phash_weight
@@ -47,7 +49,8 @@ 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(input_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, vframes=self.frame_num)
# remove dupliated key frames by phash value
if self.duplicate_rm:
frames = self.duplicate_rm_hash(frames)
@@ -109,6 +112,7 @@ class FFmpegVideoSegmentor(BaseVideoPreprocessor):
def __init__(self,
frame_size: str = '192:168',
frame_rate: int = 10,
frame_num: int = -1,
segment_method: str = 'cut_by_frame',
segment_interval: int = -1,
segment_num: int = 3,
@@ -120,6 +124,7 @@ def __init__(self,
super().__init__(*args, **kwargs)
self.frame_size = frame_size
self.frame_rate = frame_rate
self.frame_num = frame_num
self.segment_method = segment_method
self.segment_interval = segment_interval
self.segment_num = segment_num
@@ -134,7 +139,8 @@ 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(input_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, vframes=self.frame_num)
if self.max_frames_per_doc > 0:
random_id = random.sample(range(len(frames)),
k=min(self.max_frames_per_doc, len(frames)))
@@ -32,6 +32,7 @@ def __init__(self,
distance_metric: str = 'bhattacharya',
detect_method: str = 'threshold',
frame_rate: int = 10,
frame_num: int = -1,
*args,
**kwargs):
super().__init__(*args, **kwargs)
@@ -40,6 +41,7 @@ def __init__(self,
self.distance_metric = distance_metric
self.detect_method = detect_method
self.frame_rate = frame_rate
self.frame_num = frame_num
self._detector_kwargs = kwargs

def detect_shots(self, frames: 'np.ndarray') -> List[List['np.ndarray']]:
@@ -73,7 +75,8 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
all_frames = video_util.capture_frames(
input_data=doc.raw_bytes,
scale=self.frame_size,
fps=self.frame_rate)
fps=self.frame_rate,
vframes=self.frame_num)
num_frames = len(all_frames)
assert num_frames > 0
shots = self.detect_shots(all_frames)
@@ -32,6 +32,8 @@ def test_get_media_meta(self):

def test_capture_frames(self):
frames1 = video.capture_frames(input_fn=self.video_path, fps=10, scale='640:-2')
sub_frames = video.capture_frames(input_fn=self.video_path, fps=10, scale='640:-2', vframes=5)
self.assertEqual((sub_frames == frames1[:5]).all(), True)

with open(self.video_path, 'rb') as f:
data = f.read()
@@ -54,7 +56,9 @@ def test_encode_video(self):
def test_gif_encode(self):
gif_data = gif.encode_video(images=self.frames, frame_rate=10)
frames = gif.capture_frames(input_data=gif_data)
sub_frames = gif.capture_frames(input_data=gif_data, vframes=5)
self.assertEqual(self.frames.shape, frames.shape)
self.assertEqual((sub_frames == frames[:5]).all(), True)

def test_capture_audio(self):
audio_data1 = audio.capture_audio(input_fn=self.video_path)

0 comments on commit a599982

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