Skip to content
Permalink
Browse files

refactor(video-decoder): update video decoder signature and add more …

…options
  • Loading branch information...
numb3r3 committed Sep 27, 2019
1 parent f6c263a commit 535085ef749345ed4db99c39f5c54134f6df541b
Showing with 35 additions and 26 deletions.
  1. +32 −23 gnes/preprocessor/video/video_decode.py
  2. +3 −3 gnes/preprocessor/video/video_encoder.py
@@ -13,50 +13,59 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from gnes.preprocessor.base import BaseVideoPreprocessor
from gnes.proto import gnes_pb2, array2blob
from gnes.preprocessor.io_utils import video
from ...proto import gnes_pb2, array2blob

from ..base import BaseVideoPreprocessor
from ..io_utils import video as video_util


class VideoDecodePreprocessor(BaseVideoPreprocessor):
store_args_kwargs = True

def __init__(self,
frame_rate: int = 10,
frame_size: str = None,
vframes: int = -1,
scale: str = None,
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.scale = scale
self.drop_raw_data = drop_raw_data
self.chunk_spliter = chunk_spliter

def apply(self, doc: 'gnes_pb2.Document') -> None:
super().apply(doc)

all_frames = []
if doc.WhichOneof('raw_data'):
video_frames = []
raw_type = type(getattr(doc, doc.WhichOneof('raw_data')))
if doc.raw_bytes:
all_frames = video.capture_frames(
video_frames = video_util.capture_frames(
input_data=doc.raw_bytes,
scale=self.scale,
scale=self.frame_size,
fps=self.frame_rate,
vframes=self.vframes)
elif raw_type == gnes_pb2.NdArray:
all_frames = blob2array(doc.raw_video)
if self.vframes > 0:
all_frames = video_frames[0:self.vframes, :].copy()

num_frames = len(all_frames)
if num_frames > 0:
c = doc.chunks.add()
c.doc_id = doc.doc_id
c.blob.CopyFrom(array2blob(all_frames))
c.offset = 0
c.weight = 1.0
if not self.drop_raw_data:
doc.raw_video.CopyFrom(array2blob(video_frames))
else:
self.logger.error('bad document: "raw_bytes" or "raw_video" is empty!')
self.logger.error('the document "raw_bytes" is empty!')

if self.chunk_spliter == 'frame_split':
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':
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_bytes" is empty!')
self.logger.error('bad document: "raw_data" is empty!')
@@ -13,9 +13,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from gnes.preprocessor.base import BaseVideoPreprocessor
from gnes.preprocessor.io_utils import video, gif
from gnes.proto import gnes_pb2, blob2array
from ...proto import gnes_pb2, blob2array
from ..base import BaseVideoPreprocessor
from ..io_utils import video, gif


class VideoEncoderPreprocessor(BaseVideoPreprocessor):

0 comments on commit 535085e

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