import os, base64, pathlib, tempfile, warnings, webbrowser
from .util import gridify, import_torch, _create_folder, _to_tensor, _to_float, _verify_chw
import imageio.v3 as iio
import numpy as np
torch = import_torch()
# Embedding much more than this bloats notebooks, matplotlib uses the same limit
_EMBED_LIMIT = 20 * 1024 * 1024
[docs]
@torch.no_grad()
def show_video(
tensor: torch.Tensor | np.ndarray,
fps: int = 30,
columns: int | None = None,
max_width: int | None = None,
padding: int = 3,
pad_value: float = 0.0,
quality: float = 8.0,
) -> None:
"""
Shows tensor as a video. Accepts both **(T,H,W)**, **(T,3,H,W)** and **(\\*,T,3,H,W)** tensors.
Tensor should be a float tensor with values in [0,1]. Clips the values otherwise.
In a notebook, the video is displayed inline. Otherwise, it is opened in the
default web browser. Both use the same html5 player, with play/pause and seeking.
Args:
tensor : (T,H,W) or (T,3,H,W) or (\\*,T,3,H,W) float tensor or numpy array
fps : fps of the video (default 30)
columns : number of columns to use for the grid of videos (default 8 or less)
max_width : maximum width of the video
padding : number of pixels between videos in the grid
pad_value : inter-padding value for the grid of videos
quality : encoding quality between 1 and 10, 10 being lossless (default 8)
"""
tensor = _format_video(
tensor, columns=columns, fps=fps, max_width=max_width, padding=padding, pad_value=pad_value
) # (T,3,H',W') ready to show
_display_video(_encode_video(tensor, fps=fps, quality=quality))
[docs]
@torch.no_grad()
def save_video(
tensor: torch.Tensor | np.ndarray,
folder: str,
name: str = "videotensor",
fps: int = 30,
columns: int | None = None,
max_width: int | None = None,
padding: int = 3,
pad_value: float = 0.0,
create_folder: bool = True,
quality: float = 8.0,
) -> None:
"""
Shows tensor as a video. Accepts both **(T,H,W)**, **(T,3,H,W)** and **(\\*,T,3,H,W)** tensors.
Tensor should be a float tensor with values in [0,1]. Clips the values otherwise.
Saved as H.264 (libx264) in an mp4 container.
Args:
tensor : (T,H,W) or (T,3,H,W) or (\\*,T,3,H,W) float tensor or numpy array
folder : path to save the video
name : name of the video
fps : fps of the video (default 30)
columns : number of columns to use for the grid of videos (default 8 or less)
max_width : maximum width of the image
padding : number of pixels between images in the grid
pad_value : inter-padding value for the grid of images
create_folder : whether to create the folder if it does not exist (default True)
quality : encoding quality between 1 and 10, 10 being lossless (default 8)
"""
_create_folder(folder, create_folder)
tensor = _format_video(
tensor, columns=columns, fps=fps, max_width=max_width, padding=padding, pad_value=pad_value
) # (T,3,H',W') ready to save
_encode_video(tensor, fps=fps, quality=quality, uri=os.path.join(folder, f"{name}.mp4"))
def _format_video(tensor, columns=None, fps=30, max_width=None, padding=3, pad_value=0.0):
"""
Formats tensor as a video. Accepts both (T,H,W), (T,3,H,W) and (\\*,T,3,H,W) float tensors.
Assumes that the tensor value are in [0,1], clips them otherwise.
Args:
tensor : (T,H,W) or (T,3,H,W) or (\\*,T,3,H,W) float tensor or numpy array
columns : number of columns to use for the grid of videos (default 8 or less)
fps : fps of the video (default 30)
max_width : maximum width of the image
padding : number of pixels between images in the grid
pad_value : inter-padding value for the grid of images
"""
tensor = _to_tensor(tensor).detach().cpu()
tensor = _to_float(tensor)
if len(tensor.shape) >= 4:
# (T,H,W) videos have no channel dimension to place, anything bigger ends with (C,H,W)
tensor = _verify_chw(tensor)
extra_params = dict(columns=columns, fps=fps, max_width=max_width, pad_value=pad_value, padding=padding)
if len(tensor.shape) == 3:
# add channel dimension
tensor = tensor[:, None, :, :].expand(-1, 3, -1, -1) # (T,3,H,W)
return _format_video(tensor, **extra_params)
elif len(tensor.shape) == 4:
if tensor.shape[1] == 1:
print("Assuming gray-scale video")
tensor = tensor.expand(-1, 3, -1, -1) # (T,3,H,W)
assert tensor.shape[1] == 3, f"Tensor shape should be (T,H,W), (T,3,H,W) or (*,T,3,H,W) !"
# A single video
return tensor
elif len(tensor.shape) == 5:
tensor = gridify(tensor, max_width=max_width, columns=columns, padding=padding, pad_value=pad_value)
return _format_video(tensor, **extra_params)
elif len(tensor.shape) > 5:
tensor = tensor.reshape((-1, *tensor.shape[-4:]))
return _format_video(tensor, **extra_params)
else:
raise ValueError(
f"Tensor shape should be (T,H,W), (T,3,H,W) or (*,T,3,H,W), but got : {tensor.shape} !"
)
@torch.no_grad()
def _encode_video(tensor: torch.Tensor, fps: int, quality: float, uri: str = "<bytes>") -> bytes | None:
"""
Encodes a video as H.264 (libx264) in an mp4 container.
Args:
tensor : (T,3,H,W) float tensor, ready to encode
fps : fps of the video
quality : encoding quality between 1 and 10, 10 being lossless
uri : path to write to. Defaults to encoding in memory
Returns:
the mp4 bytes if uri is "<bytes>", None otherwise
"""
# libx264 cannot encode odd dimensions in yuv420p, and ffmpeg fails silently if we let it try.
H, W = tensor.shape[-2], tensor.shape[-1]
if H % 2 or W % 2:
tensor = torch.nn.functional.pad(tensor, (0, W % 2, 0, H % 2), mode="replicate")
frames = tensor.permute(0, 2, 3, 1) # (T,H',W',3)
frames = frames.clamp(0.0, 1.0).mul(255).round().to(torch.uint8).numpy()
return iio.imwrite(
uri,
frames,
extension=".mp4" if uri == "<bytes>" else None,
plugin="FFMPEG",
fps=fps,
codec="libx264",
quality=quality,
macro_block_size=1, # otherwise frames are resized up to a multiple of 16
)
def _in_notebook() -> bool:
"""
Whether we are running inside a jupyter notebook, where html can be displayed inline.
"""
try:
from IPython import get_ipython # always installed in a notebook, it runs the kernel
except ImportError:
return False
shell = get_ipython()
return shell is not None and shell.__class__.__name__ == "ZMQInteractiveShell"
def _video_html(mp4: bytes) -> str:
"""
Wraps mp4 bytes in a self-contained html5 video player.
Args:
mp4 : encoded mp4 video
Returns:
html string, embedding the video as a base64 data uri
"""
encoded = base64.b64encode(mp4).decode("ascii")
# muted, otherwise browsers refuse to autoplay
return f'<video controls loop autoplay muted src="data:video/mp4;base64,{encoded}"></video>'
def _display_video(mp4: bytes) -> None:
"""
Displays an encoded video, inline in a notebook, or in the default browser otherwise.
Args:
mp4 : encoded mp4 video
"""
if len(mp4) > _EMBED_LIMIT:
warnings.warn(
f"Embedding a {len(mp4) / 1024**2:.0f}MB video. Use max_width, fewer frames or a lower"
" quality to shrink it, or save_video to write it to a file instead.",
stacklevel=3,
)
html = _video_html(mp4)
if _in_notebook():
from IPython.display import HTML, display
display(HTML(html))
return
# webbrowser picks the right browser on windows, macos and linux, so we never look at the platform
page = pathlib.Path(tempfile.mkdtemp()) / "showtens_video.html"
page.write_text(html, encoding="utf-8")
webbrowser.open(page.as_uri())