def test_capture_frames(self):
frames1 = video.capture_frames(input_fn=self.video_path, fps=10, scale='640:-2')
with open(self.video_path, 'rb') as f:
data = f.read()
frames2 = video.capture_frames(input_data=data, fps=10, scale='-1:360')
self.assertEqual(frames1.shape, frames2.shape)
def apply(self, doc: 'gnes_pb2.Document') -> None:
super().apply(doc)
video_frames = []
if doc.WhichOneof('raw_data'):
raw_type = type(getattr(doc, doc.WhichOneof('raw_data')))
if doc.raw_bytes:
video_frames = video.capture_frames(input_data=doc.raw_bytes,
scale=self.frame_size,
fps=self.frame_rate,
vframes=self.vframes)
elif raw_type == gnes_pb2.NdArray:
video_frames = blob2array(doc.raw_video)
if self.vframes > 0:
video_frames = video_frames[0:self.vframes, :].copy()
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:
if shot_len >= 2 * self.sframes:
step = math.ceil(shot_len / self.sframes)
frames = frames[::step]
else:
idx = np.sort(
np.random.choice(shot_len,
self.sframes,
replace=False))
frames = [frames[idx_] for idx_ in idx]
chunk_data = np.array(frames)
c.blob.CopyFrom(array2blob(chunk_data))
else:
self.logger.error(
'bad document: "raw_bytes" or "raw_video" is empty!')
else:
self.logger.error('bad document: "raw_data" is empty!')
if self.drop_raw_data:
self.logger.info("document raw data will be cleaned!")
doc.ClearField('raw_data')
def apply(self, doc: 'gnes_pb2.Document') -> None:
super().apply(doc)
video_frames = []
if doc.WhichOneof('raw_data'):
raw_type = type(getattr(doc, doc.WhichOneof('raw_data')))
if doc.raw_bytes:
video_frames = video.capture_frames(input_data=doc.raw_bytes,
scale=self.frame_size,
fps=self.frame_rate,
vframes=self.frame_num)
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, :]
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
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!')
else:
self.logger.error('bad document: "raw_data" is empty!')
if self.drop_raw_data:
self.logger.info("document raw data will be cleaned!")
doc.ClearField('raw_data')
def setUp(self):
self.dirname = os.path.dirname(__file__)
self.video_path = os.path.join(self.dirname, 'videos', 'test.mp4')
self.frames = video.capture_frames(input_fn=self.video_path, fps=10, scale='640:360')