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