mediapy

   1# Copyright 2026 The mediapy Authors.
   2#
   3# Licensed under the Apache License, Version 2.0 (the "License");
   4# you may not use this file except in compliance with the License.
   5# You may obtain a copy of the License at
   6#
   7#     http://www.apache.org/licenses/LICENSE-2.0
   8#
   9# Unless required by applicable law or agreed to in writing, software
  10# distributed under the License is distributed on an "AS IS" BASIS,
  11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12# See the License for the specific language governing permissions and
  13# limitations under the License.
  14
  15"""`mediapy`: Read/write/show images and videos in an IPython/Jupyter notebook.
  16
  17[**[GitHub source]**](https://github.com/google/mediapy)  
  18[**[API docs]**](https://google.github.io/mediapy/)  
  19[**[PyPI package]**](https://pypi.org/project/mediapy/)  
  20[**[Colab
  21example]**](https://colab.research.google.com/github/google/mediapy/blob/main/mediapy_examples.ipynb)
  22
  23See the [example
  24notebook](https://github.com/google/mediapy/blob/main/mediapy_examples.ipynb),
  25or better yet, [**open it in
  26Colab**](https://colab.research.google.com/github/google/mediapy/blob/main/mediapy_examples.ipynb).
  27
  28## Image examples
  29
  30Display an image (2D or 3D `numpy` array):
  31```python
  32checkerboard = np.kron([[0, 1] * 16, [1, 0] * 16] * 16, np.ones((4, 4)))
  33show_image(checkerboard)
  34```
  35
  36Read and display an image (either local or from the Web):
  37```python
  38IMAGE = 'https://github.com/hhoppe/data/raw/main/image.png'
  39show_image(read_image(IMAGE))
  40```
  41
  42Read and display an image from a local file:
  43```python
  44!wget -q -O /tmp/burano.png {IMAGE}
  45show_image(read_image('/tmp/burano.png'))
  46```
  47
  48Show titled images side-by-side:
  49```python
  50rng = np.random.default_rng(1)
  51images = {
  52    'original': checkerboard,
  53    'darkened': checkerboard * 0.7,
  54    'random': rng.random((32, 32, 3)),
  55}
  56show_images(images, vmin=0.0, vmax=1.0, border=True, height=64)
  57```
  58
  59Compare two images using an interactive slider:
  60```python
  61compare_images([checkerboard, rng.random((128, 128, 3))])
  62```
  63
  64## Video examples
  65
  66Display a video (an iterable of images, e.g., a 3D or 4D array):
  67```python
  68video = moving_circle((100, 100), num_images=10)
  69show_video(video, fps=10)
  70```
  71
  72Show the video frames side-by-side:
  73```python
  74show_images(video, columns=6, border=True, height=64)
  75```
  76
  77Show the frames with their indices:
  78```python
  79show_images({f'{i}': image for i, image in enumerate(video)}, width=32)
  80```
  81
  82Read and display a video (either local or from the Web):
  83```python
  84VIDEO = 'https://github.com/hhoppe/data/raw/main/video.mp4'
  85show_video(read_video(VIDEO))
  86```
  87
  88Create and display a looping two-frame GIF video:
  89```python
  90image1 = resize_image(np.random.default_rng(1).random((10, 10, 3)), (50, 50))
  91show_video([image1, image1 * 0.8], fps=2, codec='gif')
  92```
  93
  94Darken a video frame-by-frame:
  95```python
  96output_path = '/tmp/out.mp4'
  97with VideoReader(VIDEO) as r:
  98  darken_image = lambda image: to_float01(image) * 0.5
  99  with VideoWriter(output_path, shape=r.shape, fps=r.fps, bps=r.bps) as w:
 100    for image in r:
 101      w.add_image(darken_image(image))
 102```
 103"""
 104
 105from __future__ import annotations
 106
 107__docformat__ = 'google'
 108__version__ = '1.2.7'
 109__version_info__ = tuple(int(num) for num in __version__.split('.'))
 110
 111import base64
 112from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence
 113import contextlib
 114import functools
 115import importlib
 116import io
 117import itertools
 118import math
 119import numbers
 120import os  # Package only needed for typing.TYPE_CHECKING.
 121import pathlib
 122import re
 123import shlex
 124import shutil
 125import subprocess
 126import sys
 127import tempfile
 128import typing
 129from typing import Any, TypeVar
 130import urllib.request
 131import warnings
 132
 133import IPython.display
 134import matplotlib.pyplot
 135import numpy as np
 136import numpy.typing as npt
 137import PIL.Image
 138import PIL.ImageOps
 139
 140
 141if not hasattr(PIL.Image, 'Resampling'):  # Allow Pillow<9.0.
 142  PIL.Image.Resampling = PIL.Image  # type: ignore
 143
 144# Selected and reordered here for pdoc documentation.
 145__all__ = [  # noqa: RUF022
 146    'show_image',
 147    'show_images',
 148    'compare_images',
 149    'show_video',
 150    'show_videos',
 151    'read_image',
 152    'write_image',
 153    'read_video',
 154    'write_video',
 155    'VideoReader',
 156    'VideoWriter',
 157    'VideoMetadata',
 158    'compress_image',
 159    'decompress_image',
 160    'compress_video',
 161    'decompress_video',
 162    'html_from_compressed_image',
 163    'html_from_compressed_video',
 164    'resize_image',
 165    'resize_video',
 166    'to_rgb',
 167    'to_type',
 168    'to_float01',
 169    'to_uint8',
 170    'set_output_height',
 171    'set_max_output_height',
 172    'color_ramp',
 173    'moving_circle',
 174    'set_show_save_dir',
 175    'set_ffmpeg',
 176    'video_is_available',
 177]
 178
 179if TYPE_CHECKING:
 180  _ArrayLike = npt.ArrayLike
 181  _DTypeLike = npt.DTypeLike
 182  _NDArray = npt.NDArray[Any]
 183  _DType = np.dtype[Any]
 184  _BaseArray = npt.NDArray[Any]  # Satisfies mypy.
 185else:
 186  # Create named types for use in the `pdoc` documentation.
 187  _ArrayLike = TypeVar('_ArrayLike')
 188  _DTypeLike = TypeVar('_DTypeLike')
 189  _NDArray = TypeVar('_NDArray')
 190  _DType = TypeVar('_DType')  # pylint: disable=invalid-name
 191  _BaseArray = np.ndarray  # Avoids runtime errors in Python 3.14+.
 192
 193_IPYTHON_HTML_SIZE_LIMIT = 10**10  # Unlimited seems to be OK now.
 194_T = TypeVar('_T')
 195_Path = str | os.PathLike[str]
 196
 197# Raw PCM sample formats (as in `ffmpeg -formats`) for the audio dtypes that
 198# `VideoWriter` is able to feed to `ffmpeg`.  Multi-byte formats lack the
 199# byte-order suffix ('le' or 'be'), which is appended for the native order.
 200_FFMPEG_AUDIO_FORMAT_FROM_DTYPE = {
 201    np.uint8: 'u8',
 202    np.int16: 's16',
 203    np.int32: 's32',
 204    np.float32: 'f32',
 205    np.float64: 'f64',
 206}
 207
 208_IMAGE_COMPARISON_HTML = """\
 209<script
 210  defer
 211  src="https://unpkg.com/img-comparison-slider@7/dist/index.js"
 212></script>
 213<link
 214  rel="stylesheet"
 215  href="https://unpkg.com/img-comparison-slider@7/dist/styles.css"
 216/>
 217
 218<img-comparison-slider>
 219  <img slot="first" src="data:image/png;base64,{b64_1}" />
 220  <img slot="second" src="data:image/png;base64,{b64_2}" />
 221</img-comparison-slider>
 222"""
 223
 224# ** Miscellaneous.
 225
 226
 227class _Config:
 228  ffmpeg_name_or_path: _Path = 'ffmpeg'
 229  show_save_dir: _Path | None = None
 230
 231
 232_config = _Config()
 233
 234
 235def _open(path: _Path, *args: Any, **kwargs: Any) -> Any:
 236  """Opens the file; this is a hook for the built-in `open()`."""
 237  return open(path, *args, **kwargs)  # noqa: PTH123
 238
 239
 240def _path_is_local(path: _Path) -> bool:
 241  """Returns True if the path is in the filesystem accessible by `ffmpeg`."""
 242  del path
 243  return True
 244
 245
 246def _search_for_ffmpeg_path() -> str | None:
 247  """Returns a path to the ffmpeg program, or None if not found."""
 248  if filename := shutil.which(_config.ffmpeg_name_or_path):
 249    return str(pathlib.Path(filename))
 250  return None
 251
 252
 253def _print_err(*args: str, **kwargs: Any) -> None:
 254  """Prints arguments to stderr immediately."""
 255  kwargs = {**dict(file=sys.stderr, flush=True), **kwargs}
 256  print(*args, **kwargs)
 257
 258
 259def _chunked(
 260    iterable: Iterable[_T], n: int | None = None
 261) -> Iterator[tuple[_T, ...]]:
 262  """Returns elements collected as tuples of length at most `n` if not None."""
 263
 264  def take(n: int | None, iterable: Iterable[_T]) -> tuple[_T, ...]:
 265    return tuple(itertools.islice(iterable, n))
 266
 267  return iter(functools.partial(take, n, iter(iterable)), ())
 268
 269
 270def _peek_first(iterator: Iterable[_T]) -> tuple[_T, Iterable[_T]]:
 271  """Given an iterator, returns first element and re-initialized iterator.
 272
 273  >>> first_image, images = _peek_first(moving_circle())
 274
 275  Args:
 276    iterator: An input iterator or iterable.
 277
 278  Returns:
 279    A tuple (first_element, iterator_reinitialized) containing:
 280      first_element: The first element of the input.
 281      iterator_reinitialized: A clone of the original iterator/iterable.
 282  """
 283  # Inspired from https://stackoverflow.com/a/12059829/1190077
 284  peeker, iterator_reinitialized = itertools.tee(iterator)
 285  first = next(peeker)
 286  return first, iterator_reinitialized
 287
 288
 289def _check_2d_shape(shape: tuple[int, int]) -> None:
 290  """Checks that `shape` is of the form (height, width) with two integers."""
 291  if len(shape) != 2:
 292    raise ValueError(f'Shape {shape} is not of the form (height, width).')
 293  if not all(isinstance(i, numbers.Integral) for i in shape):
 294    raise ValueError(f'Shape {shape} contains non-integers.')
 295
 296
 297def _run(args: str | Sequence[str]) -> None:
 298  """Executes command, printing output from stdout and stderr.
 299
 300  Args:
 301    args: Command to execute, which can be either a string or a sequence of word
 302      strings, as in `subprocess.run()`.  If `args` is a string, the shell is
 303      invoked to interpret it.
 304
 305  Raises:
 306    RuntimeError: If the command's exit code is nonzero.
 307  """
 308  proc = subprocess.run(
 309      args,
 310      shell=isinstance(args, str),
 311      stdout=subprocess.PIPE,
 312      stderr=subprocess.STDOUT,
 313      check=False,
 314      text=True,
 315  )
 316  print(proc.stdout, end='', flush=True)
 317  if proc.returncode:
 318    raise RuntimeError(
 319        f"Command '{proc.args}' failed with code {proc.returncode}."
 320    )
 321
 322
 323def _display_html(text: str, /) -> None:
 324  """In a Jupyter notebook, display the HTML `text`."""
 325  IPython.display.display(IPython.display.HTML(text))  # type: ignore
 326
 327
 328def set_ffmpeg(name_or_path: _Path) -> None:
 329  """Specifies the name or path for the `ffmpeg` external program.
 330
 331  The `ffmpeg` program is required for compressing and decompressing video.
 332  (It is used in `read_video`, `write_video`, `show_video`, `show_videos`,
 333  etc.)
 334
 335  Args:
 336    name_or_path: Either a filename within a directory of `os.environ['PATH']`
 337      or a filepath.  The default setting is 'ffmpeg'.
 338  """
 339  _config.ffmpeg_name_or_path = name_or_path
 340
 341
 342def set_output_height(num_pixels: int) -> None:
 343  """Overrides the height of the current output cell, if using Colab."""
 344  with contextlib.suppress(ModuleNotFoundError, AttributeError):
 345    # We want to fail gracefully for non-Colab IPython notebooks.
 346    output = importlib.import_module('google.colab.output')
 347    s = f'google.colab.output.setIframeHeight("{num_pixels}px")'
 348    output.eval_js(s)
 349
 350
 351def set_max_output_height(num_pixels: int) -> None:
 352  """Sets the maximum height of the current output cell, if using Colab."""
 353  with contextlib.suppress(ModuleNotFoundError, AttributeError):
 354    # We want to fail gracefully for non-Colab IPython notebooks.
 355    output = importlib.import_module('google.colab.output')
 356    s = (
 357        'google.colab.output.setIframeHeight('
 358        f'0, true, {{maxHeight: {num_pixels}}})'
 359    )
 360    output.eval_js(s)
 361
 362
 363# ** Type conversions.
 364
 365
 366def _as_valid_media_type(dtype: _DTypeLike) -> _DType:
 367  """Returns validated media data type."""
 368  dtype = np.dtype(dtype)
 369  if not issubclass(dtype.type, (np.unsignedinteger, np.floating)):
 370    raise ValueError(
 371        f'Type {dtype} is not a valid media data type (uint or float).'
 372    )
 373  return dtype
 374
 375
 376def _as_valid_media_array(x: _ArrayLike) -> _NDArray:
 377  """Converts to ndarray (if not already), and checks validity of data type."""
 378  a = np.asarray(x)
 379  if a.dtype == bool:
 380    a = a.astype(np.uint8) * np.iinfo(np.uint8).max
 381  _as_valid_media_type(a.dtype)
 382  return a
 383
 384
 385def to_type(array: _ArrayLike, dtype: _DTypeLike) -> _NDArray:
 386  """Returns media array converted to specified type.
 387
 388  A "media array" is one in which the dtype is either a floating-point type
 389  (np.float32 or np.float64) or an unsigned integer type.  The array values are
 390  assumed to lie in the range [0.0, 1.0] for floating-point values, and in the
 391  full range for unsigned integers, e.g. [0, 255] for np.uint8.
 392
 393  Conversion between integers and floats maps uint(0) to 0.0 and uint(MAX) to
 394  1.0.  The input array may also be of type bool, whereby True maps to
 395  uint(MAX) or 1.0.  The values are scaled and clamped as appropriate during
 396  type conversions.
 397
 398  Args:
 399    array: Input array-like object (floating-point, unsigned int, or bool).
 400    dtype: Desired output type (floating-point or unsigned int).
 401
 402  Returns:
 403    Array `a` if it is already of the specified dtype, else a converted array.
 404  """
 405  a = np.asarray(array)
 406  dtype2: np.dtype[Any] = np.dtype(dtype)
 407  del array, dtype
 408  if a.dtype != bool:
 409    _as_valid_media_type(a.dtype)  # Verify that 'a' has a valid dtype.
 410  if a.dtype == bool:
 411    result = a.astype(dtype2)
 412    if np.issubdtype(dtype2, np.unsignedinteger):
 413      # Force in-place multiplication to preserve the exact dtype.
 414      np.multiply(result, np.iinfo(dtype2).max, out=result)
 415  elif a.dtype == dtype2:
 416    result = a
 417  elif np.issubdtype(dtype2, np.unsignedinteger):
 418    if np.issubdtype(a.dtype, np.unsignedinteger):
 419      src_max: float = np.iinfo(a.dtype).max
 420    else:
 421      a = np.clip(a, 0.0, 1.0)
 422      src_max = 1.0
 423    dst_max = np.iinfo(dtype2).max
 424    if dst_max <= np.iinfo(np.uint16).max:
 425      scale = np.array(dst_max / src_max, dtype=np.float32)
 426      result = (a * scale + 0.5).astype(dtype2)
 427    elif dst_max <= np.iinfo(np.uint32).max:
 428      result = (a.astype(np.float64) * (dst_max / src_max) + 0.5).astype(dtype2)
 429    else:
 430      # https://stackoverflow.com/a/66306123/
 431      a = a.astype(np.float64) * (dst_max / src_max) + 0.5
 432      dst = np.atleast_1d(a)
 433      values_too_large = dst >= np.float64(dst_max)
 434      with np.errstate(invalid='ignore'):
 435        dst = dst.astype(dtype2)
 436      dst[values_too_large] = dst_max
 437      result = dst if a.ndim > 0 else dst[0]
 438  else:
 439    assert np.issubdtype(dtype2, np.floating)
 440    result = a.astype(dtype2)
 441    if np.issubdtype(a.dtype, np.unsignedinteger):
 442      result = result / dtype2.type(np.iinfo(a.dtype).max)
 443  return result
 444
 445
 446def to_float01(a: _ArrayLike, dtype: _DTypeLike = np.float32) -> _NDArray:
 447  """If array has unsigned integers, rescales them to the range [0.0, 1.0].
 448
 449  Scaling is such that uint(0) maps to 0.0 and uint(MAX) maps to 1.0.  See
 450  `to_type`.
 451
 452  Args:
 453    a: Input array.
 454    dtype: Desired floating-point type if rescaling occurs.
 455
 456  Returns:
 457    A new array of dtype values in the range [0.0, 1.0] if the input array `a`
 458    contains unsigned integers; otherwise, array `a` is returned unchanged.
 459  """
 460  a = np.asarray(a)
 461  dtype = np.dtype(dtype)
 462  if not np.issubdtype(dtype, np.floating):
 463    raise ValueError(f'Type {dtype} is not floating-point.')
 464  if np.issubdtype(a.dtype, np.floating):
 465    return a
 466  return to_type(a, dtype)
 467
 468
 469def to_uint8(a: _ArrayLike) -> _NDArray:
 470  """Returns array converted to uint8 values; see `to_type`."""
 471  return to_type(a, np.uint8)
 472
 473
 474# ** Functions to generate example image and video data.
 475
 476
 477def color_ramp(
 478    shape: tuple[int, int] = (64, 64), *, dtype: _DTypeLike = np.float32
 479) -> _NDArray:
 480  """Returns an image of a red-green color gradient.
 481
 482  This is useful for quick experimentation and testing.  See also
 483  `moving_circle` to generate a sample video.
 484
 485  Args:
 486    shape: 2D spatial dimensions (height, width) of generated image.
 487    dtype: Type (uint or floating) of resulting pixel values.
 488  """
 489  _check_2d_shape(shape)
 490  dtype = _as_valid_media_type(dtype)
 491  yx = (np.moveaxis(np.indices(shape), 0, -1) + 0.5) / shape
 492  image = np.insert(yx, 2, 0.0, axis=-1)
 493  return to_type(image, dtype)
 494
 495
 496def moving_circle(
 497    shape: tuple[int, int] = (256, 256),
 498    num_images: int = 10,
 499    *,
 500    dtype: _DTypeLike = np.float32,
 501) -> _NDArray:
 502  """Returns a video of a circle moving in front of a color ramp.
 503
 504  This is useful for quick experimentation and testing.  See also `color_ramp`
 505  to generate a sample image.
 506
 507  >>> show_video(moving_circle((480, 640), 60), fps=60)
 508
 509  Args:
 510    shape: 2D spatial dimensions (height, width) of generated video.
 511    num_images: Number of video frames.
 512    dtype: Type (uint or floating) of resulting pixel values.
 513  """
 514  _check_2d_shape(shape)
 515  dtype = np.dtype(dtype)
 516
 517  def generate_image(image_index: int) -> _NDArray:
 518    """Returns a video frame image."""
 519    image = color_ramp(shape, dtype=dtype)
 520    yx = np.moveaxis(np.indices(shape), 0, -1)
 521    center = shape[0] * 0.6, shape[1] * (image_index + 0.5) / num_images
 522    radius_squared = (min(shape) * 0.1) ** 2
 523    inside = np.sum((yx - center) ** 2, axis=-1) < radius_squared
 524    white_circle_color = 1.0, 1.0, 1.0
 525    if np.issubdtype(dtype, np.unsignedinteger):
 526      white_circle_color = to_type([white_circle_color], dtype)[0]
 527    image[inside] = white_circle_color
 528    return image
 529
 530  return np.array([generate_image(i) for i in range(num_images)])
 531
 532
 533# ** Color-space conversions.
 534
 535# Same matrix values as in two sources:
 536# https://github.com/scikit-image/scikit-image/blob/master/skimage/color/colorconv.py#L377
 537# https://github.com/tensorflow/tensorflow/blob/r1.14/tensorflow/python/ops/image_ops_impl.py#L2754
 538_YUV_FROM_RGB_MATRIX = np.array(
 539    [
 540        [0.299, -0.14714119, 0.61497538],
 541        [0.587, -0.28886916, -0.51496512],
 542        [0.114, 0.43601035, -0.10001026],
 543    ],
 544    dtype=np.float32,
 545)
 546_RGB_FROM_YUV_MATRIX = np.linalg.inv(_YUV_FROM_RGB_MATRIX)
 547_YUV_CHROMA_OFFSET = np.array([0.0, 0.5, 0.5], dtype=np.float32)
 548
 549
 550def yuv_from_rgb(rgb: _ArrayLike) -> _NDArray:
 551  """Returns the RGB image/video mapped to YUV [0,1] color space.
 552
 553  Note that the "YUV" color space used by video compressors is actually YCbCr!
 554
 555  Args:
 556    rgb: Input image in sRGB space.
 557  """
 558  rgb = to_float01(rgb)
 559  if rgb.shape[-1] != 3:
 560    raise ValueError(f'The last dimension in {rgb.shape} is not 3.')
 561  return rgb @ _YUV_FROM_RGB_MATRIX + _YUV_CHROMA_OFFSET
 562
 563
 564def rgb_from_yuv(yuv: _ArrayLike) -> _NDArray:
 565  """Returns the YUV image/video mapped to RGB [0,1] color space."""
 566  yuv = to_float01(yuv)
 567  if yuv.shape[-1] != 3:
 568    raise ValueError(f'The last dimension in {yuv.shape} is not 3.')
 569  return (yuv - _YUV_CHROMA_OFFSET) @ _RGB_FROM_YUV_MATRIX
 570
 571
 572# Same matrix values as in
 573# https://github.com/scikit-image/scikit-image/blob/master/skimage/color/colorconv.py#L1654
 574# and https://en.wikipedia.org/wiki/YUV#Studio_swing_for_BT.601
 575_YCBCR_FROM_RGB_MATRIX = np.array(
 576    [
 577        [65.481, 128.553, 24.966],
 578        [-37.797, -74.203, 112.0],
 579        [112.0, -93.786, -18.214],
 580    ],
 581    dtype=np.float32,
 582).transpose()
 583_RGB_FROM_YCBCR_MATRIX = np.linalg.inv(_YCBCR_FROM_RGB_MATRIX)
 584_YCBCR_OFFSET = np.array([16.0, 128.0, 128.0], dtype=np.float32)
 585# Note that _YCBCR_FROM_RGB_MATRIX =~ _YUV_FROM_RGB_MATRIX * [219, 256, 182];
 586# https://en.wikipedia.org/wiki/YUV: "Y' values are conventionally shifted and
 587# scaled to the range [16, 235] (referred to as studio swing or 'TV levels')";
 588# "studio range of 16-240 for U and V".  (Where does value 182 come from?)
 589
 590
 591def ycbcr_from_rgb(rgb: _ArrayLike) -> _NDArray:
 592  """Returns the RGB image/video mapped to YCbCr [0,1] color space.
 593
 594  The YCbCr color space is the one called "YUV" by video compressors.
 595
 596  Args:
 597    rgb: Input image in sRGB space.
 598  """
 599  rgb = to_float01(rgb)
 600  if rgb.shape[-1] != 3:
 601    raise ValueError(f'The last dimension in {rgb.shape} is not 3.')
 602  return (rgb @ _YCBCR_FROM_RGB_MATRIX + _YCBCR_OFFSET) / 255.0
 603
 604
 605def rgb_from_ycbcr(ycbcr: _ArrayLike) -> _NDArray:
 606  """Returns the YCbCr image/video mapped to RGB [0,1] color space."""
 607  ycbcr = to_float01(ycbcr)
 608  if ycbcr.shape[-1] != 3:
 609    raise ValueError(f'The last dimension in {ycbcr.shape} is not 3.')
 610  return (ycbcr * 255.0 - _YCBCR_OFFSET) @ _RGB_FROM_YCBCR_MATRIX
 611
 612
 613# ** Image processing.
 614
 615
 616def _pil_image(image: _ArrayLike, mode: str | None = None) -> PIL.Image.Image:
 617  """Returns a PIL image given a numpy matrix (either uint8 or float [0,1])."""
 618  image = _as_valid_media_array(image)
 619  if image.ndim not in (2, 3):
 620    raise ValueError(f'Image shape {image.shape} is neither 2D nor 3D.')
 621  pil_image: PIL.Image.Image = PIL.Image.fromarray(image, mode=mode)
 622  return pil_image
 623
 624
 625def resize_image(image: _ArrayLike, shape: tuple[int, int]) -> _NDArray:
 626  """Resizes image to specified spatial dimensions using a Lanczos filter.
 627
 628  Args:
 629    image: Array-like 2D or 3D object, where dtype is uint or floating-point.
 630    shape: 2D spatial dimensions (height, width) of output image.
 631
 632  Returns:
 633    A resampled image whose spatial dimensions match `shape`.
 634  """
 635  image = _as_valid_media_array(image)
 636  if image.ndim not in (2, 3):
 637    raise ValueError(f'Image shape {image.shape} is neither 2D nor 3D.')
 638  _check_2d_shape(shape)
 639
 640  # A PIL image can be multichannel only if it has 3 or 4 uint8 channels,
 641  # and it can be resized only if it is uint8 or float32.
 642  supported_single_channel = (
 643      np.issubdtype(image.dtype, np.floating) or image.dtype == np.uint8
 644  ) and image.ndim == 2
 645  supported_multichannel = (
 646      image.dtype == np.uint8 and image.ndim == 3 and image.shape[2] in (3, 4)
 647  )
 648  if supported_single_channel or supported_multichannel:
 649    h, w = shape
 650    return np.array(
 651        _pil_image(image).resize((w, h), resample=PIL.Image.Resampling.LANCZOS),
 652        dtype=image.dtype,
 653    )
 654  if image.ndim == 2:
 655    # We convert to floating-point for resizing and convert back.
 656    return to_type(resize_image(to_float01(image), shape), image.dtype)
 657  # We resize each image channel individually.
 658  return np.dstack(
 659      [resize_image(channel, shape) for channel in np.moveaxis(image, -1, 0)]
 660  )
 661
 662
 663# ** Video processing.
 664
 665
 666def resize_video(video: Iterable[_NDArray], shape: tuple[int, int]) -> _NDArray:
 667  """Resizes `video` to specified spatial dimensions using a Lanczos filter.
 668
 669  Args:
 670    video: Iterable of images.
 671    shape: 2D spatial dimensions (height, width) of output video.
 672
 673  Returns:
 674    A resampled video whose spatial dimensions match `shape`.
 675  """
 676  _check_2d_shape(shape)
 677  return np.array([resize_image(image, shape) for image in video])
 678
 679
 680# ** General I/O.
 681
 682
 683def _is_url(path_or_url: _Path) -> bool:
 684  return isinstance(path_or_url, str) and path_or_url.startswith(
 685      ('http://', 'https://', 'file://')
 686  )
 687
 688
 689def read_contents(path_or_url: _Path) -> bytes:
 690  """Returns the contents of the file specified by either a path or URL."""
 691  data: bytes
 692  if _is_url(path_or_url):
 693    assert isinstance(path_or_url, str)
 694    headers = {'User-Agent': 'Chrome'}
 695    request = urllib.request.Request(path_or_url, headers=headers)
 696    with urllib.request.urlopen(request) as response:
 697      data = response.read()
 698  else:
 699    with _open(path_or_url, 'rb') as f:
 700      data = f.read()
 701  return data
 702
 703
 704@contextlib.contextmanager
 705def _read_via_local_file(path_or_url: _Path) -> Generator[str, None, None]:
 706  """Context to copy a remote file locally to read from it.
 707
 708  Args:
 709    path_or_url: File, which may be remote.
 710
 711  Yields:
 712    The name of a local file which may be a copy of a remote file.
 713  """
 714  if _is_url(path_or_url) or not _path_is_local(path_or_url):
 715    suffix = pathlib.Path(path_or_url).suffix
 716    with tempfile.TemporaryDirectory() as directory_name:
 717      tmp_path = pathlib.Path(directory_name) / f'file{suffix}'
 718      tmp_path.write_bytes(read_contents(path_or_url))
 719      yield str(tmp_path)
 720  else:
 721    yield str(path_or_url)
 722
 723
 724@contextlib.contextmanager
 725def _write_via_local_file(path: _Path) -> Generator[str, None, None]:
 726  """Context to write a temporary local file and subsequently copy it remotely.
 727
 728  Args:
 729    path: File, which may be remote.
 730
 731  Yields:
 732    The name of a local file which may be subsequently copied remotely.
 733  """
 734  if _path_is_local(path):
 735    yield str(path)
 736  else:
 737    suffix = pathlib.Path(path).suffix
 738    with tempfile.TemporaryDirectory() as directory_name:
 739      tmp_path = pathlib.Path(directory_name) / f'file{suffix}'
 740      yield str(tmp_path)
 741      with _open(path, mode='wb') as f:
 742        f.write(tmp_path.read_bytes())
 743
 744
 745@contextlib.contextmanager
 746def _audio_via_local_file(audio: _NDArray) -> Generator[str, None, None]:
 747  """Context to write audio samples to a temporary local file.
 748
 749  Args:
 750    audio: Array of raw PCM audio samples, with shape (N,) for mono or (N, C)
 751      for C interleaved channels.
 752
 753  Yields:
 754    The name of a local file containing the raw samples.
 755  """
 756  with tempfile.TemporaryDirectory() as directory_name:
 757    tmp_path = pathlib.Path(directory_name) / 'audio.raw'
 758    tmp_path.write_bytes(audio.tobytes())
 759    yield str(tmp_path)
 760
 761
 762class set_show_save_dir:  # pylint: disable=invalid-name
 763  """Save all titled output from `show_*()` calls into files.
 764
 765  If the specified `directory` is not None, all titled images and videos
 766  displayed by `show_image`, `show_images`, `show_video`, and `show_videos` are
 767  also saved as files within the directory.
 768
 769  It can be used either to set the state or as a context manager:
 770
 771  >>> set_show_save_dir('/tmp')
 772  >>> show_image(color_ramp(), title='image1')  # Creates /tmp/image1.png.
 773  >>> show_video(moving_circle(), title='video2')  # Creates /tmp/video2.mp4.
 774  >>> set_show_save_dir(None)
 775
 776  >>> with set_show_save_dir('/tmp'):
 777  ...   show_image(color_ramp(), title='image1')  # Creates /tmp/image1.png.
 778  ...   show_video(moving_circle(), title='video2')  # Creates /tmp/video2.mp4.
 779  """
 780
 781  def __init__(self, directory: _Path | None):
 782    self._old_show_save_dir = _config.show_save_dir
 783    _config.show_save_dir = directory
 784
 785  def __enter__(self) -> None:
 786    pass
 787
 788  def __exit__(self, *_: Any) -> None:
 789    _config.show_save_dir = self._old_show_save_dir
 790
 791
 792# ** Image I/O.
 793
 794
 795def read_image(
 796    path_or_url: _Path,
 797    *,
 798    apply_exif_transpose: bool = True,
 799    dtype: _DTypeLike | None = None,
 800) -> _NDArray:
 801  """Returns an image read from a file path or URL.
 802
 803  Decoding is performed using `PIL`, which supports `uint8` images with 1, 3,
 804  or 4 channels and `uint16` images with a single channel.
 805
 806  Args:
 807    path_or_url: Path of input file.
 808    apply_exif_transpose: If True, rotate image according to EXIF orientation.
 809    dtype: Data type of the returned array.  If None, `np.uint8` or `np.uint16`
 810      is inferred automatically.
 811  """
 812  data = read_contents(path_or_url)
 813  return decompress_image(data, dtype, apply_exif_transpose)
 814
 815
 816def write_image(
 817    path: _Path, image: _ArrayLike, fmt: str = 'png', **kwargs: Any
 818) -> None:
 819  """Writes an image to a file.
 820
 821  Encoding is performed using `PIL`, which supports `uint8` images with 1, 3,
 822  or 4 channels and `uint16` images with a single channel.
 823
 824  File format is explicitly provided by `fmt` and not inferred by `path`.
 825
 826  Args:
 827    path: Path of output file.
 828    image: Array-like object.  If its type is float, it is converted to np.uint8
 829      using `to_uint8` (thus clamping to the input to the range [0.0, 1.0]).
 830      Otherwise it must be np.uint8 or np.uint16.
 831    fmt: Desired compression encoding, e.g. 'png'.
 832    **kwargs: Additional parameters for `PIL.Image.save()`.
 833  """
 834  image = _as_valid_media_array(image)
 835  if np.issubdtype(image.dtype, np.floating):
 836    image = to_uint8(image)
 837  with _open(path, 'wb') as f:
 838    _pil_image(image).save(f, format=fmt, **kwargs)
 839
 840
 841def to_rgb(
 842    array: _ArrayLike,
 843    *,
 844    vmin: float | None = None,
 845    vmax: float | None = None,
 846    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
 847) -> _NDArray:
 848  """Maps scalar values to RGB using value bounds and a color map.
 849
 850  Args:
 851    array: Scalar values, with arbitrary shape.
 852    vmin: Explicit min value for remapping; if None, it is obtained as the
 853      minimum finite value of `array`.
 854    vmax: Explicit max value for remapping; if None, it is obtained as the
 855      maximum finite value of `array`.
 856    cmap: A `pyplot` color map or callable, to map from 1D value to 3D or 4D
 857      color.
 858
 859  Returns:
 860    A new array in which each element is affinely mapped from [vmin, vmax]
 861    to [0.0, 1.0] and then color-mapped.
 862  """
 863  a = _as_valid_media_array(array)
 864  del array
 865  # For future numpy version 1.7.0:
 866  # vmin = np.amin(a, where=np.isfinite(a)) if vmin is None else vmin
 867  # vmax = np.amax(a, where=np.isfinite(a)) if vmax is None else vmax
 868  vmin = np.amin(np.where(np.isfinite(a), a, np.inf)) if vmin is None else vmin
 869  vmax = np.amax(np.where(np.isfinite(a), a, -np.inf)) if vmax is None else vmax
 870  a = (a.astype('float') - vmin) / (vmax - vmin + np.finfo(float).eps)  # pylint: disable=no-member
 871  if isinstance(cmap, str):
 872    if hasattr(matplotlib, 'colormaps'):
 873      rgb_from_scalar: Any = matplotlib.colormaps[cmap]  # Newer version.
 874    else:
 875      # pylint: disable-next=no-member
 876      rgb_from_scalar = matplotlib.pyplot.cm.get_cmap(cmap)  # type: ignore[attr-defined, unused-ignore]
 877  else:
 878    rgb_from_scalar = cmap
 879  a = cast(_NDArray, rgb_from_scalar(a))
 880  # If there is a fully opaque alpha channel, remove it.
 881  if a.shape[-1] == 4 and np.all(to_float01(a[..., 3])) == 1.0:
 882    a = a[..., :3]
 883  return a
 884
 885
 886def compress_image(
 887    image: _ArrayLike, *, fmt: str = 'png', **kwargs: Any
 888) -> bytes:
 889  """Returns a buffer containing a compressed image.
 890
 891  Args:
 892    image: Array in a format supported by `PIL`, e.g. np.uint8 or np.uint16.
 893    fmt: Desired compression encoding, e.g. 'png'.
 894    **kwargs: Options for `PIL.save()`, e.g. `optimize=True` for greater
 895      compression.
 896  """
 897  image = _as_valid_media_array(image)
 898  with io.BytesIO() as output:
 899    _pil_image(image).save(output, format=fmt, **kwargs)
 900    return output.getvalue()
 901
 902
 903def decompress_image(
 904    data: bytes,
 905    dtype: _DTypeLike | None = None,
 906    apply_exif_transpose: bool = True,
 907) -> _NDArray:
 908  """Returns an image from a compressed data buffer.
 909
 910  Decoding is performed using `PIL`, which supports `uint8` images with 1, 3,
 911  or 4 channels and `uint16` images with a single channel.
 912
 913  Args:
 914    data: Buffer containing compressed image.
 915    dtype: Data type of the returned array.  If None, `np.uint8` or `np.uint16`
 916      is inferred automatically.
 917    apply_exif_transpose: If True, rotate image according to EXIF orientation.
 918  """
 919  pil_image: PIL.Image.Image = PIL.Image.open(io.BytesIO(data))
 920  if apply_exif_transpose:
 921    pil_image = PIL.ImageOps.exif_transpose(pil_image)  # Future: in_place=True.
 922  if dtype is None:
 923    dtype = np.uint16 if pil_image.mode.startswith('I') else np.uint8
 924  return np.array(pil_image, dtype=dtype)
 925
 926
 927def html_from_compressed_image(
 928    data: bytes,
 929    width: int,
 930    height: int,
 931    *,
 932    title: str | None = None,
 933    border: bool | str = False,
 934    pixelated: bool = True,
 935    fmt: str = 'png',
 936) -> str:
 937  """Returns an HTML string with an image tag containing encoded data.
 938
 939  Args:
 940    data: Compressed image bytes.
 941    width: Width of HTML image in pixels.
 942    height: Height of HTML image in pixels.
 943    title: Optional text shown centered above image.
 944    border: If `bool`, whether to place a black boundary around the image, or if
 945      `str`, the boundary CSS style.
 946    pixelated: If True, sets the CSS style to 'image-rendering: pixelated;'.
 947    fmt: Compression encoding.
 948  """
 949  b64 = base64.b64encode(data).decode('utf-8')
 950  if isinstance(border, str):
 951    border = f'{border}; '
 952  elif border:
 953    border = 'border:1px solid black; '
 954  else:
 955    border = ''
 956  s_pixelated = 'pixelated' if pixelated else 'auto'
 957  s = (
 958      f'<img width="{width}" height="{height}"'
 959      f' style="{border}image-rendering:{s_pixelated}; object-fit:cover;"'
 960      f' src="data:image/{fmt};base64,{b64}"/>'
 961  )
 962  if title is not None:
 963    s = f"""<div style="display:flex; align-items:left;">
 964      <div style="display:flex; flex-direction:column; align-items:center;">
 965      <div>{title}</div><div>{s}</div></div></div>"""
 966  return s
 967
 968
 969def _get_width_height(
 970    width: int | None, height: int | None, shape: tuple[int, int]
 971) -> tuple[int, int]:
 972  """Returns (width, height) given optional parameters and image shape."""
 973  assert len(shape) == 2, shape
 974  if width and height:
 975    return width, height
 976  if width and not height:
 977    return width, int(width * (shape[0] / shape[1]) + 0.5)
 978  if height and not width:
 979    return int(height * (shape[1] / shape[0]) + 0.5), height
 980  return shape[1], shape[0]
 981
 982
 983def _ensure_mapped_to_rgb(
 984    image: _ArrayLike,
 985    *,
 986    vmin: float | None = None,
 987    vmax: float | None = None,
 988    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
 989) -> _NDArray:
 990  """Ensure image is mapped to RGB."""
 991  image = _as_valid_media_array(image)
 992  if not (image.ndim == 2 or (image.ndim == 3 and image.shape[2] in (1, 3, 4))):
 993    raise ValueError(
 994        f'Image with shape {image.shape} is neither a 2D array'
 995        ' nor a 3D array with 1, 3, or 4 channels.'
 996    )
 997  if image.ndim == 3 and image.shape[2] == 1:
 998    image = image[:, :, 0]
 999  if image.ndim == 2:
1000    image = to_rgb(image, vmin=vmin, vmax=vmax, cmap=cmap)
1001  return image
1002
1003
1004def show_image(
1005    image: _ArrayLike, *, title: str | None = None, **kwargs: Any
1006) -> str | None:
1007  """Displays an image in the notebook and optionally saves it to a file.
1008
1009  See `show_images`.
1010
1011  >>> rng = np.random.default_rng(1)
1012  >>> show_image(rng.random((100, 100)))
1013  >>> show_image(rng.integers(0, 256, size=(80, 80, 3), dtype='uint8'))
1014  >>> show_image(rng.random((10, 10)) - 0.5, cmap='bwr', height=100)
1015  >>> show_image(read_image('/tmp/image.png'))
1016  >>> url = 'https://github.com/hhoppe/data/raw/main/image.png'
1017  >>> show_image(read_image(url))
1018
1019  Args:
1020    image: 2D array-like, or 3D array-like with 1, 3, or 4 channels.
1021    title: Optional text shown centered above the image.
1022    **kwargs: See `show_images`.
1023
1024  Returns:
1025    html string if `return_html` is `True`.
1026  """
1027  return show_images([np.asarray(image)], [title], **kwargs)
1028
1029
1030def show_images(
1031    images: Iterable[_ArrayLike] | Mapping[str, _ArrayLike],
1032    titles: Iterable[str | None] | None = None,
1033    *,
1034    width: int | None = None,
1035    height: int | None = None,
1036    downsample: bool = True,
1037    columns: int | None = None,
1038    vmin: float | None = None,
1039    vmax: float | None = None,
1040    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
1041    border: bool | str = False,
1042    ylabel: str = '',
1043    html_class: str = 'show_images',
1044    pixelated: bool | None = None,
1045    return_html: bool = False,
1046) -> str | None:
1047  """Displays a row of images in the IPython/Jupyter notebook.
1048
1049  If a directory has been specified using `set_show_save_dir`, also saves each
1050  titled image to a file in that directory based on its title.
1051
1052  >>> image1 = np.random.default_rng(1).random((64, 64, 3))
1053  >>> image2 = color_ramp((64, 64))
1054  >>> show_images([image1, image2])
1055  >>> show_images({'random image': image1, 'color ramp': image2}, height=128)
1056  >>> show_images([image1, image2] * 5, columns=4, border=True)
1057
1058  Args:
1059    images: Iterable of images, or dictionary of `{title: image}`.  Each image
1060      must be either a 2D array or a 3D array with 1, 3, or 4 channels.
1061    titles: Optional strings shown above the corresponding images.
1062    width: Optional, overrides displayed width (in pixels).
1063    height: Optional, overrides displayed height (in pixels).
1064    downsample: If True, each image whose width or height is greater than the
1065      specified `width` or `height` is resampled to the display resolution. This
1066      improves antialiasing and reduces the size of the notebook.
1067    columns: Optional, maximum number of images per row.
1068    vmin: For single-channel image, explicit min value for display.
1069    vmax: For single-channel image, explicit max value for display.
1070    cmap: For single-channel image, `pyplot` color map or callable to map 1D to
1071      3D color.
1072    border: If `bool`, whether to place a black boundary around the image, or if
1073      `str`, the boundary CSS style.
1074    ylabel: Text (rotated by 90 degrees) shown on the left of each row.
1075    html_class: CSS class name used in definition of HTML element.
1076    pixelated: If True, sets the CSS style to 'image-rendering: pixelated;'; if
1077      False, sets 'image-rendering: auto'; if None, uses pixelated rendering
1078      only on images for which `width` or `height` introduces magnification.
1079    return_html: If `True` return the raw HTML `str` instead of displaying.
1080
1081  Returns:
1082    html string if `return_html` is `True`.
1083  """
1084  list_titles: list[str | None]
1085  if isinstance(images, Mapping):
1086    if titles is not None:
1087      raise ValueError('Cannot have images dictionary and titles parameter.')
1088    list_images0 = list(images.values())
1089    list_titles = cast(list[str | None], list(images.keys()))
1090  else:
1091    list_images0 = list(images)
1092    list_titles = [None] * len(list_images0) if titles is None else list(titles)
1093    if len(list_images0) != len(list_titles):
1094      raise ValueError(
1095          'Number of images does not match number of titles'
1096          f' ({len(list_images0)} vs {len(list_titles)}).'
1097      )
1098
1099  list_images = [
1100      _ensure_mapped_to_rgb(image, vmin=vmin, vmax=vmax, cmap=cmap)
1101      for image in list_images0
1102  ]
1103  del list_images0
1104
1105  def maybe_downsample(image: _NDArray) -> _NDArray:
1106    shape = image.shape[0], image.shape[1]
1107    w, h = _get_width_height(width, height, shape)
1108    if w < shape[1] or h < shape[0]:
1109      image = resize_image(image, (h, w))
1110    return image
1111
1112  if downsample:
1113    list_images = [maybe_downsample(image) for image in list_images]
1114  png_datas = [compress_image(to_uint8(image)) for image in list_images]
1115
1116  for title, png_data in zip(list_titles, png_datas):
1117    if title is not None and _config.show_save_dir:
1118      path = pathlib.Path(_config.show_save_dir) / f'{title}.png'
1119      with _open(path, mode='wb') as f:
1120        f.write(png_data)
1121
1122  def html_from_compressed_images() -> str:
1123    html_strings = []
1124    for image, title, png_data in zip(list_images, list_titles, png_datas):
1125      shape = image.shape[0], image.shape[1]
1126      w, h = _get_width_height(width, height, shape)
1127      magnified = h > image.shape[0] or w > image.shape[1]
1128      pixelated2 = pixelated if pixelated is not None else magnified
1129      html_strings.append(
1130          html_from_compressed_image(
1131              png_data, w, h, title=title, border=border, pixelated=pixelated2
1132          )
1133      )
1134    # Create single-row tables each with no more than 'columns' elements.
1135    table_strings = []
1136    for row_html_strings in _chunked(html_strings, columns):
1137      td = '<td style="padding:1px;">'
1138      s = ''.join(f'{td}{e}</td>' for e in row_html_strings)
1139      if ylabel:
1140        style = 'writing-mode:vertical-lr; transform:rotate(180deg);'
1141        s = f'{td}<span style="{style}">{ylabel}</span></td>' + s
1142      table_strings.append(
1143          f'<table class="{html_class}"'
1144          f' style="border-spacing:0px;"><tr>{s}</tr></table>'
1145      )
1146    return ''.join(table_strings)
1147
1148  s = html_from_compressed_images()
1149  while len(s) > _IPYTHON_HTML_SIZE_LIMIT * 0.5:
1150    warnings.warn('mediapy: subsampling images to reduce HTML size')
1151    list_images = [image[::2, ::2] for image in list_images]
1152    png_datas = [compress_image(to_uint8(image)) for image in list_images]
1153    s = html_from_compressed_images()
1154  if return_html:
1155    return s
1156  _display_html(s)
1157  return None
1158
1159
1160def compare_images(
1161    images: Iterable[_ArrayLike],
1162    *,
1163    vmin: float | None = None,
1164    vmax: float | None = None,
1165    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
1166) -> None:
1167  """Compare two images using an interactive slider.
1168
1169  Displays an HTML slider component to interactively swipe between two images.
1170  The slider functionality requires that the web browser have Internet access.
1171  See additional info in `https://github.com/sneas/img-comparison-slider`.
1172
1173  >>> image1 = np.random.default_rng(1).random((64, 64, 3))
1174  >>> image2 = color_ramp((64, 64))
1175  >>> compare_images([image1, image2])
1176
1177  Args:
1178    images: Iterable of images.  Each image must be either a 2D array or a 3D
1179      array with 1, 3, or 4 channels.  There must be exactly two images.
1180    vmin: For single-channel image, explicit min value for display.
1181    vmax: For single-channel image, explicit max value for display.
1182    cmap: For single-channel image, `pyplot` color map or callable to map 1D to
1183      3D color.
1184  """
1185  list_images = [
1186      _ensure_mapped_to_rgb(image, vmin=vmin, vmax=vmax, cmap=cmap)
1187      for image in images
1188  ]
1189  if len(list_images) != 2:
1190    raise ValueError('The number of images must be 2.')
1191  png_datas = [compress_image(to_uint8(image)) for image in list_images]
1192  b64_1, b64_2 = [
1193      base64.b64encode(png_data).decode('utf-8') for png_data in png_datas
1194  ]
1195  s = _IMAGE_COMPARISON_HTML.replace('{b64_1}', b64_1).replace('{b64_2}', b64_2)
1196  _display_html(s)
1197
1198
1199# ** Video I/O.
1200
1201
1202def _filename_suffix_from_codec(codec: str) -> str:
1203  if codec == 'gif':
1204    return '.gif'
1205  if codec == 'vp9':
1206    return '.webm'
1207
1208  return '.mp4'
1209
1210
1211def _get_ffmpeg_path() -> str:
1212  path = _search_for_ffmpeg_path()
1213  if not path:
1214    raise RuntimeError(
1215        f"Program '{_config.ffmpeg_name_or_path}' is not found;"
1216        " perhaps install ffmpeg using 'apt install ffmpeg'."
1217    )
1218  return path
1219
1220
1221@typing.overload
1222def _run_ffmpeg(
1223    ffmpeg_args: Sequence[str],
1224    stdin: int | None = None,
1225    stdout: int | None = None,
1226    stderr: int | None = None,
1227    encoding: None = None,  # No encoding -> bytes
1228    allowed_input_files: Sequence[str] | None = None,
1229    allowed_output_files: Sequence[str] | None = None,
1230    sandbox_max_run_time_secs: int | None = None,
1231) -> subprocess.Popen[bytes]:
1232  ...
1233
1234
1235@typing.overload
1236def _run_ffmpeg(
1237    ffmpeg_args: Sequence[str],
1238    stdin: int | None = None,
1239    stdout: int | None = None,
1240    stderr: int | None = None,
1241    encoding: str = ...,  # Encoding -> str
1242    allowed_input_files: Sequence[str] | None = None,
1243    allowed_output_files: Sequence[str] | None = None,
1244    sandbox_max_run_time_secs: int | None = None,
1245) -> subprocess.Popen[str]:
1246  ...
1247
1248
1249def _run_ffmpeg(
1250    ffmpeg_args: Sequence[str],
1251    stdin: int | None = None,
1252    stdout: int | None = None,
1253    stderr: int | None = None,
1254    encoding: str | None = None,
1255    allowed_input_files: Sequence[str] | None = None,
1256    allowed_output_files: Sequence[str] | None = None,
1257    sandbox_max_run_time_secs: int | None = None,
1258) -> subprocess.Popen[bytes] | subprocess.Popen[str]:
1259  """Runs ffmpeg with the given args.
1260
1261  Args:
1262    ffmpeg_args: The args to pass to ffmpeg.
1263    stdin: Same as in `subprocess.Popen`.
1264    stdout: Same as in `subprocess.Popen`.
1265    stderr: Same as in `subprocess.Popen`.
1266    encoding: Same as in `subprocess.Popen`.
1267    allowed_input_files: The input files to allow for ffmpeg.
1268    allowed_output_files: The output files to allow for ffmpeg.
1269    sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
1270      If None, the default limit is 30 minutes.
1271
1272  Returns:
1273    The subprocess.Popen object with running ffmpeg process.
1274  """
1275  argv = []
1276  # In open source, keep env=None to preserve default behavior.
1277  # Context: https://github.com/google/mediapy/pull/62
1278  env: Any = None  # pylint: disable=unused-variable
1279  ffmpeg_path = _get_ffmpeg_path()
1280
1281  # Sandbox max runtime, allowed input and output files are not supported in
1282  # open source.
1283  del allowed_input_files
1284  del allowed_output_files
1285  del sandbox_max_run_time_secs
1286
1287  argv.append(ffmpeg_path)
1288  argv.extend(ffmpeg_args)
1289
1290  return subprocess.Popen(
1291      argv,
1292      stdin=stdin,
1293      stdout=stdout,
1294      stderr=stderr,
1295      encoding=encoding,
1296      env=env,
1297  )
1298
1299
1300def video_is_available() -> bool:
1301  """Returns True if the program `ffmpeg` is found.
1302
1303  See also `set_ffmpeg`.
1304  """
1305  return _search_for_ffmpeg_path() is not None
1306
1307
1308class VideoMetadata(NamedTuple):
1309  """Represents the data stored in a video container header.
1310
1311  Attributes:
1312    num_images: Number of frames that is expected from the video stream.  This
1313      is estimated from the framerate and the duration stored in the video
1314      header, so it might be inexact.  We set the value to -1 if number of
1315      frames is not found in the header.
1316    shape: The dimensions (height, width) of each video frame.
1317    fps: The framerate in frames per second.
1318    bps: The estimated bitrate of the video stream in bits per second, retrieved
1319      from the video header.
1320  """
1321
1322  num_images: int
1323  shape: tuple[int, int]
1324  fps: float
1325  bps: int | None
1326
1327
1328def _get_video_metadata(path: _Path) -> VideoMetadata:
1329  """Returns attributes of video stored in the specified local file."""
1330  if not pathlib.Path(path).is_file():
1331    raise RuntimeError(f"Video file '{path}' is not found.")
1332
1333  command = [
1334      '-nostdin',
1335      '-i',
1336      str(path),
1337      '-acodec',
1338      'copy',
1339      # Necessary to get "frame= *(\d+)" using newer ffmpeg versions.
1340      # Previously, was `'-vcodec', 'copy'`
1341      '-vf',
1342      'select=1',
1343      '-vsync',
1344      '0',
1345      '-f',
1346      'null',
1347      '-',
1348  ]
1349  with _run_ffmpeg(
1350      command,
1351      allowed_input_files=[str(path)],
1352      stderr=subprocess.PIPE,
1353      encoding='utf-8',
1354  ) as proc:
1355    _, err = proc.communicate()
1356  bps = fps = num_images = width = height = rotation = None
1357  before_output_info = True
1358  for line in err.split('\n'):
1359    if line.startswith('Output '):
1360      before_output_info = False
1361    if match := re.search(r', bitrate: *([\d.]+) kb/s', line):
1362      bps = int(match.group(1)) * 1000
1363    if matches := re.findall(r'frame= *(\d+) ', line):
1364      num_images = int(matches[-1])
1365    if 'Stream #0:' in line and ': Video:' in line and before_output_info:
1366      if not (match := re.search(r', (\d+)x(\d+)', line)):
1367        raise RuntimeError(f'Unable to parse video dimensions in line {line}')
1368      width, height = int(match.group(1)), int(match.group(2))
1369      # Some videos have a 'k' suffix for kiloframes per second. Thus an extra
1370      # `(k?)` is used for this case. This usually happens when we don't know
1371      # the exact framerate. However, in order to not raise an error here, we
1372      # will try to parse the framerate as x1000.
1373      if match := re.search(r', ([\d.]+)(k?) fps', line):
1374        number = float(match.group(1))
1375        fps = number * 1000 if match.group(2) == 'k' else number
1376      elif str(path).endswith('.gif'):
1377        # Some GIF files lack a framerate attribute; use a reasonable default.
1378        fps = 10
1379      else:
1380        raise RuntimeError(f'Unable to parse video framerate in line {line}')
1381    if match := re.fullmatch(r'\s*rotate\s*:\s*(\d+)', line):
1382      rotation = int(match.group(1))
1383    if match := re.fullmatch(r'.*rotation of (-?\d+).*\sdegrees', line):
1384      rotation = int(match.group(1))
1385  if not num_images:
1386    num_images = -1
1387  if not width:
1388    raise RuntimeError(f'Unable to parse video header: {err}')
1389  # By default, ffmpeg enables "-autorotate"; we just fix the dimensions.
1390  if rotation in (90, 270, -90, -270):
1391    width, height = height, width
1392  assert height is not None and width is not None
1393  shape = height, width
1394  assert fps is not None
1395  return VideoMetadata(num_images, shape, fps, bps)
1396
1397
1398class _VideoIO:
1399  """Base class for `VideoReader` and `VideoWriter`."""
1400
1401  def _get_pix_fmt(self, dtype: _DType, image_format: str) -> str:
1402    """Returns ffmpeg pix_fmt given data type and image format."""
1403    native_endian_suffix = {'little': 'le', 'big': 'be'}[sys.byteorder]
1404    pix_fmt_map: dict[type[np.generic], dict[str, str]] = {
1405        np.uint8: {
1406            'rgb': 'rgb24',
1407            'yuv': 'yuv444p',
1408            'gray': 'gray',
1409        },
1410        np.uint16: {
1411            'rgb': 'rgb48' + native_endian_suffix,
1412            'yuv': 'yuv444p16' + native_endian_suffix,
1413            'gray': 'gray16' + native_endian_suffix,
1414        },
1415    }
1416    scalar_type = cast(type[np.generic], dtype.type)
1417    return pix_fmt_map[scalar_type][image_format]
1418
1419
1420class VideoReader(_VideoIO):
1421  """Context to read a compressed video as an iterable over its images.
1422
1423  >>> with VideoReader('/tmp/river.mp4') as reader:
1424  ...   print(f'Video has {reader.num_images} images with shape={reader.shape},'
1425  ...         f' at {reader.fps} frames/sec and {reader.bps} bits/sec.')
1426  ...   for image in reader:
1427  ...     print(image.shape)
1428
1429  >>> with VideoReader('/tmp/river.mp4') as reader:
1430  ...   video = np.array(tuple(reader))
1431
1432  >>> url = 'https://github.com/hhoppe/data/raw/main/video.mp4'
1433  >>> with VideoReader(url) as reader:
1434  ...   show_video(reader)
1435
1436  Attributes:
1437    path_or_url: Location of input video.
1438    output_format: Format of output images (default 'rgb').  If 'rgb', each
1439      image has shape=(height, width, 3) with R, G, B values.  If 'yuv', each
1440      image has shape=(height, width, 3) with Y, U, V values.  If 'gray', each
1441      image has shape=(height, width).
1442    dtype: Data type for output images.  The default is `np.uint8`.  Use of
1443      `np.uint16` allows reading 10-bit or 12-bit data without precision loss.
1444    metadata: Object storing the information retrieved from the video header.
1445      Its attributes are copied as attributes in this class.
1446    num_images: Number of frames that is expected from the video stream.  This
1447      is estimated from the framerate and the duration stored in the video
1448      header, so it might be inexact.
1449    shape: The dimensions (height, width) of each video frame.
1450    fps: The framerate in frames per second.
1451    bps: The estimated bitrate of the video stream in bits per second, retrieved
1452      from the video header.
1453    stream_index: The stream index to read from. The default is 0.
1454    sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
1455      If None, the default limit is 30 minutes. Unused in open source.
1456  """
1457
1458  path_or_url: _Path
1459  output_format: str
1460  dtype: _DType
1461  metadata: VideoMetadata
1462  num_images: int
1463  shape: tuple[int, int]
1464  fps: float
1465  bps: int | None
1466  stream_index: int
1467  _num_bytes_per_image: int
1468
1469  def __init__(
1470      self,
1471      path_or_url: _Path,
1472      *,
1473      stream_index: int = 0,
1474      output_format: str = 'rgb',
1475      dtype: _DTypeLike = np.uint8,
1476      sandbox_max_run_time_secs: int | None = None,
1477  ):
1478    if output_format not in {'rgb', 'yuv', 'gray'}:
1479      raise ValueError(
1480          f'Output format {output_format} is not rgb, yuv, or gray.'
1481      )
1482    self.path_or_url = path_or_url
1483    self.output_format = output_format
1484    self.stream_index = stream_index
1485    self.dtype = np.dtype(dtype)
1486    if self.dtype.type not in (np.uint8, np.uint16):
1487      raise ValueError(f'Type {dtype} is not np.uint8 or np.uint16.')
1488    self.sandbox_max_run_time_secs = sandbox_max_run_time_secs
1489    self._read_via_local_file: Any = None
1490    self._popen: subprocess.Popen[bytes] | None = None
1491    self._proc: subprocess.Popen[bytes] | None = None
1492
1493  def __enter__(self) -> VideoReader:
1494    try:
1495      self._read_via_local_file = _read_via_local_file(self.path_or_url)
1496      # pylint: disable-next=no-member
1497      tmp_name = self._read_via_local_file.__enter__()
1498
1499      self.metadata = _get_video_metadata(tmp_name)
1500      self.num_images, self.shape, self.fps, self.bps = self.metadata
1501      pix_fmt = self._get_pix_fmt(self.dtype, self.output_format)
1502      num_channels = {'rgb': 3, 'yuv': 3, 'gray': 1}[self.output_format]
1503      bytes_per_channel = self.dtype.itemsize
1504      self._num_bytes_per_image = (
1505          math.prod(self.shape) * num_channels * bytes_per_channel
1506      )
1507
1508      command = [
1509          '-v',
1510          'panic',
1511          '-nostdin',
1512          '-i',
1513          tmp_name,
1514          '-vcodec',
1515          'rawvideo',
1516          '-f',
1517          'image2pipe',
1518          '-map',
1519          f'0:v:{self.stream_index}',
1520          '-pix_fmt',
1521          pix_fmt,
1522          '-vsync',
1523          'vfr',
1524          '-',
1525      ]
1526      self._popen = _run_ffmpeg(
1527          command,
1528          stdout=subprocess.PIPE,
1529          stderr=subprocess.PIPE,
1530          allowed_input_files=[tmp_name],
1531          sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
1532      )
1533      self._proc = self._popen.__enter__()
1534    except Exception:
1535      self.__exit__(None, None, None)
1536      raise
1537    return self
1538
1539  def __exit__(self, *_: Any) -> None:
1540    self.close()
1541
1542  def read(self) -> _NDArray | None:
1543    """Reads a video image frame (or None if at end of file).
1544
1545    Returns:
1546      A numpy array in the format specified by `output_format`, i.e., a 3D
1547      array with 3 color channels, except for format 'gray' which is 2D.
1548
1549    Raises:
1550      RuntimeError: If there is an error reading from the output file.
1551    """
1552    assert self._proc, 'Error: reading from an already closed context.'
1553    stdout = self._proc.stdout
1554    assert stdout is not None
1555    data = stdout.read(self._num_bytes_per_image)
1556    if not data:  # Due to either end-of-file or subprocess error.
1557      self.close()  # Raises exception if subprocess had error.
1558      return None  # To indicate end-of-file.
1559    if len(data) != self._num_bytes_per_image:
1560      self._proc.wait()
1561      stderr = self._proc.stderr
1562      stderr_output = ''
1563      if stderr is not None:
1564        stderr_output = stderr.read().decode('utf-8', errors='replace').strip()
1565      raise RuntimeError(
1566          f'ffmpeg exited with code {self._proc.returncode}.\nIncomplete'
1567          f' frame read: expected {self._num_bytes_per_image} bytes, but got'
1568          f' {len(data)}.\nffmpeg stderr:\n{stderr_output}'
1569      )
1570    image: _NDArray = np.frombuffer(data, dtype=self.dtype)
1571    if self.output_format == 'rgb':
1572      image = image.reshape(*self.shape, 3)
1573    elif self.output_format == 'yuv':  # Convert from planar YUV to pixel YUV.
1574      image = np.moveaxis(image.reshape(3, *self.shape), 0, 2)
1575    elif self.output_format == 'gray':  # Generate 2D rather than 3D ndimage.
1576      image = image.reshape(*self.shape)
1577    else:
1578      raise AssertionError
1579    return image
1580
1581  def __iter__(self) -> Iterator[_NDArray]:
1582    while True:
1583      image = self.read()
1584      if image is None:
1585        return
1586      yield image
1587
1588  def close(self) -> None:
1589    """Terminates video reader.  (Called automatically at end of context.)"""
1590    if self._popen:
1591      self._popen.__exit__(None, None, None)
1592      self._popen = None
1593      self._proc = None
1594    if self._read_via_local_file:
1595      # pylint: disable-next=no-member
1596      self._read_via_local_file.__exit__(None, None, None)
1597      self._read_via_local_file = None
1598
1599
1600class VideoWriter(_VideoIO):
1601  """Context to write a compressed video.
1602
1603  >>> shape = 480, 640
1604  >>> with VideoWriter('/tmp/v.mp4', shape, fps=60) as writer:
1605  ...   for image in moving_circle(shape, num_images=60):
1606  ...     writer.add_image(image)
1607  >>> show_video(read_video('/tmp/v.mp4'))
1608
1609
1610  Bitrate control may be specified using at most one of: `bps`, `qp`, or `crf`.
1611  If none are specified, `qp` is set to a default value.
1612  See https://slhck.info/video/2017/03/01/rate-control.html
1613
1614  If codec is 'gif', the args `bps`, `qp`, `crf`, and `encoded_format` are
1615  ignored.
1616
1617  Attributes:
1618    path: Output video.  Its suffix (e.g. '.mp4') determines the video container
1619      format.  The suffix must be '.gif' if the codec is 'gif'.
1620    shape: 2D spatial dimensions (height, width) of video image frames.  The
1621      dimensions must be even if 'encoded_format' has subsampled chroma (e.g.,
1622      'yuv420p' or 'yuv420p10le').
1623    codec: Compression algorithm as defined by "ffmpeg -codecs" (e.g., 'h264',
1624      'hevc', 'vp9', or 'gif').
1625    metadata: Optional VideoMetadata object whose `fps` and `bps` attributes are
1626      used if not specified as explicit parameters.
1627    fps: Frames-per-second framerate (default is 60.0 except 25.0 for 'gif').
1628    bps: Requested average bits-per-second bitrate (default None).
1629    qp: Quantization parameter for video compression quality (default None).
1630    crf: Constant rate factor for video compression quality (default None).
1631    ffmpeg_args: Additional arguments for `ffmpeg` command, e.g. '-g 30' to
1632      introduce I-frames, or '-bf 0' to omit B-frames.
1633    input_format: Format of input images (default 'rgb').  If 'rgb', each image
1634      has shape=(height, width, 3) or (height, width).  If 'yuv', each image has
1635      shape=(height, width, 3) with Y, U, V values.  If 'gray', each image has
1636      shape=(height, width).
1637    dtype: Expected data type for input images (any float input images are
1638      converted to `dtype`).  The default is `np.uint8`.  Use of `np.uint16` is
1639      necessary when encoding >8 bits/channel.
1640    encoded_format: Pixel format as defined by `ffmpeg -pix_fmts`, e.g.,
1641      'yuv420p' (2x2-subsampled chroma), 'yuv444p' (full-res chroma),
1642      'yuv420p10le' (10-bit per channel), etc.  The default (None) selects
1643      'yuv420p' if all shape dimensions are even, else 'yuv444p'.
1644    sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
1645      If None, the default limit is 30 minutes. Unused in open source.
1646    audio: Optional audio data as a NumPy array.  It should have shape (N,) for
1647      mono or (N, C) for multi-channel audio, where N is the number of samples
1648      and C is the number of channels.  The dtype must be one of `np.uint8`,
1649      `np.int16`, `np.int32`, `np.float32`, or `np.float64`, in native byte
1650      order.  Float samples are nominally in [-1.0, 1.0], and values outside
1651      this range typically clip; integer samples span the full range of their
1652      type, where `np.uint8` is unsigned with silence at 128.  The two streams
1653      are not truncated to a common length: if the audio is shorter than the
1654      video, the remainder is silent; if it is longer, the audio continues past
1655      the end of the video stream (most players keep showing the last frame).
1656    audio_sample_rate: Sample rate of the audio in Hz. Required if `audio` is
1657      provided.
1658    audio_codec: Audio compression algorithm as defined by "ffmpeg -codecs"
1659      (default 'aac').  It must be supported by the video container, e.g., 'aac'
1660      for MP4 ('h264' or 'hevc') or 'libopus' for WebM ('vp9').  Ignored if
1661      `audio` is None.
1662  """
1663
1664  def __init__(
1665      self,
1666      path: _Path,
1667      shape: tuple[int, int],
1668      *,
1669      codec: str = 'h264',
1670      metadata: VideoMetadata | None = None,
1671      fps: float | None = None,
1672      bps: int | None = None,
1673      qp: int | None = None,
1674      crf: float | None = None,
1675      ffmpeg_args: str | Sequence[str] = '',
1676      input_format: str = 'rgb',
1677      dtype: _DTypeLike = np.uint8,
1678      encoded_format: str | None = None,
1679      sandbox_max_run_time_secs: int | None = None,
1680      audio: _NDArray | None = None,
1681      audio_sample_rate: int | None = None,
1682      audio_codec: str = 'aac',
1683  ) -> None:
1684    _check_2d_shape(shape)
1685    if fps is None and metadata:
1686      fps = metadata.fps
1687    if fps is None:
1688      fps = 25.0 if codec == 'gif' else 60.0
1689    if fps <= 0.0:
1690      raise ValueError(f'Frame-per-second value {fps} is invalid.')
1691    if bps is None and metadata:
1692      bps = metadata.bps
1693    bps = bps if bps is not None else None
1694    if bps is not None and bps <= 0:
1695      raise ValueError(f'Bitrate value {bps} is invalid.')
1696    if qp is not None and (not isinstance(qp, int) or qp < 0):
1697      raise ValueError(
1698          f'Quantization parameter {qp} cannot be negative. It must be a'
1699          ' non-negative integer.'
1700      )
1701    num_rate_specifications = sum(x is not None for x in (bps, qp, crf))
1702    if num_rate_specifications > 1:
1703      raise ValueError(
1704          f'Must specify at most one of bps, qp, or crf ({bps}, {qp}, {crf}).'
1705      )
1706    ffmpeg_args = (
1707        shlex.split(ffmpeg_args)
1708        if isinstance(ffmpeg_args, str)
1709        else list(ffmpeg_args)
1710    )
1711    if input_format not in {'rgb', 'yuv', 'gray'}:
1712      raise ValueError(f'Input format {input_format} is not rgb, yuv, or gray.')
1713    dtype = np.dtype(dtype)
1714    if dtype.type not in (np.uint8, np.uint16):
1715      raise ValueError(f'Type {dtype} is not np.uint8 or np.uint16.')
1716    if audio is not None:
1717      if audio_sample_rate is None:
1718        raise ValueError('The audio_sample_rate must be set if audio is set.')
1719      if audio.ndim not in (1, 2):
1720        raise ValueError(f'Audio shape {audio.shape} is not (N,) or (N, C).')
1721      if audio.dtype.type not in _FFMPEG_AUDIO_FORMAT_FROM_DTYPE:
1722        raise ValueError(f'Audio type {audio.dtype} is unsupported.')
1723      if not audio.dtype.isnative:
1724        raise ValueError(
1725            f'Audio type {audio.dtype} is not in native byte order.'
1726        )
1727      if codec == 'gif':
1728        raise ValueError('Audio is not supported with the gif codec.')
1729    self.path = pathlib.Path(path)
1730    self.shape = shape
1731    all_dimensions_are_even = all(dim % 2 == 0 for dim in shape)
1732    if encoded_format is None:
1733      encoded_format = 'yuv420p' if all_dimensions_are_even else 'yuv444p'
1734    if not all_dimensions_are_even and encoded_format.startswith(
1735        ('yuv42', 'yuvj42')
1736    ):
1737      raise ValueError(
1738          f'With encoded_format {encoded_format}, video dimensions must be'
1739          f' even, but shape is {shape}.'
1740      )
1741    self.fps = fps
1742    self.codec = codec
1743    self.bps = bps
1744    self.qp = qp
1745    self.crf = crf
1746    self.ffmpeg_args = ffmpeg_args
1747    self.input_format = input_format
1748    self.dtype = dtype
1749    self.encoded_format = encoded_format
1750    self.sandbox_max_run_time_secs = sandbox_max_run_time_secs
1751    self.audio = audio
1752    self.audio_sample_rate = audio_sample_rate
1753    self.audio_codec = audio_codec
1754    if num_rate_specifications == 0 and not ffmpeg_args:
1755      qp = 20 if math.prod(self.shape) <= 640 * 480 else 28
1756    self._bitrate_args = (
1757        (['-vb', f'{bps}'] if bps is not None else [])
1758        + (['-qp', f'{qp}'] if qp is not None else [])
1759        + (['-vb', '0', '-crf', f'{crf}'] if crf is not None else [])
1760    )
1761    if self.codec == 'gif':
1762      if self.path.suffix != '.gif':
1763        raise ValueError(f"File '{self.path}' does not have a .gif suffix.")
1764      self.encoded_format = 'pal8'
1765      self._bitrate_args = []
1766      video_filter = 'split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse'
1767      # Less common (and likely less useful) is a per-frame color palette:
1768      # video_filter = ('split[s0][s1];[s0]palettegen=stats_mode=single[p];'
1769      #                 '[s1][p]paletteuse=new=1')
1770      self.ffmpeg_args = ['-vf', video_filter, '-f', 'gif'] + self.ffmpeg_args
1771    self._exit_stack: contextlib.ExitStack | None = None
1772    self._proc: subprocess.Popen[bytes] | None = None
1773
1774  def __enter__(self) -> VideoWriter:
1775    input_pix_fmt = self._get_pix_fmt(self.dtype, self.input_format)
1776    try:
1777      self._exit_stack = contextlib.ExitStack()
1778      tmp_name = self._exit_stack.enter_context(
1779          _write_via_local_file(self.path)
1780      )
1781
1782      # Writing to stdout using ('-f', 'mp4', '-') would require
1783      # ('-movflags', 'frag_keyframe+empty_moov') and the result is nonportable.
1784      height, width = self.shape
1785
1786      audio_input_args = []
1787      audio_output_args = ['-an']
1788      allowed_input_files = []
1789
1790      if self.audio is not None:
1791        audio_path = self._exit_stack.enter_context(
1792            _audio_via_local_file(self.audio)
1793        )
1794        channels = 1 if self.audio.ndim == 1 else self.audio.shape[1]
1795        audio_format = _FFMPEG_AUDIO_FORMAT_FROM_DTYPE[self.audio.dtype.type]
1796        if self.audio.dtype.itemsize > 1:
1797          audio_format += {'little': 'le', 'big': 'be'}[sys.byteorder]
1798        audio_input_args = [
1799            '-f',
1800            audio_format,
1801            '-ar',
1802            str(self.audio_sample_rate),
1803            '-ac',
1804            str(channels),
1805            '-i',
1806            audio_path,
1807        ]
1808        audio_output_args = ['-c:a', self.audio_codec]
1809        allowed_input_files.append(audio_path)
1810
1811      command = (
1812          [
1813              '-v',
1814              'error',
1815              '-f',
1816              'rawvideo',
1817              '-vcodec',
1818              'rawvideo',
1819              '-pix_fmt',
1820              input_pix_fmt,
1821              '-s',
1822              f'{width}x{height}',
1823              '-r',
1824              f'{self.fps}',
1825              '-i',
1826              '-',
1827          ]
1828          + audio_input_args
1829          + audio_output_args
1830          + [
1831              '-vcodec',
1832              self.codec,
1833              '-pix_fmt',
1834              self.encoded_format,
1835          ]
1836          + self._bitrate_args
1837          + self.ffmpeg_args
1838          + ['-y', tmp_name]
1839      )
1840      self._proc = self._exit_stack.enter_context(
1841          _run_ffmpeg(
1842              command,
1843              stdin=subprocess.PIPE,
1844              stderr=subprocess.PIPE,
1845              # `_run_ffmpeg` omits the sandbox flag only for None, so an empty
1846              # list would pass an empty '--sandbox_read_access_files'.
1847              allowed_input_files=allowed_input_files or None,
1848              allowed_output_files=[tmp_name],
1849              sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
1850          )
1851      )
1852    except Exception:
1853      self.__exit__(None, None, None)
1854      raise
1855    return self
1856
1857  def __exit__(self, *_: Any) -> None:
1858    self.close()
1859
1860  def add_image(self, image: _NDArray) -> None:
1861    """Writes a video frame.
1862
1863    Args:
1864      image: Array whose dtype and first two dimensions must match the `dtype`
1865        and `shape` specified in `VideoWriter` initialization.  If
1866        `input_format` is 'gray', the image must be 2D.  For the 'rgb'
1867        input_format, the image may be either 2D (interpreted as grayscale) or
1868        3D with three (R, G, B) channels.  For the 'yuv' input_format, the image
1869        must be 3D with three (Y, U, V) channels.
1870
1871    Raises:
1872      RuntimeError: If there is an error writing to the output file.
1873    """
1874    assert self._proc, 'Error: writing to an already closed context.'
1875    if issubclass(image.dtype.type, (np.floating, np.bool_)):
1876      image = to_type(image, self.dtype)
1877    if image.dtype != self.dtype:
1878      raise ValueError(f'Image type {image.dtype} != {self.dtype}.')
1879    if self.input_format == 'gray':
1880      if image.ndim != 2:
1881        raise ValueError(f'Image dimensions {image.shape} are not 2D.')
1882    else:
1883      if image.ndim == 2 and self.input_format == 'rgb':
1884        image = np.dstack((image, image, image))
1885      if not (image.ndim == 3 and image.shape[2] == 3):
1886        raise ValueError(f'Image dimensions {image.shape} are invalid.')
1887    if image.shape[:2] != self.shape:
1888      raise ValueError(
1889          f'Image dimensions {image.shape[:2]} do not match'
1890          f' those of the initialized video {self.shape}.'
1891      )
1892    if self.input_format == 'yuv':  # Convert from per-pixel YUV to planar YUV.
1893      image = np.moveaxis(image, 2, 0)
1894    data = image.tobytes()
1895    stdin = self._proc.stdin
1896    assert stdin is not None
1897    if stdin.write(data) != len(data):
1898      self._proc.wait()
1899      stderr = self._proc.stderr
1900      assert stderr is not None
1901      s = stderr.read().decode('utf-8')
1902      raise RuntimeError(f"Error writing '{self.path}': {s}")
1903
1904  def close(self) -> None:
1905    """Finishes writing the video.  (Called automatically at end of context.)"""
1906    if self._exit_stack is None:
1907      return
1908    # Unwinding the stack terminates the `ffmpeg` process, removes the
1909    # temporary audio file, and copies the encoded video to a remote `path`.
1910    # Raising within the `with` propagates the error into those contexts, so
1911    # an incomplete video is discarded rather than copied.
1912    with self._exit_stack:
1913      self._exit_stack = None
1914      if self._proc is not None:
1915        proc, self._proc = self._proc, None
1916        stdin = proc.stdin
1917        assert stdin is not None
1918        stdin.close()
1919        if proc.wait():
1920          stderr = proc.stderr
1921          assert stderr is not None
1922          s = stderr.read().decode('utf-8')
1923          raise RuntimeError(f"Error writing '{self.path}': {s}")
1924
1925
1926class _VideoArray(_BaseArray):
1927  """Wrapper to add a VideoMetadata `metadata` attribute to a numpy array."""
1928
1929  metadata: VideoMetadata | None
1930
1931  def __new__(
1932      cls: type[_VideoArray],
1933      input_array: _NDArray,
1934      metadata: VideoMetadata | None = None,
1935  ) -> _VideoArray:
1936    obj: _VideoArray = np.asarray(input_array).view(cls)
1937    obj.metadata = metadata
1938    return obj
1939
1940  def __array_finalize__(self, obj: Any) -> None:
1941    if obj is None:
1942      return
1943    self.metadata = getattr(obj, 'metadata', None)
1944
1945
1946def read_video(path_or_url: _Path, **kwargs: Any) -> _VideoArray:
1947  """Returns an array containing all images read from a compressed video file.
1948
1949  >>> video = read_video('/tmp/river.mp4')
1950  >>> print(f'The framerate is {video.metadata.fps} frames/s.')
1951  >>> show_video(video)
1952
1953  >>> url = 'https://github.com/hhoppe/data/raw/main/video.mp4'
1954  >>> show_video(read_video(url))
1955
1956  Args:
1957    path_or_url: Input video file.
1958    **kwargs: Additional parameters for `VideoReader`.
1959
1960  Returns:
1961    A 4D `numpy` array with dimensions (frame, height, width, channel), or a 3D
1962    array if `output_format` is specified as 'gray'.  The returned array has an
1963    attribute `metadata` containing `VideoMetadata` information.  This enables
1964    `show_video` to retrieve the framerate in `metadata.fps`.  Note that the
1965    metadata attribute is lost in most subsequent `numpy` operations.
1966  """
1967  with VideoReader(path_or_url, **kwargs) as reader:
1968    return _VideoArray(np.array(tuple(reader)), metadata=reader.metadata)
1969
1970
1971def write_video(path: _Path, images: Iterable[_NDArray], **kwargs: Any) -> None:
1972  """Writes images to a compressed video file.
1973
1974  >>> video = moving_circle((480, 640), num_images=60)
1975  >>> write_video('/tmp/v.mp4', video, fps=60, qp=18)
1976  >>> show_video(read_video('/tmp/v.mp4'))
1977
1978  Args:
1979    path: Output video file.
1980    images: Iterable over video frames, e.g. a 4D array or a list of 2D or 3D
1981      arrays.
1982    **kwargs: Additional parameters for `VideoWriter`.
1983  """
1984  first_image, images = _peek_first(images)
1985  shape = first_image.shape[0], first_image.shape[1]
1986  dtype = first_image.dtype
1987  if dtype == bool:
1988    dtype = np.dtype(np.uint8)
1989  elif np.issubdtype(dtype, np.floating):
1990    dtype = np.dtype(np.uint16)
1991  kwargs = {'metadata': getattr(images, 'metadata', None), **kwargs}
1992  with VideoWriter(path, shape=shape, dtype=dtype, **kwargs) as writer:
1993    for image in images:
1994      writer.add_image(image)
1995
1996
1997def compress_video(
1998    images: Iterable[_NDArray], *, codec: str = 'h264', **kwargs: Any
1999) -> bytes:
2000  """Returns a buffer containing a compressed video.
2001
2002  The video container is 'gif' for 'gif' codec, 'webm' for 'vp9' codec,
2003  and mp4 otherwise.
2004
2005  >>> video = read_video('/tmp/river.mp4')
2006  >>> data = compress_video(video, bps=10_000_000)
2007  >>> print(len(data))
2008
2009  >>> data = compress_video(moving_circle((100, 100), num_images=10), fps=10)
2010
2011  Args:
2012    images: Iterable over video frames.
2013    codec: Compression algorithm as defined by `ffmpeg -codecs` (e.g., 'h264',
2014      'hevc', 'vp9', or 'gif').
2015    **kwargs: Additional parameters for `VideoWriter`.
2016
2017  Returns:
2018    A bytes buffer containing the compressed video.
2019  """
2020  suffix = _filename_suffix_from_codec(codec)
2021  with tempfile.TemporaryDirectory() as directory_name:
2022    tmp_path = pathlib.Path(directory_name) / f'file{suffix}'
2023    write_video(tmp_path, images, codec=codec, **kwargs)
2024    return tmp_path.read_bytes()
2025
2026
2027def decompress_video(data: bytes, **kwargs: Any) -> _NDArray:
2028  """Returns video images from an MP4-compressed data buffer."""
2029  with tempfile.TemporaryDirectory() as directory_name:
2030    tmp_path = pathlib.Path(directory_name) / 'file.mp4'
2031    tmp_path.write_bytes(data)
2032    return read_video(tmp_path, **kwargs)
2033
2034
2035def html_from_compressed_video(
2036    data: bytes,
2037    width: int,
2038    height: int,
2039    *,
2040    title: str | None = None,
2041    border: bool | str = False,
2042    loop: bool = True,
2043    autoplay: bool = True,
2044) -> str:
2045  """Returns an HTML string with a video tag containing H264-encoded data.
2046
2047  Args:
2048    data: MP4-compressed video bytes.
2049    width: Width of HTML video in pixels.
2050    height: Height of HTML video in pixels.
2051    title: Optional text shown centered above the video.
2052    border: If `bool`, whether to place a black boundary around the image, or if
2053      `str`, the boundary CSS style.
2054    loop: If True, the playback repeats forever.
2055    autoplay: If True, video playback starts without having to click.
2056  """
2057  b64 = base64.b64encode(data).decode('utf-8')
2058  if isinstance(border, str):
2059    border = f'{border}; '
2060  elif border:
2061    border = 'border:1px solid black; '
2062  else:
2063    border = ''
2064  options = (
2065      f'controls width="{width}" height="{height}"'
2066      f' style="{border}object-fit:cover;"'
2067      f'{" loop" if loop else ""}'
2068      f'{" autoplay muted" if autoplay else ""}'
2069  )
2070  s = f"""<video {options}>
2071      <source src="data:video/mp4;base64,{b64}" type="video/mp4"/>
2072      This browser does not support the video tag.
2073      </video>"""
2074  if title is not None:
2075    s = f"""<div style="display:flex; align-items:left;">
2076      <div style="display:flex; flex-direction:column; align-items:center;">
2077      <div>{title}</div><div>{s}</div></div></div>"""
2078  return s
2079
2080
2081def show_video(
2082    images: Iterable[_NDArray], *, title: str | None = None, **kwargs: Any
2083) -> str | None:
2084  """Displays a video in the IPython notebook and optionally saves it to a file.
2085
2086  See `show_videos`.
2087
2088  >>> video = read_video('https://github.com/hhoppe/data/raw/main/video.mp4')
2089  >>> show_video(video, title='River video')
2090
2091  >>> show_video(moving_circle((80, 80), num_images=10), fps=5, border=True)
2092
2093  >>> show_video(read_video('/tmp/river.mp4'))
2094
2095  Args:
2096    images: Iterable of video frames (e.g., a 4D array or a list of 2D or 3D
2097      arrays).
2098    title: Optional text shown centered above the video.
2099    **kwargs: See `show_videos`.
2100
2101  Returns:
2102    html string if `return_html` is `True`.
2103  """
2104  return show_videos([images], [title], **kwargs)
2105
2106
2107def show_videos(
2108    videos: Iterable[Iterable[_NDArray]] | Mapping[str, Iterable[_NDArray]],
2109    titles: Iterable[str | None] | None = None,
2110    *,
2111    width: int | None = None,
2112    height: int | None = None,
2113    downsample: bool = True,
2114    columns: int | None = None,
2115    fps: float | None = None,
2116    bps: int | None = None,
2117    qp: int | None = None,
2118    codec: str = 'h264',
2119    ylabel: str = '',
2120    html_class: str = 'show_videos',
2121    return_html: bool = False,
2122    audios: Iterable[_NDArray] | Mapping[str, _NDArray] | None = None,
2123    audio_sample_rate: int | None = None,
2124    audio_codec: str = 'aac',
2125    **kwargs: Any,
2126) -> str | None:
2127  """Displays a row of videos in the IPython notebook.
2128
2129  Creates HTML with `<video>` tags containing embedded H264-encoded bytestrings.
2130  If `codec` is set to 'gif', we instead use `<img>` tags containing embedded
2131  GIF-encoded bytestrings.  Note that the resulting GIF animations skip frames
2132  when the `fps` period is not a multiple of 10 ms units (GIF frame delay
2133  units).  Encoding at `fps` = 20.0, 25.0, or 50.0 works fine.
2134
2135  If a directory has been specified using `set_show_save_dir`, also saves each
2136  titled video to a file in that directory based on its title.
2137
2138  Args:
2139    videos: Iterable of videos, or dictionary of `{title: video}`.  Each video
2140      must be an iterable of images.  If a video object has a `metadata`
2141      (`VideoMetadata`) attribute, its `fps` field provides a default framerate.
2142    titles: Optional strings shown above the corresponding videos.
2143    width: Optional, overrides displayed width (in pixels).
2144    height: Optional, overrides displayed height (in pixels).
2145    downsample: If True, each video whose width or height is greater than the
2146      specified `width` or `height` is resampled to the display resolution. This
2147      improves antialiasing and reduces the size of the notebook.
2148    columns: Optional, maximum number of videos per row.
2149    fps: Frames-per-second framerate (default is 60.0 except 25.0 for GIF).
2150    bps: Bits-per-second bitrate (default None).
2151    qp: Quantization parameter for video compression quality (default None).
2152    codec: Compression algorithm; must be either 'h264' or 'gif'.
2153    ylabel: Text (rotated by 90 degrees) shown on the left of each row.
2154    html_class: CSS class name used in definition of HTML element.
2155    return_html: If `True` return the raw HTML `str` instead of displaying.
2156    audios: Optional iterable of audio tracks, or dictionary of `{title:
2157      audio}`; see `VideoWriter`.  Each track is attached to the corresponding
2158      video, so an iterable must have one entry (possibly None) per video.  With
2159      a dictionary, its keys must match the video titles exactly (a value may be
2160      None for no audio).  Audio is incompatible with `codec` 'gif'.  With the
2161      default `autoplay=True`, the video starts muted (browsers block unmuted
2162      autoplay); unmute it in the player controls, or pass `autoplay=False`.
2163    audio_sample_rate: Sample rate of the audio in Hz, shared by all tracks.
2164      Required if `audios` is provided.
2165    audio_codec: Audio compression algorithm (default 'aac'); see `VideoWriter`.
2166    **kwargs: Additional parameters (`border`, `loop`, `autoplay`) for
2167      `html_from_compressed_video`.
2168
2169  Returns:
2170    html string if `return_html` is `True`.
2171  """
2172  list_titles: list[str | None]
2173  if isinstance(videos, Mapping):
2174    if titles is not None:
2175      raise ValueError(
2176          'Cannot have both a video dictionary and a titles parameter.'
2177      )
2178    list_videos = list(videos.values())
2179    list_titles = cast(list[str | None], list(videos.keys()))
2180  else:
2181    list_videos = list(cast('Iterable[_NDArray]', videos))
2182    list_titles = [None] * len(list_videos) if titles is None else list(titles)
2183    if len(list_videos) != len(list_titles):
2184      raise ValueError(
2185          'Number of videos does not match number of titles'
2186          f' ({len(list_videos)} vs {len(list_titles)}).'
2187      )
2188
2189  list_audios: list[_NDArray | None]
2190  if audios is None:
2191    list_audios = [None] * len(list_videos)
2192  elif isinstance(audios, Mapping):
2193    audio_of_title = cast('Mapping[str, _NDArray]', audios)
2194    missing = set(list_titles).difference(audio_of_title)
2195    extra = set(audio_of_title).difference(list_titles)
2196    if missing or extra:
2197      raise ValueError(
2198          'The audios dictionary keys must match the video titles (use None as'
2199          f' the value for no audio); missing: {sorted(missing, key=str)},'
2200          f' extra: {sorted(extra, key=str)}.'
2201      )
2202    # The check above ensures that each title is a key, so none is None.
2203    list_audios = [
2204        audio_of_title[title] for title in list_titles if title is not None
2205    ]
2206  else:
2207    list_audios = list(audios)
2208
2209  if len(list_videos) != len(list_audios):
2210    raise ValueError(
2211        'Number of videos does not match number of audio'
2212        f' ({len(list_videos)} vs {len(list_audios)}).'
2213    )
2214
2215  if codec not in {'h264', 'gif'}:
2216    raise ValueError(f'Codec {codec} is neither h264 or gif.')
2217
2218  html_strings = []
2219  for video, title, video_audio in zip(list_videos, list_titles, list_audios):
2220    metadata: VideoMetadata | None = getattr(video, 'metadata', None)
2221    first_image, video = _peek_first(video)
2222    first_h, first_w = first_image.shape[:2]
2223    w, h = _get_width_height(width, height, (first_h, first_w))
2224    if downsample and (w < first_w or h < first_h):
2225      # Not resize_video() because each image may have different depth and type.
2226      video = [resize_image(image, (h, w)) for image in video]
2227      first_image = video[0]
2228      first_h, first_w = first_image.shape[:2]
2229    data = compress_video(
2230        video,
2231        metadata=metadata,
2232        fps=fps,
2233        bps=bps,
2234        qp=qp,
2235        codec=codec,
2236        audio=video_audio,
2237        audio_sample_rate=audio_sample_rate,
2238        audio_codec=audio_codec,
2239    )
2240    if title is not None and _config.show_save_dir:
2241      suffix = _filename_suffix_from_codec(codec)
2242      path = pathlib.Path(_config.show_save_dir) / f'{title}{suffix}'
2243      with _open(path, mode='wb') as f:
2244        f.write(data)
2245    if codec == 'gif':
2246      pixelated = h > first_h or w > first_w
2247      html_string = html_from_compressed_image(
2248          data, w, h, title=title, fmt='gif', pixelated=pixelated, **kwargs
2249      )
2250    else:
2251      html_string = html_from_compressed_video(
2252          data, w, h, title=title, **kwargs
2253      )
2254    html_strings.append(html_string)
2255
2256  # Create single-row tables each with no more than 'columns' elements.
2257  table_strings = []
2258  for row_html_strings in _chunked(html_strings, columns):
2259    td = '<td style="padding:1px;">'
2260    s = ''.join(f'{td}{e}</td>' for e in row_html_strings)
2261    if ylabel:
2262      style = 'writing-mode:vertical-lr; transform:rotate(180deg);'
2263      s = f'{td}<span style="{style}">{ylabel}</span></td>' + s
2264    table_strings.append(
2265        f'<table class="{html_class}"'
2266        f' style="border-spacing:0px;"><tr>{s}</tr></table>'
2267    )
2268  s = ''.join(table_strings)
2269  if return_html:
2270    return s
2271  _display_html(s)
2272  return None
2273
2274
2275# Local Variables:
2276# fill-column: 80
2277# End: