Source code for jina.types.document.mixins.video

from typing import Union, BinaryIO

import numpy as np

from .helper import _deprecate_by
from ....helper import T


[docs]class VideoDataMixin: """Provide helper functions for :class:`Document` to support video data. """
[docs] def load_uri_to_video_blob(self: T, only_keyframes: bool = False) -> T: """Convert a :attr:`.uri` to a video ndarray :attr:`.blob`. :param only_keyframes: only keep the keyframes in the video :return: Document itself after processed """ import av with av.open(self.uri) as container: if only_keyframes: stream = container.streams.video[0] stream.codec_context.skip_frame = 'NONKEY' frames = [] for frame in container.decode(video=0): img = frame.to_image() frames.append(np.asarray(img)) self.blob = np.moveaxis(np.stack(frames), 1, 2) return self
[docs] def dump_video_blob_to_file( self: T, file: Union[str, BinaryIO], frame_rate: int = 30, codec: str = 'h264' ) -> T: """Save :attr:`.blob` as a video mp4/h264 file. :param file: The file to open, which can be either a string or a file-like object. :param frame_rate: frames per second :param codec: the name of a decoder/encoder :return: itself after processed """ if ( self.blob.ndim != 4 or self.blob.shape[-1] != 3 or self.blob.dtype != np.uint8 ): raise ValueError( f'expects `.blob` with dtype=uint8 and ndim=4 and the last dimension is 3, ' f'but receiving {self.blob.shape} in {self.blob.dtype}' ) video_blob = np.moveaxis(np.clip(self.blob, 0, 255), 1, 2) import av with av.open(file, mode='w') as container: stream = container.add_stream(codec, rate=frame_rate) stream.width = self.blob.shape[1] stream.height = self.blob.shape[2] stream.pix_fmt = 'yuv420p' for b in video_blob: frame = av.VideoFrame.from_ndarray(b, format='rgb24') for packet in stream.encode(frame): container.mux(packet) for packet in stream.encode(): container.mux(packet) return self
convert_uri_to_video_blob = _deprecate_by(load_uri_to_video_blob)