from ..base import BaseVideoPreprocessor
from ..io_utils import video as video_util
from ...proto import array2blob
[docs]class VideoDecoderPreprocessor(BaseVideoPreprocessor):
store_args_kwargs = True
def __init__(self,
frame_rate: int = 10,
frame_size: str = None,
vframes: int = -1,
drop_raw_data: bool = False,
chunk_spliter: str = None,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.frame_rate = frame_rate
self.frame_size = frame_size
self.vframes = vframes
self.drop_raw_data = drop_raw_data
self.chunk_spliter = chunk_spliter
[docs] def apply(self, doc: 'gnes_pb2.Document') -> None:
super().apply(doc)
if doc.WhichOneof('raw_data'):
video_frames = []
# raw_type = type(getattr(doc, doc.WhichOneof('raw_data')))
if doc.raw_bytes:
video_frames = video_util.capture_frames(
input_data=doc.raw_bytes,
scale=self.frame_size,
fps=self.frame_rate,
vframes=self.vframes)
if not self.drop_raw_data:
doc.raw_video.CopyFrom(array2blob(video_frames))
else:
self.logger.error('the document "raw_bytes" is empty!')
if self.chunk_spliter == 'base':
for i, frame in enumerate(video_frames):
c = doc.chunks.add()
c.doc_id = doc.doc_id
c.blob.CopyFrom(array2blob(frame))
c.offset = i
c.weight = 1.0
elif self.chunk_spliter == 'none':
pass
elif self.chunk_spliter == 'shot':
raise NotImplementedError
else:
chunk = doc.chunks.add()
chunk.doc_id = doc.doc_id
chunk.blob.CopyFrom(array2blob(video_frames))
chunk.offset = 0
chunk.weight = 1.0
else:
self.logger.error('bad document: "raw_data" is empty!')