-###Video backend for xarray based on the xarray rasterio backend
-
-
-import os
-import warnings
-import tempfile
-
-import numpy as np
-import numcodecs
-import av
-
-from xarray import DataArray, Dataset
-from xarray.core import indexing
-from xarray.core.utils import is_scalar
-from xarray.backends.common import BackendArray
-from xarray.backends.file_manager import CachingFileManager
-from xarray.backends.locks import SerializableLock
-
-from .exceptions import VideoReadError
-
-VIDEO_LOCK = SerializableLock()
-TEMPDIR = os.path.join(tempfile.gettempdir(), "xarray_video")
-if not os.path.exists(TEMPDIR):
- os.mkdir(TEMPDIR)
-
-compressor = numcodecs.registry.get_codec(dict(id="h264"))
-
-
-def _key_length(key, length):
- if isinstance(key, slice):
- return len(range(*key.indices(length)))
- elif is_scalar(key):
- return 1
- else:
- return length
-
-
-[docs]class VideoArrayWrapper(BackendArray):
- """A wrapper around video dataset objects"""
-
- def __init__(self, manager, lock, shape):
- self.manager = manager
- self.lock = lock
-
- reader = manager.acquire()
- stream = reader.streams.video[0]
-
- self._shape = shape
- self._dtype = np.dtype("uint8")
-
- ts0 = int((100 * av.time_base) / stream.average_rate) + stream.start_time
- reader.seek(ts0)
- for frame in reader.decode(stream):
- dt = frame.dts - stream.start_time
- self._can_seek = dt > 0
- break
- manager.close()
-
- @property
- def dtype(self):
- return self._dtype
-
- @property
- def shape(self):
- return self._shape
-
- def _getitem(self, key):
- assert len(key) == 4, "video DataArrays should always be 4D"
-
- frame_key, y_key, x_key, band_key = key
-
- if isinstance(frame_key, slice):
- f0 = frame_key.start or 0
- f1 = frame_key.stop or self._shape[0]
- fstep = frame_key.step or 1
- elif is_scalar(frame_key):
- f0 = frame_key
- f1 = frame_key + 1
- fstep = 1
- else:
- f0 = 0
- f1 = self._shape[0]
- fstep = 1
- nf = len(range(f0, f1, fstep))
- ny = _key_length(y_key, self._shape[1])
- nx = _key_length(x_key, self._shape[2])
- nb = _key_length(band_key, self._shape[3])
-
- data = np.zeros((nf, ny, nx, nb), dtype="uint8")
- reader = self.manager.acquire()
- stream = reader.streams.video[0]
- if self._can_seek:
- ts0 = int((f0 * av.time_base) / stream.average_rate) + stream.start_time
- reader.seek(ts0)
- frame_start = -1
- else:
- frame_start = 0
- ind0 = 0
- for i, frame in enumerate(reader.decode(video=0)):
- if frame_start < 0:
- dts = frame.dts
- if (
- dts is None
- ): # Some packets at start have dts=None, same for fluxhing packets at end
- if packet.buffer_size > 0:
- dts = 0
- else:
- dts = 1e10
- frame_start = int(dts * stream.time_base * stream.rate)
- ind = frame_start + i
- if ind < f0:
- continue
- elif ind >= f1:
- break
- elif ind % fstep == 0:
- data[ind0] = frame.to_ndarray(format="rgb24")[y_key, x_key, band_key]
- ind0 += 1
- self.manager.close()
- data = np.squeeze(data)
- return data
-
- def __getitem__(self, key):
- return indexing.explicit_indexing_adapter(
- key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
- )
-
-
-def _open_video(filename, mode):
- return av.open(filename, mode=mode)
-
-
-def _write_video(filename, array, fps=25, metadata={}):
-
- writer = av.open(filename, mode="w", format="mp4")
-
- nf, ny, nx, nb = array.shape
-
- stream = writer.add_stream("h264", rate=fps)
- stream.thread_type = "AUTO"
-
- stream.width = nx
- stream.height = ny
- stream.pix_fmt = "yuv420p"
-
- for frame_i in array:
- frame = av.VideoFrame.from_ndarray(frame_i, format="rgb24")
- for packet in stream.encode(frame):
- writer.mux(packet)
-
- # Flush stream
- for packet in stream.encode():
- writer.mux(packet)
-
- writer.close()
-
-
-[docs]def open_video(filename, start_time=None, **kwargs):
- """Video file into an xarray dataset.
-
- This reads a video into an xarray dataset with the video in a DataArray.
- If a start time is provided, a time axis will be created for the frames.
-
- Args:
- filename (string): filename of videos to open
- start_time (:class:`numpy.datetime64`): Start time of video
-
- Returns:
- dataset (:class:`xarray.Dataset`): Dataset with video as a DataArray
-
- Raises:
- VideoReadError: Missing or incompatible files
- """
-
- manager = CachingFileManager(
- _open_video,
- filename,
- lock=VIDEO_LOCK,
- mode="r",
- kwargs=kwargs,
- )
- reader = manager.acquire()
- stream = reader.streams.video[0]
- codec = stream.codec_context
- frames = stream.frames
- # If the frame count is not in metadata, this is likely a matroska file. Then seeking will likely not work either
- # Solution is to scan the file using to demux to count the frames
- if frames == 0:
- for packet in reader.demux(stream):
- if packet.buffer_size > 0:
- frames += 1
-
- fps = int(stream.average_rate)
- width = codec.width
- height = codec.height
-
- coords = {"channel": ["R", "G", "B"]}
- coords["pixel_x"] = np.arange(width)
- coords["pixel_y"] = np.arange(height)
- if start_time:
- times = np.datetime64(start_time) + np.arange(
- 0, 1000 * frames / fps, 1000 / fps
- ).astype("<m8[ms]")
- coords["time"] = ("frame", times)
- else:
- coords["frame"] = np.arange(frames)
-
- # Attributes
- attrs = {"fps": fps, "_video": codec.name}
- data = indexing.LazilyIndexedArray(
- VideoArrayWrapper(
- manager,
- VIDEO_LOCK,
- (
- frames,
- height,
- width,
- 3,
- ),
- )
- )
-
- dataset = Dataset(
- data_vars={
- "video": DataArray(
- data=data,
- dims=("frame", "pixel_y", "pixel_x", "channel"),
- coords=coords,
- attrs=attrs,
- )
- },
- )
- if start_time:
- dataset = dataset.set_xindex("time")
-
- # Set the default zarr compressor and assign preferred chunk sizes
- dataset["video"].encoding = {
- "preferred_chunks": {"channel": 3, "pixel_y": height, "pixel_x": width},
- }
-
- # Make the file closeable
- dataset.set_close(manager.close)
-
- return dataset
-
-import os
-import tempfile
-import numpy
-import io
-
-from numcodecs.abc import Codec
-from numcodecs.compat import ndarray_copy, ensure_contiguous_ndarray
-
-import av
-
-
-[docs]class H264(Codec):
- """Codec providing compression using h264 via pyav
-
- Parameters
- ----------
- fps : int (optional)
- Frames per second in compressed chunk (default 25)
- """
-
- codec_id = "h264"
-
-
-
-[docs] def encode(self, buf):
-
- # normalise input
- nf, ny, nx, nb = buf.shape
-
- # write chunk as video file
- data = io.BytesIO()
- writer = av.open(data, mode="w", format="mp4")
-
- stream = writer.add_stream("h264", rate=self.fps)
- writer.streams.video[0].thread_type = "AUTO"
-
- stream.width = nx
- stream.height = ny
- stream.pix_fmt = "yuv420p"
-
- for frame_i in buf:
- frame = av.VideoFrame.from_ndarray(frame_i, format="rgb24")
- for packet in stream.encode(frame):
- writer.mux(packet)
-
- # Flush stream
- for packet in stream.encode():
- writer.mux(packet)
-
- writer.close()
- data.seek(0)
- return data.read()
-
-[docs] def decode(self, buf, out=None):
- # normalise inputs
- buf = ensure_contiguous_ndarray(buf)
-
- stream = io.BytesIO(buf)
- reader = av.open(stream)
- reader.streams.video[0].thread_type = "AUTO"
-
- codec = reader.streams[0].codec_context
- frames = reader.streams[0].frames
- width = codec.width
- height = codec.height
-
- if out is not None:
- out = ensure_contiguous_ndarray(out)
-
- data = numpy.zeros((frames, height, width, 3), dtype="uint8")
- for i, frame in enumerate(reader.decode(video=0)):
- data[i] = frame.to_ndarray(format="rgb24")
-
- reader.close()
- return ndarray_copy(data, out)
-
-import xarray
-import numpy
-import av
-import matplotlib.pyplot as plt
-
-from .exceptions import VideoDisplayError, VideoWriteError
-from .backend import _write_video
-
-_MAX_VIDEO_SIZE = 1024000000
-
-
-[docs]@xarray.register_dataarray_accessor("video")
-class VideoArray:
- """Video extension for :class:`xarray.DataArray`.
-
- Implements operations on a dataset which includes video data
- """
-
-
-
- def _check_video(self, ndim=4):
- if len(self._arr.shape) != ndim or self._arr.shape[-1] != 3:
- raise VideoError(
- f"Expected array with {ndim} dimensions and with 3 channels in last dimension"
- )
- if self._arr.dtype != "uint8":
- raise VideoError("Expected uint8 dtype")
- return numpy.prod(self._arr.shape[:-1])
-
-[docs] def plot(self, **kwargs):
- """Plot DataArray frame as an image using matplotlib.imshow
-
- Raises:
- VideoError: if DataArray does not represent an image
- """
- try:
- import matplotlib.pyplot as plt
- except:
- raise VideoDisplayError("Need matplotlib installed to plot video")
-
- self._check_video(3)
- plt.imshow(self._arr.values, **kwargs)
- plt.show()
-
-[docs] def play(self, interval=10, repeat=False, **kwargs):
- """Play DataArray as a video using matplotlib.animation.
-
- All the video data is loaded into buffered in memory before rendering. Large video sequences cannot be played with this function.
- Note that this function is unlikely to play the video at the correct frame rate.
-
- Args:
- start_time (Union[datetime.datetime,str,:class:`numpy.datetime64`], "optional*): Start time of video array
- interval (int, *optional*): Interval in milliseconds between each frame render (default 0)
- repeat (bool, *optional*): Repeat animation (default False)
-
- Kwargs:
- kwargs are passed to matplotlib.pyplot.figure
-
- Raises:
- VideoError: if DataArray does not represent an video or the video array is too large
- """
- try:
- import matplotlib.pyplot as plt
- import matplotlib.animation as animation
-
- except:
- raise VideoDisplayError("Need matplotlib installed to play video")
-
- nsize = self._check_video()
- if nsize > _MAX_VIDEO_SIZE:
- raise VideoDisplayError("Video too large to buffer")
-
- fps = self._arr.attrs.get("fps", 1)
- buffer = self._arr.values
-
- fig = plt.figure(**kwargs)
- ax = fig.gca()
- im = ax.imshow(buffer[0], animated=True)
- frame_coords = []
- for c in self._arr.coords:
- if "frame" in self._arr.coords[c].dims:
- frame_coords.append(self._arr.coords[c])
-
- def update_frame(i):
- frame = buffer[i]
- im.set_array(frame)
- ftime = i / fps
- ax.set_title(
- " ".join([c.name + ": " + str(c[i].values) for c in frame_coords])
- )
- return (im,)
-
- ani = animation.FuncAnimation(
- fig,
- update_frame,
- frames=range(len(self._arr)),
- interval=interval,
- repeat=repeat,
- )
- plt.show()
-
-[docs] def to_video(self, filename):
- """Write DataArray to a video file
-
- Args:
- filename (string): name of output file
-
- Raises:
- VideoWriteError: Incompatible DataArray or file write error
-
- """
-
- self._check_video()
-
- try:
- _write_video(
- filename, [self._arr.values], fps=self._arr.attrs.get("fps", 25)
- )
- except Exception as e:
- raise VideoWriteError(f"Error writing to file {filename}: {e}")
-
-import xarray
-import warnings
-
-from .exceptions import VideoWriteError
-from .backend import compressor, _write_video
-
-_DEFAULT_CHUNK_SIZE = 512000000
-
-
-[docs]@xarray.register_dataset_accessor("video")
-class VideoDataset:
- """Video extension for :class:`xarray.Dataset`.
-
- Implements operations on a dataset which includes video data
- """
-
-
-
-[docs] def to_zarr(self, *args, chunk_sizes={}, **kwargs):
- (
- """Write to zarr using a video codec for compatible data variables.
-
- Kwargs:
- chunk_sizes (dict, *optional*): preferred chunk_sizes for frame, pixel_y and pixel_x dimensions
-
- """
- + xarray.core.dataset.Dataset.to_zarr.__doc__
- )
- encoding = {}
- for v in self._dset.data_vars:
- dv = self._dset.data_vars[v]
- if len(dv.shape) == 4 and dv.shape[3] == 3:
- nf, ny, nx, nb = dv.shape
- ny0 = chunk_sizes.get("pixel_y", ny)
- nx0 = chunk_sizes.get("pixel_x", nx)
- nf0 = chunk_sizes.get("frame", _DEFAULT_CHUNK_SIZE // ny0 // nx0 // 3)
- encoding[v] = {
- "compressor": compressor,
- "chunks": [nf0, ny0, nx0, 3],
- }
- self._dset.to_zarr(
- *args,
- encoding=encoding,
- **kwargs,
- )
-
-[docs] def to_video(self, filename, data_var=None):
- """Write Dataset to a video file.
-
- Args:
- filename (string): name of output file
- data_var (string, Optional): Data variable to write as video streams.
-
- Raises:
- VideoWriteError: Incompatible DataArray or file write error
-
- """
- try:
- output_stream = None
- for v in self._dset.data_vars:
- if data_var and v != data_var:
- continue
- video_array = self._dset.data_vars[v]
- if (
- len(video_array.shape) == 4
- and video_array.shape[3] == 3
- and video_array.dtype == "uint8"
- ):
- output_stream = video_array.values
- if output_stream is None:
- warnings.warn("No compatible DataArrays found")
- else:
- _write_video(
- filename,
- output_stream,
- fps=video_array.attrs.get("fps", 25),
- )
-
- except Exception as e:
- raise VideoWriteError(f"Error writing to file {filename}: {e}")
-