Source code for araviq6.array2qvideoframe

"""
Array-frame conversion
======================

:mod:`araviq6.array2qvideoframe` provides functions to convert numpy array to
``QVideoFrame``.

To convert ``QVideoFrame`` to numpy array, convert the frame to ``QImage`` by
``QVideoFrame.toImage()`` and use :mod:`qimage2ndarray` package.

.. note::
   This module imitates https://github.com/hmeine/qimage2ndarray.

.. autofunction:: array2qvideoframe

"""

import numpy as np
import numpy.typing as npt
import sys
from qimage2ndarray import _normalize255  # type: ignore[import]
from araviq6.qt_compat import QtCore, QtMultimedia, get_frame_data
from typing import Optional, Union, Tuple


__all__ = [
    "array2qvideoframe",
]


class ArrayInterfaceAroundQVideoFrame(object):
    __slots__ = ("__qvideoframe", "__array_interface__")

    def __init__(self, frame: QtMultimedia.QVideoFrame, bytes_per_pixel: int):
        self.__qvideoframe = frame

        self.__array_interface__ = dict(
            shape=(frame.height(), frame.width()),
            typestr="|u%d" % bytes_per_pixel,
            data=get_frame_data(frame),
            strides=(frame.bytesPerLine(0), bytes_per_pixel),
            version=3,
        )


def qvideoframeview(frame: QtMultimedia.QVideoFrame) -> np.ndarray:
    pixelFormat = frame.surfaceFormat().pixelFormat()
    if pixelFormat == QtMultimedia.QVideoFrameFormat.PixelFormat.Format_BGRA8888:
        bits = 32
    elif pixelFormat == QtMultimedia.QVideoFrameFormat.PixelFormat.Format_BGRX8888:
        bits = 32
    else:
        raise TypeError(f"Invalid pixel format: {pixelFormat}")
    interface = ArrayInterfaceAroundQVideoFrame(frame, bits // 8)
    return np.asarray(interface)


def byte_view(
    frame: QtMultimedia.QVideoFrame, byteorder: Optional[str] = "little"
) -> npt.NDArray[np.uint8]:
    raw = qvideoframeview(frame)
    result = raw.view(np.uint8).reshape(raw.shape + (-1,))
    if byteorder is not None and byteorder != sys.byteorder:
        result = result[..., ::-1]
    return result


def rgb_view(
    frame: QtMultimedia.QVideoFrame, byteorder: Optional[str] = "big"
) -> npt.NDArray[np.uint8]:
    if byteorder is None:
        byteorder = sys.byteorder
    bytes = byte_view(frame, byteorder)

    if byteorder == "little":
        result = bytes[..., :3]  # strip A off BGRA
    else:
        result = bytes[..., 1:]  # strip A off ARGB
    return result


def alpha_view(frame: QtMultimedia.QVideoFrame) -> npt.NDArray[np.uint8]:
    bytes = byte_view(frame, byteorder=None)
    if sys.byteorder == "little":
        ret = bytes[..., 3]
    else:
        ret = bytes[..., 0]
    return ret


[docs] def array2qvideoframe( array: np.ndarray, normalize: Union[bool, int, Tuple[int, int]] = False ) -> QtMultimedia.QVideoFrame: """ Convert a 2D or 3D numpy array into 32-bit ``QVideoFrame``. The dimensions of a 3D array are ``(width, height, channels)``, and the channels can be 1, 2, 3 or 4. 2D array with ``(width, height)`` dimension is converted to ``(width, height, 1)``. Number of the channels is interpreted as follows: ========= =================== #channels interpretation ========= =================== 1 scalar/gray 2 scalar/gray + alpha 3 RGB 4 RGB + alpha ========= =================== Note that the scalar data will be converted into gray RGB triples. The parameter *normalize* can be used to normalize an frame's value range to 0-255. If *normalize* = ``(nmin, nmax)``: Scale & clip frame values from ``nmin..nmax`` to ``0..255`` If *normalize* = ``nmax``: Lets ``nmin`` default to zero, i.e. scale & clip the range ``0..nmax`` to ``0..255`` If *normalize* = ``True``: Scale frame values to ``0..255``, except for boolean arays, where ``False`` and ``True`` are mapped to ``0`` and ``255``. Same as passing ``(gray.min(), gray.max())`` If `array` contains masked values, the corresponding pixels will be transparent in the result. Thus, the result be of ``Format_BGRA8888`` if the input already contains an alhpa channel (i.e., has shape ``(H, W, 4)``) or if there are masked pixels, and ``Format_BGRX8888`` otherwise. """ dim = np.ndim(array) if dim == 2: array = array[..., None] elif dim != 3: raise ValueError( f"invalid number of dimensions for array2qvideoframe (got {dim} dimensions)" ) h, w, ch = array.shape if ch not in (1, 2, 3, 4): raise ValueError( f"invalid number of channels for array2qvideoframe: (got {ch} channels)" ) hasAlpha = np.ma.is_masked(array) or ch in (2, 4) if hasAlpha: pixelFormat = QtMultimedia.QVideoFrameFormat.PixelFormat.Format_BGRA8888 else: pixelFormat = QtMultimedia.QVideoFrameFormat.PixelFormat.Format_BGRX8888 frameFormat = QtMultimedia.QVideoFrameFormat(QtCore.QSize(w, h), pixelFormat) frame = QtMultimedia.QVideoFrame(frameFormat) frame.map(QtMultimedia.QVideoFrame.MapMode.WriteOnly) array = _normalize255(array, normalize) rgb = rgb_view(frame) if ch >= 3: rgb[:] = array[..., :3] else: rgb[:] = array[..., :1] # scalar data alpha = alpha_view(frame) if ch in (2, 4): alpha[:] = array[..., -1] else: alpha[:] = 255 if np.ma.is_masked(array): alpha[:] *= np.logical_not( np.any(array.mask, axis=-1) # type: ignore[attr-defined] ) frame.unmap() # save the modified memory to frame instance return frame