Skip to content
Permalink
Browse files

feat(preprocessor): add sframes for shots frame number

  • Loading branch information...
raccoonliukai committed Sep 19, 2019
1 parent 968f6c7 commit 6833a27ca0a8dc75949c973ecfa0bee0942b25c4
Showing with 10 additions and 4 deletions.
  1. +10 −4 gnes/preprocessor/video/shotdetect.py
@@ -32,6 +32,7 @@ def __init__(self,
frame_size: str = None,
frame_rate: int = 10,
frame_num: int = -1,
sframes: int = -1,
drop_raw_data: bool = False,
*args,
**kwargs):
@@ -42,6 +43,7 @@ def __init__(self,
self.detect_method = detect_method
self.frame_rate = frame_rate
self.frame_num = frame_num
self.sframes = sframes
self.drop_raw_data = drop_raw_data
self._detector_kwargs = kwargs

@@ -85,19 +87,23 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
elif raw_type == gnes_pb2.NdArray:
video_frames = blob2array(doc.raw_video)
if self.frame_num > 0:
stepwise = len(video_frames) / self.frame_num
video_frames = video_frames[0::stepwise, :]
video_frames = video_frames[0:self.frame_num, :]

num_frames = len(video_frames)
if num_frames > 0:
shots = self.detect_shots(video_frames)
for ci, frames in enumerate(shots):
c = doc.chunks.add()
c.doc_id = doc.doc_id
c.offset = ci
shot_len = len(frames)
c.weight = shot_len / num_frames
if self.sframes > 0 and shot_len > self.sframes:
start_id = (shot_len - self.sframes) / 2
end_id = start_id + self.sframes
frames = frames[start_id:end_id]
chunk_data = np.array(frames)
c.blob.CopyFrom(array2blob(chunk_data))
c.offset = ci
c.weight = len(frames) / num_frames
else:
self.logger.error(
'bad document: "raw_bytes" or "raw_video" is empty!')

0 comments on commit 6833a27

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