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# Raw PCM sample formats (as in `ffmpeg -formats`) for the audio dtypes that
 195# `VideoWriter` is able to feed to `ffmpeg`.  Multi-byte formats lack the
 196# byte-order suffix ('le' or 'be'), which is appended for the native order.
 197_FFMPEG_AUDIO_FORMAT_FROM_DTYPE = {
 198    np.uint8: 'u8',
 199    np.int16: 's16',
 200    np.int32: 's32',
 201    np.float32: 'f32',
 202    np.float64: 'f64',
 203}
 204
 205_IMAGE_COMPARISON_HTML = """\
 206<script
 207  defer
 208  src="https://unpkg.com/img-comparison-slider@7/dist/index.js"
 209></script>
 210<link
 211  rel="stylesheet"
 212  href="https://unpkg.com/img-comparison-slider@7/dist/styles.css"
 213/>
 214
 215<img-comparison-slider>
 216  <img slot="first" src="data:image/png;base64,{b64_1}" />
 217  <img slot="second" src="data:image/png;base64,{b64_2}" />
 218</img-comparison-slider>
 219"""
 220
 221# ** Miscellaneous.
 222
 223
 224class _Config:
 225  ffmpeg_name_or_path: _Path = 'ffmpeg'
 226  show_save_dir: _Path | None = None
 227
 228
 229_config = _Config()
 230
 231
 232def _open(path: _Path, *args: Any, **kwargs: Any) -> Any:
 233  """Opens the file; this is a hook for the built-in `open()`."""
 234  return open(path, *args, **kwargs)
 235
 236
 237def _path_is_local(path: _Path) -> bool:
 238  """Returns True if the path is in the filesystem accessible by `ffmpeg`."""
 239  del path
 240  return True
 241
 242
 243def _search_for_ffmpeg_path() -> str | None:
 244  """Returns a path to the ffmpeg program, or None if not found."""
 245  if filename := shutil.which(_config.ffmpeg_name_or_path):
 246    return str(filename)
 247  return None
 248
 249
 250def _print_err(*args: str, **kwargs: Any) -> None:
 251  """Prints arguments to stderr immediately."""
 252  kwargs = {**dict(file=sys.stderr, flush=True), **kwargs}
 253  print(*args, **kwargs)
 254
 255
 256def _chunked(
 257    iterable: Iterable[_T], n: int | None = None
 258) -> Iterator[tuple[_T, ...]]:
 259  """Returns elements collected as tuples of length at most `n` if not None."""
 260
 261  def take(n: int | None, iterable: Iterable[_T]) -> tuple[_T, ...]:
 262    return tuple(itertools.islice(iterable, n))
 263
 264  return iter(functools.partial(take, n, iter(iterable)), ())
 265
 266
 267def _peek_first(iterator: Iterable[_T]) -> tuple[_T, Iterable[_T]]:
 268  """Given an iterator, returns first element and re-initialized iterator.
 269
 270  >>> first_image, images = _peek_first(moving_circle())
 271
 272  Args:
 273    iterator: An input iterator or iterable.
 274
 275  Returns:
 276    A tuple (first_element, iterator_reinitialized) containing:
 277      first_element: The first element of the input.
 278      iterator_reinitialized: A clone of the original iterator/iterable.
 279  """
 280  # Inspired from https://stackoverflow.com/a/12059829/1190077
 281  peeker, iterator_reinitialized = itertools.tee(iterator)
 282  first = next(peeker)
 283  return first, iterator_reinitialized
 284
 285
 286def _check_2d_shape(shape: tuple[int, int]) -> None:
 287  """Checks that `shape` is of the form (height, width) with two integers."""
 288  if len(shape) != 2:
 289    raise ValueError(f'Shape {shape} is not of the form (height, width).')
 290  if not all(isinstance(i, numbers.Integral) for i in shape):
 291    raise ValueError(f'Shape {shape} contains non-integers.')
 292
 293
 294def _run(args: str | Sequence[str]) -> None:
 295  """Executes command, printing output from stdout and stderr.
 296
 297  Args:
 298    args: Command to execute, which can be either a string or a sequence of word
 299      strings, as in `subprocess.run()`.  If `args` is a string, the shell is
 300      invoked to interpret it.
 301
 302  Raises:
 303    RuntimeError: If the command's exit code is nonzero.
 304  """
 305  proc = subprocess.run(
 306      args,
 307      shell=isinstance(args, str),
 308      stdout=subprocess.PIPE,
 309      stderr=subprocess.STDOUT,
 310      check=False,
 311      universal_newlines=True,
 312  )
 313  print(proc.stdout, end='', flush=True)
 314  if proc.returncode:
 315    raise RuntimeError(
 316        f"Command '{proc.args}' failed with code {proc.returncode}."
 317    )
 318
 319
 320def _display_html(text: str, /) -> None:
 321  """In a Jupyter notebook, display the HTML `text`."""
 322  IPython.display.display(IPython.display.HTML(text))  # type: ignore
 323
 324
 325def set_ffmpeg(name_or_path: _Path) -> None:
 326  """Specifies the name or path for the `ffmpeg` external program.
 327
 328  The `ffmpeg` program is required for compressing and decompressing video.
 329  (It is used in `read_video`, `write_video`, `show_video`, `show_videos`,
 330  etc.)
 331
 332  Args:
 333    name_or_path: Either a filename within a directory of `os.environ['PATH']`
 334      or a filepath.  The default setting is 'ffmpeg'.
 335  """
 336  _config.ffmpeg_name_or_path = name_or_path
 337
 338
 339def set_output_height(num_pixels: int) -> None:
 340  """Overrides the 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 = f'google.colab.output.setIframeHeight("{num_pixels}px")'
 345    output.eval_js(s)
 346  except (ModuleNotFoundError, AttributeError):
 347    pass
 348
 349
 350def set_max_output_height(num_pixels: int) -> None:
 351  """Sets the maximum height of the current output cell, if using Colab."""
 352  try:
 353    # We want to fail gracefully for non-Colab IPython notebooks.
 354    output = importlib.import_module('google.colab.output')
 355    s = (
 356        'google.colab.output.setIframeHeight('
 357        f'0, true, {{maxHeight: {num_pixels}}})'
 358    )
 359    output.eval_js(s)
 360  except (ModuleNotFoundError, AttributeError):
 361    pass
 362
 363
 364# ** Type conversions.
 365
 366
 367def _as_valid_media_type(dtype: _DTypeLike) -> _DType:
 368  """Returns validated media data type."""
 369  dtype = np.dtype(dtype)
 370  if not issubclass(dtype.type, (np.unsignedinteger, np.floating)):
 371    raise ValueError(
 372        f'Type {dtype} is not a valid media data type (uint or float).'
 373    )
 374  return dtype
 375
 376
 377def _as_valid_media_array(x: _ArrayLike) -> _NDArray:
 378  """Converts to ndarray (if not already), and checks validity of data type."""
 379  a = np.asarray(x)
 380  if a.dtype == bool:
 381    a = a.astype(np.uint8) * np.iinfo(np.uint8).max
 382  _as_valid_media_type(a.dtype)
 383  return a
 384
 385
 386def to_type(array: _ArrayLike, dtype: _DTypeLike) -> _NDArray:
 387  """Returns media array converted to specified type.
 388
 389  A "media array" is one in which the dtype is either a floating-point type
 390  (np.float32 or np.float64) or an unsigned integer type.  The array values are
 391  assumed to lie in the range [0.0, 1.0] for floating-point values, and in the
 392  full range for unsigned integers, e.g. [0, 255] for np.uint8.
 393
 394  Conversion between integers and floats maps uint(0) to 0.0 and uint(MAX) to
 395  1.0.  The input array may also be of type bool, whereby True maps to
 396  uint(MAX) or 1.0.  The values are scaled and clamped as appropriate during
 397  type conversions.
 398
 399  Args:
 400    array: Input array-like object (floating-point, unsigned int, or bool).
 401    dtype: Desired output type (floating-point or unsigned int).
 402
 403  Returns:
 404    Array `a` if it is already of the specified dtype, else a converted array.
 405  """
 406  a = np.asarray(array)
 407  dtype = np.dtype(dtype)
 408  del array
 409  if a.dtype != bool:
 410    _as_valid_media_type(a.dtype)  # Verify that 'a' has a valid dtype.
 411  if a.dtype == bool:
 412    result = a.astype(dtype)
 413    if np.issubdtype(dtype, np.unsignedinteger):
 414      result = result * dtype.type(np.iinfo(dtype).max)  # pyrefly: ignore[no-matching-overload]
 415  elif a.dtype == dtype:
 416    result = a
 417  elif np.issubdtype(dtype, 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(dtype).max  # pyrefly: ignore[no-matching-overload]
 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(dtype)
 427    elif dst_max <= np.iinfo(np.uint32).max:
 428      result = (a.astype(np.float64) * (dst_max / src_max) + 0.5).astype(dtype)
 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(dtype)
 436      dst[values_too_large] = dst_max
 437      result = dst if a.ndim > 0 else dst[0]
 438  else:
 439    assert np.issubdtype(dtype, np.floating)
 440    result = a.astype(dtype)
 441    if np.issubdtype(a.dtype, np.unsignedinteger):
 442      result = result / dtype.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    return np.array(
 650        _pil_image(image).resize(
 651            shape[::-1], resample=PIL.Image.Resampling.LANCZOS
 652        ),
 653        dtype=image.dtype,
 654    )
 655  if image.ndim == 2:
 656    # We convert to floating-point for resizing and convert back.
 657    return to_type(resize_image(to_float01(image), shape), image.dtype)
 658  # We resize each image channel individually.
 659  return np.dstack(
 660      [resize_image(channel, shape) for channel in np.moveaxis(image, -1, 0)]
 661  )
 662
 663
 664# ** Video processing.
 665
 666
 667def resize_video(video: Iterable[_NDArray], shape: tuple[int, int]) -> _NDArray:
 668  """Resizes `video` to specified spatial dimensions using a Lanczos filter.
 669
 670  Args:
 671    video: Iterable of images.
 672    shape: 2D spatial dimensions (height, width) of output video.
 673
 674  Returns:
 675    A resampled video whose spatial dimensions match `shape`.
 676  """
 677  _check_2d_shape(shape)
 678  return np.array([resize_image(image, shape) for image in video])
 679
 680
 681# ** General I/O.
 682
 683
 684def _is_url(path_or_url: _Path) -> bool:
 685  return isinstance(path_or_url, str) and path_or_url.startswith(
 686      ('http://', 'https://', 'file://')
 687  )
 688
 689
 690def read_contents(path_or_url: _Path) -> bytes:
 691  """Returns the contents of the file specified by either a path or URL."""
 692  data: bytes
 693  if _is_url(path_or_url):
 694    assert isinstance(path_or_url, str)
 695    headers = {'User-Agent': 'Chrome'}
 696    request = urllib.request.Request(path_or_url, headers=headers)
 697    with urllib.request.urlopen(request) as response:
 698      data = response.read()
 699  else:
 700    with _open(path_or_url, 'rb') as f:
 701      data = f.read()
 702  return data
 703
 704
 705@contextlib.contextmanager
 706def _read_via_local_file(path_or_url: _Path) -> Iterator[str]:
 707  """Context to copy a remote file locally to read from it.
 708
 709  Args:
 710    path_or_url: File, which may be remote.
 711
 712  Yields:
 713    The name of a local file which may be a copy of a remote file.
 714  """
 715  if _is_url(path_or_url) or not _path_is_local(path_or_url):
 716    suffix = pathlib.Path(path_or_url).suffix
 717    with tempfile.TemporaryDirectory() as directory_name:
 718      tmp_path = pathlib.Path(directory_name) / f'file{suffix}'
 719      tmp_path.write_bytes(read_contents(path_or_url))
 720      yield str(tmp_path)
 721  else:
 722    yield str(path_or_url)
 723
 724
 725@contextlib.contextmanager
 726def _write_via_local_file(path: _Path) -> Iterator[str]:
 727  """Context to write a temporary local file and subsequently copy it remotely.
 728
 729  Args:
 730    path: File, which may be remote.
 731
 732  Yields:
 733    The name of a local file which may be subsequently copied remotely.
 734  """
 735  if _path_is_local(path):
 736    yield str(path)
 737  else:
 738    suffix = pathlib.Path(path).suffix
 739    with tempfile.TemporaryDirectory() as directory_name:
 740      tmp_path = pathlib.Path(directory_name) / f'file{suffix}'
 741      yield str(tmp_path)
 742      with _open(path, mode='wb') as f:
 743        f.write(tmp_path.read_bytes())
 744
 745
 746@contextlib.contextmanager
 747def _audio_via_local_file(audio: _NDArray) -> Iterator[str]:
 748  """Context to write audio samples to a temporary local file.
 749
 750  Args:
 751    audio: Array of raw PCM audio samples, with shape (N,) for mono or (N, C)
 752      for C interleaved channels.
 753
 754  Yields:
 755    The name of a local file containing the raw samples.
 756  """
 757  with tempfile.TemporaryDirectory() as directory_name:
 758    tmp_path = pathlib.Path(directory_name) / 'audio.raw'
 759    tmp_path.write_bytes(audio.tobytes())
 760    yield str(tmp_path)
 761
 762
 763class set_show_save_dir:  # pylint: disable=invalid-name
 764  """Save all titled output from `show_*()` calls into files.
 765
 766  If the specified `directory` is not None, all titled images and videos
 767  displayed by `show_image`, `show_images`, `show_video`, and `show_videos` are
 768  also saved as files within the directory.
 769
 770  It can be used either to set the state or as a context manager:
 771
 772  >>> set_show_save_dir('/tmp')
 773  >>> show_image(color_ramp(), title='image1')  # Creates /tmp/image1.png.
 774  >>> show_video(moving_circle(), title='video2')  # Creates /tmp/video2.mp4.
 775  >>> set_show_save_dir(None)
 776
 777  >>> with set_show_save_dir('/tmp'):
 778  ...   show_image(color_ramp(), title='image1')  # Creates /tmp/image1.png.
 779  ...   show_video(moving_circle(), title='video2')  # Creates /tmp/video2.mp4.
 780  """
 781
 782  def __init__(self, directory: _Path | None):
 783    self._old_show_save_dir = _config.show_save_dir
 784    _config.show_save_dir = directory
 785
 786  def __enter__(self) -> None:
 787    pass
 788
 789  def __exit__(self, *_: Any) -> None:
 790    _config.show_save_dir = self._old_show_save_dir
 791
 792
 793# ** Image I/O.
 794
 795
 796def read_image(
 797    path_or_url: _Path,
 798    *,
 799    apply_exif_transpose: bool = True,
 800    dtype: _DTypeLike = None,  # pyrefly: ignore[bad-function-definition]
 801) -> _NDArray:
 802  """Returns an image read from a file path or URL.
 803
 804  Decoding is performed using `PIL`, which supports `uint8` images with 1, 3,
 805  or 4 channels and `uint16` images with a single channel.
 806
 807  Args:
 808    path_or_url: Path of input file.
 809    apply_exif_transpose: If True, rotate image according to EXIF orientation.
 810    dtype: Data type of the returned array.  If None, `np.uint8` or `np.uint16`
 811      is inferred automatically.
 812  """
 813  data = read_contents(path_or_url)
 814  return decompress_image(data, dtype, apply_exif_transpose)
 815
 816
 817def write_image(
 818    path: _Path, image: _ArrayLike, fmt: str = 'png', **kwargs: Any
 819) -> None:
 820  """Writes an image to a file.
 821
 822  Encoding is performed using `PIL`, which supports `uint8` images with 1, 3,
 823  or 4 channels and `uint16` images with a single channel.
 824
 825  File format is explicitly provided by `fmt` and not inferred by `path`.
 826
 827  Args:
 828    path: Path of output file.
 829    image: Array-like object.  If its type is float, it is converted to np.uint8
 830      using `to_uint8` (thus clamping to the input to the range [0.0, 1.0]).
 831      Otherwise it must be np.uint8 or np.uint16.
 832    fmt: Desired compression encoding, e.g. 'png'.
 833    **kwargs: Additional parameters for `PIL.Image.save()`.
 834  """
 835  image = _as_valid_media_array(image)
 836  if np.issubdtype(image.dtype, np.floating):
 837    image = to_uint8(image)
 838  with _open(path, 'wb') as f:
 839    _pil_image(image).save(f, format=fmt, **kwargs)
 840
 841
 842def to_rgb(
 843    array: _ArrayLike,
 844    *,
 845    vmin: float | None = None,
 846    vmax: float | None = None,
 847    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
 848) -> _NDArray:
 849  """Maps scalar values to RGB using value bounds and a color map.
 850
 851  Args:
 852    array: Scalar values, with arbitrary shape.
 853    vmin: Explicit min value for remapping; if None, it is obtained as the
 854      minimum finite value of `array`.
 855    vmax: Explicit max value for remapping; if None, it is obtained as the
 856      maximum finite value of `array`.
 857    cmap: A `pyplot` color map or callable, to map from 1D value to 3D or 4D
 858      color.
 859
 860  Returns:
 861    A new array in which each element is affinely mapped from [vmin, vmax]
 862    to [0.0, 1.0] and then color-mapped.
 863  """
 864  a = _as_valid_media_array(array)
 865  del array
 866  # For future numpy version 1.7.0:
 867  # vmin = np.amin(a, where=np.isfinite(a)) if vmin is None else vmin
 868  # vmax = np.amax(a, where=np.isfinite(a)) if vmax is None else vmax
 869  vmin = np.amin(np.where(np.isfinite(a), a, np.inf)) if vmin is None else vmin
 870  vmax = np.amax(np.where(np.isfinite(a), a, -np.inf)) if vmax is None else vmax
 871  a = (a.astype('float') - vmin) / (vmax - vmin + np.finfo(float).eps)
 872  if isinstance(cmap, str):
 873    if hasattr(matplotlib, 'colormaps'):
 874      rgb_from_scalar: Any = matplotlib.colormaps[cmap]  # Newer version.
 875    else:
 876      rgb_from_scalar = matplotlib.pyplot.cm.get_cmap(cmap)  # pylint: disable=no-member
 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, dtype: _DTypeLike = None, apply_exif_transpose: bool = True  # pyrefly: ignore[bad-function-definition]
 905) -> _NDArray:
 906  """Returns an image from a compressed data buffer.
 907
 908  Decoding is performed using `PIL`, which supports `uint8` images with 1, 3,
 909  or 4 channels and `uint16` images with a single channel.
 910
 911  Args:
 912    data: Buffer containing compressed image.
 913    dtype: Data type of the returned array.  If None, `np.uint8` or `np.uint16`
 914      is inferred automatically.
 915    apply_exif_transpose: If True, rotate image according to EXIF orientation.
 916  """
 917  pil_image: PIL.Image.Image = PIL.Image.open(io.BytesIO(data))
 918  if apply_exif_transpose:
 919    tmp_image = PIL.ImageOps.exif_transpose(pil_image)  # Future: in_place=True.
 920    assert tmp_image
 921    pil_image = tmp_image
 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]  # pyrefly: ignore[bad-return]
 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  >>> show_image(np.random.rand(100, 100))
1012  >>> show_image(np.random.randint(0, 256, size=(80, 80, 3), dtype='uint8'))
1013  >>> show_image(np.random.rand(10, 10) - 0.5, cmap='bwr', height=100)
1014  >>> show_image(read_image('/tmp/image.png'))
1015  >>> url = 'https://github.com/hhoppe/data/raw/main/image.png'
1016  >>> show_image(read_image(url))
1017
1018  Args:
1019    image: 2D array-like, or 3D array-like with 1, 3, or 4 channels.
1020    title: Optional text shown centered above the image.
1021    **kwargs: See `show_images`.
1022
1023  Returns:
1024    html string if `return_html` is `True`.
1025  """
1026  return show_images([np.asarray(image)], [title], **kwargs)
1027
1028
1029def show_images(
1030    images: Iterable[_ArrayLike] | Mapping[str, _ArrayLike],
1031    titles: Iterable[str | None] | None = None,
1032    *,
1033    width: int | None = None,
1034    height: int | None = None,
1035    downsample: bool = True,
1036    columns: int | None = None,
1037    vmin: float | None = None,
1038    vmax: float | None = None,
1039    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
1040    border: bool | str = False,
1041    ylabel: str = '',
1042    html_class: str = 'show_images',
1043    pixelated: bool | None = None,
1044    return_html: bool = False,
1045) -> str | None:
1046  """Displays a row of images in the IPython/Jupyter notebook.
1047
1048  If a directory has been specified using `set_show_save_dir`, also saves each
1049  titled image to a file in that directory based on its title.
1050
1051  >>> image1, image2 = np.random.rand(64, 64, 3), color_ramp((64, 64))
1052  >>> show_images([image1, image2])
1053  >>> show_images({'random image': image1, 'color ramp': image2}, height=128)
1054  >>> show_images([image1, image2] * 5, columns=4, border=True)
1055
1056  Args:
1057    images: Iterable of images, or dictionary of `{title: image}`.  Each image
1058      must be either a 2D array or a 3D array with 1, 3, or 4 channels.
1059    titles: Optional strings shown above the corresponding images.
1060    width: Optional, overrides displayed width (in pixels).
1061    height: Optional, overrides displayed height (in pixels).
1062    downsample: If True, each image whose width or height is greater than the
1063      specified `width` or `height` is resampled to the display resolution. This
1064      improves antialiasing and reduces the size of the notebook.
1065    columns: Optional, maximum number of images per row.
1066    vmin: For single-channel image, explicit min value for display.
1067    vmax: For single-channel image, explicit max value for display.
1068    cmap: For single-channel image, `pyplot` color map or callable to map 1D to
1069      3D color.
1070    border: If `bool`, whether to place a black boundary around the image, or if
1071      `str`, the boundary CSS style.
1072    ylabel: Text (rotated by 90 degrees) shown on the left of each row.
1073    html_class: CSS class name used in definition of HTML element.
1074    pixelated: If True, sets the CSS style to 'image-rendering: pixelated;'; if
1075      False, sets 'image-rendering: auto'; if None, uses pixelated rendering
1076      only on images for which `width` or `height` introduces magnification.
1077    return_html: If `True` return the raw HTML `str` instead of displaying.
1078
1079  Returns:
1080    html string if `return_html` is `True`.
1081  """
1082  if isinstance(images, Mapping):
1083    if titles is not None:
1084      raise ValueError('Cannot have images dictionary and titles parameter.')
1085    list_titles, list_images = list(images.keys()), list(images.values())
1086  else:
1087    list_images = list(images)
1088    list_titles = [None] * len(list_images) if titles is None else list(titles)
1089    if len(list_images) != len(list_titles):
1090      raise ValueError(
1091          'Number of images does not match number of titles'
1092          f' ({len(list_images)} vs {len(list_titles)}).'
1093      )
1094
1095  list_images = [
1096      _ensure_mapped_to_rgb(image, vmin=vmin, vmax=vmax, cmap=cmap)
1097      for image in list_images
1098  ]
1099
1100  def maybe_downsample(image: _NDArray) -> _NDArray:
1101    shape = image.shape[0], image.shape[1]
1102    w, h = _get_width_height(width, height, shape)
1103    if w < shape[1] or h < shape[0]:
1104      image = resize_image(image, (h, w))
1105    return image
1106
1107  if downsample:
1108    list_images = [maybe_downsample(image) for image in list_images]
1109  png_datas = [compress_image(to_uint8(image)) for image in list_images]
1110
1111  for title, png_data in zip(list_titles, png_datas):
1112    if title is not None and _config.show_save_dir:
1113      path = pathlib.Path(_config.show_save_dir) / f'{title}.png'
1114      with _open(path, mode='wb') as f:
1115        f.write(png_data)
1116
1117  def html_from_compressed_images() -> str:
1118    html_strings = []
1119    for image, title, png_data in zip(list_images, list_titles, png_datas):
1120      w, h = _get_width_height(width, height, image.shape[:2])  # pyrefly: ignore[missing-attribute]
1121      magnified = h > image.shape[0] or w > image.shape[1]  # pyrefly: ignore[missing-attribute]
1122      pixelated2 = pixelated if pixelated is not None else magnified
1123      html_strings.append(
1124          html_from_compressed_image(
1125              png_data, w, h, title=title, border=border, pixelated=pixelated2  # pyrefly: ignore[bad-argument-type]
1126          )
1127      )
1128    # Create single-row tables each with no more than 'columns' elements.
1129    table_strings = []
1130    for row_html_strings in _chunked(html_strings, columns):
1131      td = '<td style="padding:1px;">'
1132      s = ''.join(f'{td}{e}</td>' for e in row_html_strings)
1133      if ylabel:
1134        style = 'writing-mode:vertical-lr; transform:rotate(180deg);'
1135        s = f'{td}<span style="{style}">{ylabel}</span></td>' + s
1136      table_strings.append(
1137          f'<table class="{html_class}"'
1138          f' style="border-spacing:0px;"><tr>{s}</tr></table>'
1139      )
1140    return ''.join(table_strings)
1141
1142  s = html_from_compressed_images()
1143  while len(s) > _IPYTHON_HTML_SIZE_LIMIT * 0.5:
1144    warnings.warn('mediapy: subsampling images to reduce HTML size')
1145    list_images = [image[::2, ::2] for image in list_images]
1146    png_datas = [compress_image(to_uint8(image)) for image in list_images]
1147    s = html_from_compressed_images()
1148  if return_html:
1149    return s
1150  _display_html(s)
1151  return None
1152
1153
1154def compare_images(
1155    images: Iterable[_ArrayLike],
1156    *,
1157    vmin: float | None = None,
1158    vmax: float | None = None,
1159    cmap: str | Callable[[_ArrayLike], _NDArray] = 'gray',
1160) -> None:
1161  """Compare two images using an interactive slider.
1162
1163  Displays an HTML slider component to interactively swipe between two images.
1164  The slider functionality requires that the web browser have Internet access.
1165  See additional info in `https://github.com/sneas/img-comparison-slider`.
1166
1167  >>> image1, image2 = np.random.rand(64, 64, 3), color_ramp((64, 64))
1168  >>> compare_images([image1, image2])
1169
1170  Args:
1171    images: Iterable of images.  Each image must be either a 2D array or a 3D
1172      array with 1, 3, or 4 channels.  There must be exactly two images.
1173    vmin: For single-channel image, explicit min value for display.
1174    vmax: For single-channel image, explicit max value for display.
1175    cmap: For single-channel image, `pyplot` color map or callable to map 1D to
1176      3D color.
1177  """
1178  list_images = [
1179      _ensure_mapped_to_rgb(image, vmin=vmin, vmax=vmax, cmap=cmap)
1180      for image in images
1181  ]
1182  if len(list_images) != 2:
1183    raise ValueError('The number of images must be 2.')
1184  png_datas = [compress_image(to_uint8(image)) for image in list_images]
1185  b64_1, b64_2 = [
1186      base64.b64encode(png_data).decode('utf-8') for png_data in png_datas
1187  ]
1188  s = _IMAGE_COMPARISON_HTML.replace('{b64_1}', b64_1).replace('{b64_2}', b64_2)
1189  _display_html(s)
1190
1191
1192# ** Video I/O.
1193
1194
1195def _filename_suffix_from_codec(codec: str) -> str:
1196  if codec == 'gif':
1197    return '.gif'
1198  if codec == 'vp9':
1199    return '.webm'
1200
1201  return '.mp4'
1202
1203
1204def _get_ffmpeg_path() -> str:
1205  path = _search_for_ffmpeg_path()
1206  if not path:
1207    raise RuntimeError(
1208        f"Program '{_config.ffmpeg_name_or_path}' is not found;"
1209        " perhaps install ffmpeg using 'apt install ffmpeg'."
1210    )
1211  return path
1212
1213
1214@typing.overload
1215def _run_ffmpeg(
1216    ffmpeg_args: Sequence[str],
1217    stdin: int | None = None,
1218    stdout: int | None = None,
1219    stderr: int | None = None,
1220    encoding: None = None,  # No encoding -> bytes
1221    allowed_input_files: Sequence[str] | None = None,
1222    allowed_output_files: Sequence[str] | None = None,
1223    sandbox_max_run_time_secs: int | None = None,
1224) -> subprocess.Popen[bytes]:
1225  ...
1226
1227
1228@typing.overload
1229def _run_ffmpeg(
1230    ffmpeg_args: Sequence[str],
1231    stdin: int | None = None,
1232    stdout: int | None = None,
1233    stderr: int | None = None,
1234    encoding: str = ...,  # Encoding -> str
1235    allowed_input_files: Sequence[str] | None = None,
1236    allowed_output_files: Sequence[str] | None = None,
1237    sandbox_max_run_time_secs: int | None = None,
1238) -> subprocess.Popen[str]:
1239  ...
1240
1241
1242def _run_ffmpeg(
1243    ffmpeg_args: Sequence[str],
1244    stdin: int | None = None,
1245    stdout: int | None = None,
1246    stderr: int | None = None,
1247    encoding: str | None = None,
1248    allowed_input_files: Sequence[str] | None = None,
1249    allowed_output_files: Sequence[str] | None = None,
1250    sandbox_max_run_time_secs: int | None = None,
1251) -> subprocess.Popen[bytes] | subprocess.Popen[str]:
1252  """Runs ffmpeg with the given args.
1253
1254  Args:
1255    ffmpeg_args: The args to pass to ffmpeg.
1256    stdin: Same as in `subprocess.Popen`.
1257    stdout: Same as in `subprocess.Popen`.
1258    stderr: Same as in `subprocess.Popen`.
1259    encoding: Same as in `subprocess.Popen`.
1260    allowed_input_files: The input files to allow for ffmpeg.
1261    allowed_output_files: The output files to allow for ffmpeg.
1262    sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
1263      If None, the default limit is 30 minutes.
1264
1265  Returns:
1266    The subprocess.Popen object with running ffmpeg process.
1267  """
1268  argv = []
1269  # In open source, keep env=None to preserve default behavior.
1270  # Context: https://github.com/google/mediapy/pull/62
1271  env: Any = None  # pylint: disable=unused-variable
1272  ffmpeg_path = _get_ffmpeg_path()
1273
1274  # Sandbox max runtime, allowed input and ouput files are not supported in
1275  # open source.
1276  del allowed_input_files
1277  del allowed_output_files
1278  del sandbox_max_run_time_secs
1279
1280  argv.append(ffmpeg_path)
1281  argv.extend(ffmpeg_args)
1282
1283  return subprocess.Popen(
1284      argv,
1285      stdin=stdin,
1286      stdout=stdout,
1287      stderr=stderr,
1288      encoding=encoding,
1289      env=env,
1290  )
1291
1292
1293def video_is_available() -> bool:
1294  """Returns True if the program `ffmpeg` is found.
1295
1296  See also `set_ffmpeg`.
1297  """
1298  return _search_for_ffmpeg_path() is not None
1299
1300
1301class VideoMetadata(NamedTuple):
1302  """Represents the data stored in a video container header.
1303
1304  Attributes:
1305    num_images: Number of frames that is expected from the video stream.  This
1306      is estimated from the framerate and the duration stored in the video
1307      header, so it might be inexact.  We set the value to -1 if number of
1308      frames is not found in the header.
1309    shape: The dimensions (height, width) of each video frame.
1310    fps: The framerate in frames per second.
1311    bps: The estimated bitrate of the video stream in bits per second, retrieved
1312      from the video header.
1313  """
1314
1315  num_images: int
1316  shape: tuple[int, int]
1317  fps: float
1318  bps: int | None
1319
1320
1321def _get_video_metadata(path: _Path) -> VideoMetadata:
1322  """Returns attributes of video stored in the specified local file."""
1323  if not pathlib.Path(path).is_file():
1324    raise RuntimeError(f"Video file '{path}' is not found.")
1325
1326  command = [
1327      '-nostdin',
1328      '-i',
1329      str(path),
1330      '-acodec',
1331      'copy',
1332      # Necessary to get "frame= *(\d+)" using newer ffmpeg versions.
1333      # Previously, was `'-vcodec', 'copy'`
1334      '-vf',
1335      'select=1',
1336      '-vsync',
1337      '0',
1338      '-f',
1339      'null',
1340      '-',
1341  ]
1342  with _run_ffmpeg(
1343      command,
1344      allowed_input_files=[str(path)],
1345      stderr=subprocess.PIPE,
1346      encoding='utf-8',
1347  ) as proc:
1348    _, err = proc.communicate()
1349  bps = fps = num_images = width = height = rotation = None
1350  before_output_info = True
1351  for line in err.split('\n'):
1352    if line.startswith('Output '):
1353      before_output_info = False
1354    if match := re.search(r', bitrate: *([\d.]+) kb/s', line):
1355      bps = int(match.group(1)) * 1000
1356    if matches := re.findall(r'frame= *(\d+) ', line):
1357      num_images = int(matches[-1])
1358    if 'Stream #0:' in line and ': Video:' in line and before_output_info:
1359      if not (match := re.search(r', (\d+)x(\d+)', line)):
1360        raise RuntimeError(f'Unable to parse video dimensions in line {line}')
1361      width, height = int(match.group(1)), int(match.group(2))
1362      # Some videos have a 'k' suffix for kiloframes per second. Thus an extra
1363      # `(k?)` is used for this case. This usually happens when we don't know
1364      # the exact framerate. However, in order to not raise an error here, we
1365      # will try to parse the framerate as x1000.
1366      if match := re.search(r', ([\d.]+)(k?) fps', line):
1367        number = float(match.group(1))
1368        if match.group(2) == 'k':
1369          fps = number * 1000
1370        else:
1371          fps = number
1372      elif str(path).endswith('.gif'):
1373        # Some GIF files lack a framerate attribute; use a reasonable default.
1374        fps = 10
1375      else:
1376        raise RuntimeError(f'Unable to parse video framerate in line {line}')
1377    if match := re.fullmatch(r'\s*rotate\s*:\s*(\d+)', line):
1378      rotation = int(match.group(1))
1379    if match := re.fullmatch(r'.*rotation of (-?\d+).*\sdegrees', line):
1380      rotation = int(match.group(1))
1381  if not num_images:
1382    num_images = -1
1383  if not width:
1384    raise RuntimeError(f'Unable to parse video header: {err}')
1385  # By default, ffmpeg enables "-autorotate"; we just fix the dimensions.
1386  if rotation in (90, 270, -90, -270):
1387    width, height = height, width
1388  assert height is not None and width is not None
1389  shape = height, width
1390  assert fps is not None
1391  return VideoMetadata(num_images, shape, fps, bps)
1392
1393
1394class _VideoIO:
1395  """Base class for `VideoReader` and `VideoWriter`."""
1396
1397  def _get_pix_fmt(self, dtype: _DType, image_format: str) -> str:
1398    """Returns ffmpeg pix_fmt given data type and image format."""
1399    native_endian_suffix = {'little': 'le', 'big': 'be'}[sys.byteorder]
1400    return {
1401        np.uint8: {
1402            'rgb': 'rgb24',
1403            'yuv': 'yuv444p',
1404            'gray': 'gray',
1405        },
1406        np.uint16: {
1407            'rgb': 'rgb48' + native_endian_suffix,
1408            'yuv': 'yuv444p16' + native_endian_suffix,
1409            'gray': 'gray16' + native_endian_suffix,
1410        },
1411    }[dtype.type][image_format]  # pyrefly: ignore[bad-index]
1412
1413
1414class VideoReader(_VideoIO):
1415  """Context to read a compressed video as an iterable over its images.
1416
1417  >>> with VideoReader('/tmp/river.mp4') as reader:
1418  ...   print(f'Video has {reader.num_images} images with shape={reader.shape},'
1419  ...         f' at {reader.fps} frames/sec and {reader.bps} bits/sec.')
1420  ...   for image in reader:
1421  ...     print(image.shape)
1422
1423  >>> with VideoReader('/tmp/river.mp4') as reader:
1424  ...   video = np.array(tuple(reader))
1425
1426  >>> url = 'https://github.com/hhoppe/data/raw/main/video.mp4'
1427  >>> with VideoReader(url) as reader:
1428  ...   show_video(reader)
1429
1430  Attributes:
1431    path_or_url: Location of input video.
1432    output_format: Format of output images (default 'rgb').  If 'rgb', each
1433      image has shape=(height, width, 3) with R, G, B values.  If 'yuv', each
1434      image has shape=(height, width, 3) with Y, U, V values.  If 'gray', each
1435      image has shape=(height, width).
1436    dtype: Data type for output images.  The default is `np.uint8`.  Use of
1437      `np.uint16` allows reading 10-bit or 12-bit data without precision loss.
1438    metadata: Object storing the information retrieved from the video header.
1439      Its attributes are copied as attributes in this class.
1440    num_images: Number of frames that is expected from the video stream.  This
1441      is estimated from the framerate and the duration stored in the video
1442      header, so it might be inexact.
1443    shape: The dimensions (height, width) of each video frame.
1444    fps: The framerate in frames per second.
1445    bps: The estimated bitrate of the video stream in bits per second, retrieved
1446      from the video header.
1447    stream_index: The stream index to read from. The default is 0.
1448    sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
1449      If None, the default limit is 30 minutes. Unused in open source.
1450  """
1451
1452  path_or_url: _Path
1453  output_format: str
1454  dtype: _DType
1455  metadata: VideoMetadata
1456  num_images: int
1457  shape: tuple[int, int]
1458  fps: float
1459  bps: int | None
1460  stream_index: int
1461  _num_bytes_per_image: int
1462
1463  def __init__(
1464      self,
1465      path_or_url: _Path,
1466      *,
1467      stream_index: int = 0,
1468      output_format: str = 'rgb',
1469      dtype: _DTypeLike = np.uint8,
1470      sandbox_max_run_time_secs: int | None = None,
1471  ):
1472    if output_format not in {'rgb', 'yuv', 'gray'}:
1473      raise ValueError(
1474          f'Output format {output_format} is not rgb, yuv, or gray.'
1475      )
1476    self.path_or_url = path_or_url
1477    self.output_format = output_format
1478    self.stream_index = stream_index
1479    self.dtype = np.dtype(dtype)
1480    if self.dtype.type not in (np.uint8, np.uint16):
1481      raise ValueError(f'Type {dtype} is not np.uint8 or np.uint16.')
1482    self.sandbox_max_run_time_secs = sandbox_max_run_time_secs
1483    self._read_via_local_file: Any = None
1484    self._popen: subprocess.Popen[bytes] | None = None
1485    self._proc: subprocess.Popen[bytes] | None = None
1486
1487  def __enter__(self) -> 'VideoReader':
1488    try:
1489      self._read_via_local_file = _read_via_local_file(self.path_or_url)
1490      # pylint: disable-next=no-member
1491      tmp_name = self._read_via_local_file.__enter__()
1492
1493      self.metadata = _get_video_metadata(tmp_name)
1494      self.num_images, self.shape, self.fps, self.bps = self.metadata
1495      pix_fmt = self._get_pix_fmt(self.dtype, self.output_format)
1496      num_channels = {'rgb': 3, 'yuv': 3, 'gray': 1}[self.output_format]
1497      bytes_per_channel = self.dtype.itemsize
1498      self._num_bytes_per_image = (
1499          math.prod(self.shape) * num_channels * bytes_per_channel
1500      )
1501
1502      command = [
1503          '-v',
1504          'panic',
1505          '-nostdin',
1506          '-i',
1507          tmp_name,
1508          '-vcodec',
1509          'rawvideo',
1510          '-f',
1511          'image2pipe',
1512          '-map',
1513          f'0:v:{self.stream_index}',
1514          '-pix_fmt',
1515          pix_fmt,
1516          '-vsync',
1517          'vfr',
1518          '-',
1519      ]
1520      self._popen = _run_ffmpeg(
1521          command,
1522          stdout=subprocess.PIPE,
1523          stderr=subprocess.PIPE,
1524          allowed_input_files=[tmp_name],
1525          sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
1526      )
1527      self._proc = self._popen.__enter__()
1528    except Exception:
1529      self.__exit__(None, None, None)
1530      raise
1531    return self
1532
1533  def __exit__(self, *_: Any) -> None:
1534    self.close()
1535
1536  def read(self) -> _NDArray | None:
1537    """Reads a video image frame (or None if at end of file).
1538
1539    Returns:
1540      A numpy array in the format specified by `output_format`, i.e., a 3D
1541      array with 3 color channels, except for format 'gray' which is 2D.
1542
1543    Raises:
1544      RuntimeError: If there is an error reading from the output file.
1545    """
1546    assert self._proc, 'Error: reading from an already closed context.'
1547    stdout = self._proc.stdout
1548    assert stdout is not None
1549    data = stdout.read(self._num_bytes_per_image)
1550    if not data:  # Due to either end-of-file or subprocess error.
1551      self.close()  # Raises exception if subprocess had error.
1552      return None  # To indicate end-of-file.
1553    if len(data) != self._num_bytes_per_image:
1554      self._proc.wait()
1555      stderr = self._proc.stderr
1556      stderr_output = ''
1557      if stderr is not None:
1558        stderr_output = stderr.read().decode('utf-8', errors='replace').strip()
1559      raise RuntimeError(
1560          f'ffmpeg exited with code {self._proc.returncode}.\nIncomplete'
1561          f' frame read: expected {self._num_bytes_per_image} bytes, but got'
1562          f' {len(data)}.\nffmpeg stderr:\n{stderr_output}'
1563      )
1564    image = np.frombuffer(data, dtype=self.dtype)
1565    if self.output_format == 'rgb':
1566      image = image.reshape(*self.shape, 3)
1567    elif self.output_format == 'yuv':  # Convert from planar YUV to pixel YUV.
1568      image = np.moveaxis(image.reshape(3, *self.shape), 0, 2)
1569    elif self.output_format == 'gray':  # Generate 2D rather than 3D ndimage.
1570      image = image.reshape(*self.shape)
1571    else:
1572      raise AssertionError
1573    return image
1574
1575  def __iter__(self) -> Iterator[_NDArray]:
1576    while True:
1577      image = self.read()
1578      if image is None:
1579        return
1580      yield image
1581
1582  def close(self) -> None:
1583    """Terminates video reader.  (Called automatically at end of context.)"""
1584    if self._popen:
1585      self._popen.__exit__(None, None, None)
1586      self._popen = None
1587      self._proc = None
1588    if self._read_via_local_file:
1589      # pylint: disable-next=no-member
1590      self._read_via_local_file.__exit__(None, None, None)
1591      self._read_via_local_file = None
1592
1593
1594class VideoWriter(_VideoIO):
1595  """Context to write a compressed video.
1596
1597  >>> shape = 480, 640
1598  >>> with VideoWriter('/tmp/v.mp4', shape, fps=60) as writer:
1599  ...   for image in moving_circle(shape, num_images=60):
1600  ...     writer.add_image(image)
1601  >>> show_video(read_video('/tmp/v.mp4'))
1602
1603
1604  Bitrate control may be specified using at most one of: `bps`, `qp`, or `crf`.
1605  If none are specified, `qp` is set to a default value.
1606  See https://slhck.info/video/2017/03/01/rate-control.html
1607
1608  If codec is 'gif', the args `bps`, `qp`, `crf`, and `encoded_format` are
1609  ignored.
1610
1611  Attributes:
1612    path: Output video.  Its suffix (e.g. '.mp4') determines the video container
1613      format.  The suffix must be '.gif' if the codec is 'gif'.
1614    shape: 2D spatial dimensions (height, width) of video image frames.  The
1615      dimensions must be even if 'encoded_format' has subsampled chroma (e.g.,
1616      'yuv420p' or 'yuv420p10le').
1617    codec: Compression algorithm as defined by "ffmpeg -codecs" (e.g., 'h264',
1618      'hevc', 'vp9', or 'gif').
1619    metadata: Optional VideoMetadata object whose `fps` and `bps` attributes are
1620      used if not specified as explicit parameters.
1621    fps: Frames-per-second framerate (default is 60.0 except 25.0 for 'gif').
1622    bps: Requested average bits-per-second bitrate (default None).
1623    qp: Quantization parameter for video compression quality (default None).
1624    crf: Constant rate factor for video compression quality (default None).
1625    ffmpeg_args: Additional arguments for `ffmpeg` command, e.g. '-g 30' to
1626      introduce I-frames, or '-bf 0' to omit B-frames.
1627    input_format: Format of input images (default 'rgb').  If 'rgb', each image
1628      has shape=(height, width, 3) or (height, width).  If 'yuv', each image has
1629      shape=(height, width, 3) with Y, U, V values.  If 'gray', each image has
1630      shape=(height, width).
1631    dtype: Expected data type for input images (any float input images are
1632      converted to `dtype`).  The default is `np.uint8`.  Use of `np.uint16` is
1633      necessary when encoding >8 bits/channel.
1634    encoded_format: Pixel format as defined by `ffmpeg -pix_fmts`, e.g.,
1635      'yuv420p' (2x2-subsampled chroma), 'yuv444p' (full-res chroma),
1636      'yuv420p10le' (10-bit per channel), etc.  The default (None) selects
1637      'yuv420p' if all shape dimensions are even, else 'yuv444p'.
1638    sandbox_max_run_time_secs: The maximum time in seconds to run the sandbox.
1639      If None, the default limit is 30 minutes. Unused in open source.
1640    audio: Optional audio data as a NumPy array.  It should have shape (N,) for
1641      mono or (N, C) for multi-channel audio, where N is the number of samples
1642      and C is the number of channels.  The dtype must be one of `np.uint8`,
1643      `np.int16`, `np.int32`, `np.float32`, or `np.float64`, in native byte
1644      order.  Float samples are nominally in [-1.0, 1.0], and values outside
1645      this range typically clip; integer samples span the full range of their
1646      type, where `np.uint8` is unsigned with silence at 128.  The two streams
1647      are not truncated to a common length: if the audio is shorter than the
1648      video, the remainder is silent; if it is longer, the audio continues past
1649      the end of the video stream (most players keep showing the last frame).
1650    audio_sample_rate: Sample rate of the audio in Hz. Required if `audio` is
1651      provided.
1652    audio_codec: Audio compression algorithm as defined by "ffmpeg -codecs"
1653      (default 'aac').  It must be supported by the video container, e.g., 'aac'
1654      for MP4 ('h264' or 'hevc') or 'libopus' for WebM ('vp9').  Ignored if
1655      `audio` is None.
1656  """
1657
1658  def __init__(
1659      self,
1660      path: _Path,
1661      shape: tuple[int, int],
1662      *,
1663      codec: str = 'h264',
1664      metadata: VideoMetadata | None = None,
1665      fps: float | None = None,
1666      bps: int | None = None,
1667      qp: int | None = None,
1668      crf: float | None = None,
1669      ffmpeg_args: str | Sequence[str] = '',
1670      input_format: str = 'rgb',
1671      dtype: _DTypeLike = np.uint8,
1672      encoded_format: str | None = None,
1673      sandbox_max_run_time_secs: int | None = None,
1674      audio: _NDArray | None = None,
1675      audio_sample_rate: int | None = None,
1676      audio_codec: str = 'aac',
1677  ) -> None:
1678    _check_2d_shape(shape)
1679    if fps is None and metadata:
1680      fps = metadata.fps
1681    if fps is None:
1682      fps = 25.0 if codec == 'gif' else 60.0
1683    if fps <= 0.0:
1684      raise ValueError(f'Frame-per-second value {fps} is invalid.')
1685    if bps is None and metadata:
1686      bps = metadata.bps
1687    bps = int(bps) if bps is not None else None
1688    if bps is not None and bps <= 0:
1689      raise ValueError(f'Bitrate value {bps} is invalid.')
1690    if qp is not None and (not isinstance(qp, int) or qp < 0):
1691      raise ValueError(
1692          f'Quantization parameter {qp} cannot be negative. It must be a'
1693          ' non-negative integer.'
1694      )
1695    num_rate_specifications = sum(x is not None for x in (bps, qp, crf))
1696    if num_rate_specifications > 1:
1697      raise ValueError(
1698          f'Must specify at most one of bps, qp, or crf ({bps}, {qp}, {crf}).'
1699      )
1700    ffmpeg_args = (
1701        shlex.split(ffmpeg_args)
1702        if isinstance(ffmpeg_args, str)
1703        else list(ffmpeg_args)
1704    )
1705    if input_format not in {'rgb', 'yuv', 'gray'}:
1706      raise ValueError(f'Input format {input_format} is not rgb, yuv, or gray.')
1707    dtype = np.dtype(dtype)
1708    if dtype.type not in (np.uint8, np.uint16):
1709      raise ValueError(f'Type {dtype} is not np.uint8 or np.uint16.')
1710    if audio is not None:
1711      if audio_sample_rate is None:
1712        raise ValueError('The audio_sample_rate must be set if audio is set.')
1713      if audio.ndim not in (1, 2):
1714        raise ValueError(f'Audio shape {audio.shape} is not (N,) or (N, C).')
1715      if audio.dtype.type not in _FFMPEG_AUDIO_FORMAT_FROM_DTYPE:
1716        raise ValueError(f'Audio type {audio.dtype} is unsupported.')
1717      if not audio.dtype.isnative:
1718        raise ValueError(
1719            f'Audio type {audio.dtype} is not in native byte order.'
1720        )
1721      if codec == 'gif':
1722        raise ValueError('Audio is not supported with the gif codec.')
1723    self.path = pathlib.Path(path)
1724    self.shape = shape
1725    all_dimensions_are_even = all(dim % 2 == 0 for dim in shape)
1726    if encoded_format is None:
1727      encoded_format = 'yuv420p' if all_dimensions_are_even else 'yuv444p'
1728    if not all_dimensions_are_even and encoded_format.startswith(
1729        ('yuv42', 'yuvj42')
1730    ):
1731      raise ValueError(
1732          f'With encoded_format {encoded_format}, video dimensions must be'
1733          f' even, but shape is {shape}.'
1734      )
1735    self.fps = fps
1736    self.codec = codec
1737    self.bps = bps
1738    self.qp = qp
1739    self.crf = crf
1740    self.ffmpeg_args = ffmpeg_args
1741    self.input_format = input_format
1742    self.dtype = dtype
1743    self.encoded_format = encoded_format
1744    self.sandbox_max_run_time_secs = sandbox_max_run_time_secs
1745    self.audio = audio
1746    self.audio_sample_rate = audio_sample_rate
1747    self.audio_codec = audio_codec
1748    if num_rate_specifications == 0 and not ffmpeg_args:
1749      qp = 20 if math.prod(self.shape) <= 640 * 480 else 28
1750    self._bitrate_args = (
1751        (['-vb', f'{bps}'] if bps is not None else [])
1752        + (['-qp', f'{qp}'] if qp is not None else [])
1753        + (['-vb', '0', '-crf', f'{crf}'] if crf is not None else [])
1754    )
1755    if self.codec == 'gif':
1756      if self.path.suffix != '.gif':
1757        raise ValueError(f"File '{self.path}' does not have a .gif suffix.")
1758      self.encoded_format = 'pal8'
1759      self._bitrate_args = []
1760      video_filter = 'split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse'
1761      # Less common (and likely less useful) is a per-frame color palette:
1762      # video_filter = ('split[s0][s1];[s0]palettegen=stats_mode=single[p];'
1763      #                 '[s1][p]paletteuse=new=1')
1764      self.ffmpeg_args = ['-vf', video_filter, '-f', 'gif'] + self.ffmpeg_args
1765    self._exit_stack: contextlib.ExitStack | None = None
1766    self._proc: subprocess.Popen[bytes] | None = None
1767
1768  def __enter__(self) -> 'VideoWriter':
1769    input_pix_fmt = self._get_pix_fmt(self.dtype, self.input_format)
1770    try:
1771      self._exit_stack = contextlib.ExitStack()
1772      tmp_name = self._exit_stack.enter_context(
1773          _write_via_local_file(self.path)
1774      )
1775
1776      # Writing to stdout using ('-f', 'mp4', '-') would require
1777      # ('-movflags', 'frag_keyframe+empty_moov') and the result is nonportable.
1778      height, width = self.shape
1779
1780      audio_input_args = []
1781      audio_output_args = ['-an']
1782      allowed_input_files = []
1783
1784      if self.audio is not None:
1785        audio_path = self._exit_stack.enter_context(
1786            _audio_via_local_file(self.audio)
1787        )
1788        channels = 1 if self.audio.ndim == 1 else self.audio.shape[1]
1789        audio_format = _FFMPEG_AUDIO_FORMAT_FROM_DTYPE[self.audio.dtype.type]
1790        if self.audio.dtype.itemsize > 1:
1791          audio_format += {'little': 'le', 'big': 'be'}[sys.byteorder]
1792        audio_input_args = [
1793            '-f',
1794            audio_format,
1795            '-ar',
1796            str(self.audio_sample_rate),
1797            '-ac',
1798            str(channels),
1799            '-i',
1800            audio_path,
1801        ]
1802        audio_output_args = ['-c:a', self.audio_codec]
1803        allowed_input_files.append(audio_path)
1804
1805      command = (
1806          [
1807              '-v',
1808              'error',
1809              '-f',
1810              'rawvideo',
1811              '-vcodec',
1812              'rawvideo',
1813              '-pix_fmt',
1814              input_pix_fmt,
1815              '-s',
1816              f'{width}x{height}',
1817              '-r',
1818              f'{self.fps}',
1819              '-i',
1820              '-',
1821          ]
1822          + audio_input_args
1823          + audio_output_args
1824          + [
1825              '-vcodec',
1826              self.codec,
1827              '-pix_fmt',
1828              self.encoded_format,
1829          ]
1830          + self._bitrate_args
1831          + self.ffmpeg_args
1832          + ['-y', tmp_name]
1833      )
1834      self._proc = self._exit_stack.enter_context(
1835          _run_ffmpeg(
1836              command,
1837              stdin=subprocess.PIPE,
1838              stderr=subprocess.PIPE,
1839              # `_run_ffmpeg` omits the sandbox flag only for None, so an empty
1840              # list would pass an empty '--sandbox_read_access_files'.
1841              allowed_input_files=allowed_input_files or None,
1842              allowed_output_files=[tmp_name],
1843              sandbox_max_run_time_secs=self.sandbox_max_run_time_secs,
1844          )
1845      )
1846    except Exception:
1847      self.__exit__(None, None, None)
1848      raise
1849    return self
1850
1851  def __exit__(self, *_: Any) -> None:
1852    self.close()
1853
1854  def add_image(self, image: _NDArray) -> None:
1855    """Writes a video frame.
1856
1857    Args:
1858      image: Array whose dtype and first two dimensions must match the `dtype`
1859        and `shape` specified in `VideoWriter` initialization.  If
1860        `input_format` is 'gray', the image must be 2D.  For the 'rgb'
1861        input_format, the image may be either 2D (interpreted as grayscale) or
1862        3D with three (R, G, B) channels.  For the 'yuv' input_format, the image
1863        must be 3D with three (Y, U, V) channels.
1864
1865    Raises:
1866      RuntimeError: If there is an error writing to the output file.
1867    """
1868    assert self._proc, 'Error: writing to an already closed context.'
1869    if issubclass(image.dtype.type, (np.floating, np.bool_)):
1870      image = to_type(image, self.dtype)
1871    if image.dtype != self.dtype:
1872      raise ValueError(f'Image type {image.dtype} != {self.dtype}.')
1873    if self.input_format == 'gray':
1874      if image.ndim != 2:
1875        raise ValueError(f'Image dimensions {image.shape} are not 2D.')
1876    else:
1877      if image.ndim == 2 and self.input_format == 'rgb':
1878        image = np.dstack((image, image, image))
1879      if not (image.ndim == 3 and image.shape[2] == 3):
1880        raise ValueError(f'Image dimensions {image.shape} are invalid.')
1881    if image.shape[:2] != self.shape:
1882      raise ValueError(
1883          f'Image dimensions {image.shape[:2]} do not match'
1884          f' those of the initialized video {self.shape}.'
1885      )
1886    if self.input_format == 'yuv':  # Convert from per-pixel YUV to planar YUV.
1887      image = np.moveaxis(image, 2, 0)
1888    data = image.tobytes()
1889    stdin = self._proc.stdin
1890    assert stdin is not None
1891    if stdin.write(data) != len(data):
1892      self._proc.wait()
1893      stderr = self._proc.stderr
1894      assert stderr is not None
1895      s = stderr.read().decode('utf-8')
1896      raise RuntimeError(f"Error writing '{self.path}': {s}")
1897
1898  def close(self) -> None:
1899    """Finishes writing the video.  (Called automatically at end of context.)"""
1900    if self._exit_stack is None:
1901      return
1902    # Unwinding the stack terminates the `ffmpeg` process, removes the
1903    # temporary audio file, and copies the encoded video to a remote `path`.
1904    # Raising within the `with` propagates the error into those contexts, so
1905    # an incomplete video is discarded rather than copied.
1906    with self._exit_stack:
1907      self._exit_stack = None
1908      if self._proc is not None:
1909        proc, self._proc = self._proc, None
1910        stdin = proc.stdin
1911        assert stdin is not None
1912        stdin.close()
1913        if proc.wait():
1914          stderr = proc.stderr
1915          assert stderr is not None
1916          s = stderr.read().decode('utf-8')
1917          raise RuntimeError(f"Error writing '{self.path}': {s}")
1918
1919
1920class _VideoArray(np.ndarray):
1921  """Wrapper to add a VideoMetadata `metadata` attribute to a numpy array."""
1922
1923  metadata: VideoMetadata | None
1924
1925  def __new__(
1926      cls: Type['_VideoArray'],
1927      input_array: _NDArray,
1928      metadata: VideoMetadata | None = None,
1929  ) -> '_VideoArray':
1930    obj: _VideoArray = np.asarray(input_array).view(cls)
1931    obj.metadata = metadata
1932    return obj
1933
1934  def __array_finalize__(self, obj: Any) -> None:
1935    if obj is None:
1936      return
1937    self.metadata = getattr(obj, 'metadata', None)
1938
1939
1940def read_video(path_or_url: _Path, **kwargs: Any) -> _VideoArray:
1941  """Returns an array containing all images read from a compressed video file.
1942
1943  >>> video = read_video('/tmp/river.mp4')
1944  >>> print(f'The framerate is {video.metadata.fps} frames/s.')
1945  >>> show_video(video)
1946
1947  >>> url = 'https://github.com/hhoppe/data/raw/main/video.mp4'
1948  >>> show_video(read_video(url))
1949
1950  Args:
1951    path_or_url: Input video file.
1952    **kwargs: Additional parameters for `VideoReader`.
1953
1954  Returns:
1955    A 4D `numpy` array with dimensions (frame, height, width, channel), or a 3D
1956    array if `output_format` is specified as 'gray'.  The returned array has an
1957    attribute `metadata` containing `VideoMetadata` information.  This enables
1958    `show_video` to retrieve the framerate in `metadata.fps`.  Note that the
1959    metadata attribute is lost in most subsequent `numpy` operations.
1960  """
1961  with VideoReader(path_or_url, **kwargs) as reader:
1962    return _VideoArray(np.array(tuple(reader)), metadata=reader.metadata)
1963
1964
1965def write_video(path: _Path, images: Iterable[_NDArray], **kwargs: Any) -> None:
1966  """Writes images to a compressed video file.
1967
1968  >>> video = moving_circle((480, 640), num_images=60)
1969  >>> write_video('/tmp/v.mp4', video, fps=60, qp=18)
1970  >>> show_video(read_video('/tmp/v.mp4'))
1971
1972  Args:
1973    path: Output video file.
1974    images: Iterable over video frames, e.g. a 4D array or a list of 2D or 3D
1975      arrays.
1976    **kwargs: Additional parameters for `VideoWriter`.
1977  """
1978  first_image, images = _peek_first(images)
1979  shape = first_image.shape[0], first_image.shape[1]
1980  dtype = first_image.dtype
1981  if dtype == bool:
1982    dtype = np.dtype(np.uint8)
1983  elif np.issubdtype(dtype, np.floating):
1984    dtype = np.dtype(np.uint16)
1985  kwargs = {'metadata': getattr(images, 'metadata', None), **kwargs}
1986  with VideoWriter(path, shape=shape, dtype=dtype, **kwargs) as writer:
1987    for image in images:
1988      writer.add_image(image)
1989
1990
1991def compress_video(
1992    images: Iterable[_NDArray], *, codec: str = 'h264', **kwargs: Any
1993) -> bytes:
1994  """Returns a buffer containing a compressed video.
1995
1996  The video container is 'gif' for 'gif' codec, 'webm' for 'vp9' codec,
1997  and mp4 otherwise.
1998
1999  >>> video = read_video('/tmp/river.mp4')
2000  >>> data = compress_video(video, bps=10_000_000)
2001  >>> print(len(data))
2002
2003  >>> data = compress_video(moving_circle((100, 100), num_images=10), fps=10)
2004
2005  Args:
2006    images: Iterable over video frames.
2007    codec: Compression algorithm as defined by `ffmpeg -codecs` (e.g., 'h264',
2008      'hevc', 'vp9', or 'gif').
2009    **kwargs: Additional parameters for `VideoWriter`.
2010
2011  Returns:
2012    A bytes buffer containing the compressed video.
2013  """
2014  suffix = _filename_suffix_from_codec(codec)
2015  with tempfile.TemporaryDirectory() as directory_name:
2016    tmp_path = pathlib.Path(directory_name) / f'file{suffix}'
2017    write_video(tmp_path, images, codec=codec, **kwargs)
2018    return tmp_path.read_bytes()
2019
2020
2021def decompress_video(data: bytes, **kwargs: Any) -> _NDArray:
2022  """Returns video images from an MP4-compressed data buffer."""
2023  with tempfile.TemporaryDirectory() as directory_name:
2024    tmp_path = pathlib.Path(directory_name) / 'file.mp4'
2025    tmp_path.write_bytes(data)
2026    return read_video(tmp_path, **kwargs)
2027
2028
2029def html_from_compressed_video(
2030    data: bytes,
2031    width: int,
2032    height: int,
2033    *,
2034    title: str | None = None,
2035    border: bool | str = False,
2036    loop: bool = True,
2037    autoplay: bool = True,
2038) -> str:
2039  """Returns an HTML string with a video tag containing H264-encoded data.
2040
2041  Args:
2042    data: MP4-compressed video bytes.
2043    width: Width of HTML video in pixels.
2044    height: Height of HTML video in pixels.
2045    title: Optional text shown centered above the video.
2046    border: If `bool`, whether to place a black boundary around the image, or if
2047      `str`, the boundary CSS style.
2048    loop: If True, the playback repeats forever.
2049    autoplay: If True, video playback starts without having to click.
2050  """
2051  b64 = base64.b64encode(data).decode('utf-8')
2052  if isinstance(border, str):
2053    border = f'{border}; '
2054  elif border:
2055    border = 'border:1px solid black; '
2056  else:
2057    border = ''
2058  options = (
2059      f'controls width="{width}" height="{height}"'
2060      f' style="{border}object-fit:cover;"'
2061      f'{" loop" if loop else ""}'
2062      f'{" autoplay muted" if autoplay else ""}'
2063  )
2064  s = f"""<video {options}>
2065      <source src="data:video/mp4;base64,{b64}" type="video/mp4"/>
2066      This browser does not support the video tag.
2067      </video>"""
2068  if title is not None:
2069    s = f"""<div style="display:flex; align-items:left;">
2070      <div style="display:flex; flex-direction:column; align-items:center;">
2071      <div>{title}</div><div>{s}</div></div></div>"""
2072  return s
2073
2074
2075def show_video(
2076    images: Iterable[_NDArray], *, title: str | None = None, **kwargs: Any
2077) -> str | None:
2078  """Displays a video in the IPython notebook and optionally saves it to a file.
2079
2080  See `show_videos`.
2081
2082  >>> video = read_video('https://github.com/hhoppe/data/raw/main/video.mp4')
2083  >>> show_video(video, title='River video')
2084
2085  >>> show_video(moving_circle((80, 80), num_images=10), fps=5, border=True)
2086
2087  >>> show_video(read_video('/tmp/river.mp4'))
2088
2089  Args:
2090    images: Iterable of video frames (e.g., a 4D array or a list of 2D or 3D
2091      arrays).
2092    title: Optional text shown centered above the video.
2093    **kwargs: See `show_videos`.
2094
2095  Returns:
2096    html string if `return_html` is `True`.
2097  """
2098  return show_videos([images], [title], **kwargs)
2099
2100
2101def show_videos(
2102    videos: Iterable[Iterable[_NDArray]] | Mapping[str, Iterable[_NDArray]],
2103    titles: Iterable[str | None] | None = None,
2104    *,
2105    width: int | None = None,
2106    height: int | None = None,
2107    downsample: bool = True,
2108    columns: int | None = None,
2109    fps: float | None = None,
2110    bps: int | None = None,
2111    qp: int | None = None,
2112    codec: str = 'h264',
2113    ylabel: str = '',
2114    html_class: str = 'show_videos',
2115    return_html: bool = False,
2116    audios: Iterable[_NDArray] | Mapping[str, _NDArray] | None = None,
2117    audio_sample_rate: int | None = None,
2118    audio_codec: str = 'aac',
2119    **kwargs: Any,
2120) -> str | None:
2121  """Displays a row of videos in the IPython notebook.
2122
2123  Creates HTML with `<video>` tags containing embedded H264-encoded bytestrings.
2124  If `codec` is set to 'gif', we instead use `<img>` tags containing embedded
2125  GIF-encoded bytestrings.  Note that the resulting GIF animations skip frames
2126  when the `fps` period is not a multiple of 10 ms units (GIF frame delay
2127  units).  Encoding at `fps` = 20.0, 25.0, or 50.0 works fine.
2128
2129  If a directory has been specified using `set_show_save_dir`, also saves each
2130  titled video to a file in that directory based on its title.
2131
2132  Args:
2133    videos: Iterable of videos, or dictionary of `{title: video}`.  Each video
2134      must be an iterable of images.  If a video object has a `metadata`
2135      (`VideoMetadata`) attribute, its `fps` field provides a default framerate.
2136    titles: Optional strings shown above the corresponding videos.
2137    width: Optional, overrides displayed width (in pixels).
2138    height: Optional, overrides displayed height (in pixels).
2139    downsample: If True, each video whose width or height is greater than the
2140      specified `width` or `height` is resampled to the display resolution. This
2141      improves antialiasing and reduces the size of the notebook.
2142    columns: Optional, maximum number of videos per row.
2143    fps: Frames-per-second framerate (default is 60.0 except 25.0 for GIF).
2144    bps: Bits-per-second bitrate (default None).
2145    qp: Quantization parameter for video compression quality (default None).
2146    codec: Compression algorithm; must be either 'h264' or 'gif'.
2147    ylabel: Text (rotated by 90 degrees) shown on the left of each row.
2148    html_class: CSS class name used in definition of HTML element.
2149    return_html: If `True` return the raw HTML `str` instead of displaying.
2150    audios: Optional iterable of audio tracks, or dictionary of `{title:
2151      audio}`; see `VideoWriter`.  Each track is attached to the corresponding
2152      video, so an iterable must have one entry (possibly None) per video.  With
2153      a dictionary, its keys must match the video titles exactly (a value may be
2154      None for no audio).  Audio is incompatible with `codec` 'gif'.  With the
2155      default `autoplay=True`, the video starts muted (browsers block unmuted
2156      autoplay); unmute it in the player controls, or pass `autoplay=False`.
2157    audio_sample_rate: Sample rate of the audio in Hz, shared by all tracks.
2158      Required if `audios` is provided.
2159    audio_codec: Audio compression algorithm (default 'aac'); see `VideoWriter`.
2160    **kwargs: Additional parameters (`border`, `loop`, `autoplay`) for
2161      `html_from_compressed_video`.
2162
2163  Returns:
2164    html string if `return_html` is `True`.
2165  """
2166  if isinstance(videos, Mapping):
2167    if titles is not None:
2168      raise ValueError(
2169          'Cannot have both a video dictionary and a titles parameter.'
2170      )
2171    list_titles = list(videos.keys())
2172    list_videos = list(videos.values())
2173  else:
2174    list_videos = list(cast('Iterable[_NDArray]', videos))
2175    list_titles = [None] * len(list_videos) if titles is None else list(titles)
2176    if len(list_videos) != len(list_titles):
2177      raise ValueError(
2178          'Number of videos does not match number of titles'
2179          f' ({len(list_videos)} vs {len(list_titles)}).'
2180      )
2181
2182  if audios is None:
2183    list_audios = [None] * len(list_videos)
2184  elif isinstance(audios, Mapping):
2185    missing = set(list_titles).difference(audios)
2186    extra = set(audios).difference(list_titles)
2187    if missing or extra:
2188      raise ValueError(
2189          'The audios dictionary keys must match the video titles (use None as'
2190          f' the value for no audio); missing: {sorted(missing, key=str)},'
2191          f' extra: {sorted(extra, key=str)}.'
2192      )
2193    list_audios = [audios.get(title) for title in list_titles]  # pyrefly: ignore[bad-argument-type]
2194  else:
2195    list_audios = list(audios)
2196
2197  if len(list_videos) != len(list_audios):
2198    raise ValueError(
2199        'Number of videos does not match number of audio'
2200        f' ({len(list_videos)} vs {len(list_audios)}).'
2201    )
2202
2203  if codec not in {'h264', 'gif'}:
2204    raise ValueError(f'Codec {codec} is neither h264 or gif.')
2205
2206  html_strings = []
2207  for video, title, video_audio in zip(list_videos, list_titles, list_audios):
2208    metadata: VideoMetadata | None = getattr(video, 'metadata', None)
2209    first_image, video = _peek_first(video)
2210    w, h = _get_width_height(width, height, first_image.shape[:2])
2211    if downsample and (w < first_image.shape[1] or h < first_image.shape[0]):
2212      # Not resize_video() because each image may have different depth and type.
2213      video = [resize_image(image, (h, w)) for image in video]
2214      first_image = video[0]
2215    data = compress_video(
2216        video,
2217        metadata=metadata,
2218        fps=fps,
2219        bps=bps,
2220        qp=qp,
2221        codec=codec,
2222        audio=video_audio,
2223        audio_sample_rate=audio_sample_rate,
2224        audio_codec=audio_codec,
2225    )
2226    if title is not None and _config.show_save_dir:
2227      suffix = _filename_suffix_from_codec(codec)
2228      path = pathlib.Path(_config.show_save_dir) / f'{title}{suffix}'
2229      with _open(path, mode='wb') as f:
2230        f.write(data)
2231    if codec == 'gif':
2232      pixelated = h > first_image.shape[0] or w > first_image.shape[1]
2233      html_string = html_from_compressed_image(
2234          data, w, h, title=title, fmt='gif', pixelated=pixelated, **kwargs  # pyrefly: ignore[bad-argument-type]
2235      )
2236    else:
2237      html_string = html_from_compressed_video(
2238          data, w, h, title=title, **kwargs  # pyrefly: ignore[bad-argument-type]
2239      )
2240    html_strings.append(html_string)
2241
2242  # Create single-row tables each with no more than 'columns' elements.
2243  table_strings = []
2244  for row_html_strings in _chunked(html_strings, columns):
2245    td = '<td style="padding:1px;">'
2246    s = ''.join(f'{td}{e}</td>' for e in row_html_strings)
2247    if ylabel:
2248      style = 'writing-mode:vertical-lr; transform:rotate(180deg);'
2249      s = f'{td}<span style="{style}">{ylabel}</span></td>' + s
2250    table_strings.append(
2251        f'<table class="{html_class}"'
2252        f' style="border-spacing:0px;"><tr>{s}</tr></table>'
2253    )
2254  s = ''.join(table_strings)
2255  if return_html:
2256    return s
2257  _display_html(s)
2258  return None
2259
2260
2261# Local Variables:
2262# fill-column: 80
2263# End: