Skip to content

Instantly share code, notes, and snippets.

@willprice
Last active October 20, 2020 08:19
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save willprice/83d91ad1b79d78db977e25c3486504ab to your computer and use it in GitHub Desktop.
Save willprice/83d91ad1b79d78db977e25c3486504ab to your computer and use it in GitHub Desktop.
Convert video array to mp4 blob
from io import BytesIO
import av
import numpy as np
def video_array_to_mp4_blob(video: np.ndarray, fps: float = 24) -> bytes:
"""
Args:
video: Video array of shape :math:`(T, H, W, 3)` and range :math:`[0, 255]` with type ``np.uint8``.
fps: FPS of video
Returns:
bytes containing the encoded video as MP4/h264.
Examples:
>>> from IPython.display import HTML, display
>>> from base64 import b64encode
>>> src = 'data:video/mp4;base64,' + b64encode(video_array_to_mp4_blob(video)).decode('utf8')
>>> display(HTML(f'<video width="{video.shape[-2]}" height="{video.shape[-3]}" controls><source src="{src}"></video>'))
"""
f = BytesIO()
container = av.open(f, mode='w', format='mp4')
stream = container.add_stream('h264', rate=fps)
stream.width = video.shape[-2]
stream.height = video.shape[-3]
stream.pix_fmt = 'yuv420p'
for frame in video:
av_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
for packet in stream.encode(av_frame):
container.mux(packet)
for packet in stream.encode():
container.mux(packet)
container.close()
return f.getvalue()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment