Source code for fastplotlib.widgets.nd_widget._video

from typing import Any

import numpy as np

from ...utils.types import TupleYUV
from ._nd_image import NDImageSlicer
from ._async import run_in_thread_pool




class VideoSlicer(NDImageSlicer):
    """
    ``NDImageSlicer`` subclass for video data, used by ``NDWSubplot.add_video()``.

    Reads the frame at the current index directly. Window functions are not currently implemented for video.

    A YUV frame is a tuple of (Y, U, V) planes rather than a single array, so it is passed through as a tuple for
    an ``ImageYUVGraphic``.
    """



    async def get_window_output(self, indices: dict[str, Any]) -> TupleYUV | np.ndarray:
        """
        Get the frame at the given indices, squeezing out the slider dims.

        Parameters
        ----------
        indices: dict[str, Any]
            Reference-space value for each slider dim, ex: ``{"time": 46.397}``. Must provide a value for every
            slider dim.

        Returns
        -------
        np.ndarray | tuple[np.ndarray, ...]
            The frame, or a tuple of the (Y, U, V) planes if the underlying data returns YUV planes.

        """
        # windowed slice if user set any window funcs
        windowed_slice = await self._get_raw_data_slice(indices)

        if isinstance(windowed_slice, (tuple, list)):
            return tuple(a.squeeze() for a in windowed_slice)

        # convert to numpy array
        return np.asarray(windowed_slice).squeeze()




    async def get(self, indices: dict[str, Any]) -> TupleYUV | np.ndarray:
        """
        Similar to NDImage.get() but accounts for TupleYUV output.
        """
        # squeezed output, dims in array order
        window_output = await self.get_window_output(indices)

        # transpose into display order, the spatial_func gets the frame as it is rendered
        if isinstance(window_output, tuple):
            window_output = tuple(
                a.transpose(*self.display_dims_indices) for a in window_output
            )
        else:
            window_output = window_output.transpose(*self.display_dims_indices)

        if self.spatial_func is not None:
            window_output = await run_in_thread_pool(
                self._executor, self._spatial_func, window_output
            )

        return window_output