Source code for torchaug.ta_tensors._batch_videos

# @Copyright: CEA-LIST/DIASI/SIALV/ (2023-    )
# @Author: CEA-LIST/DIASI/SIALV/ <julien.denize@cea.fr>
# @License: CECILL-C

from __future__ import annotations

from typing import Any, Optional, Sequence, Union

import torch
from torch._C import device, dtype

from ._ta_tensor import TATensor
from ._video import Video




class BatchVideos(TATensor):
    """:class:`torch.Tensor` subclass for batch of videos.

    Args:
        data: Any data that can be turned into a tensor with :func:`torch.as_tensor`.
        dtype: Desired data type. If omitted, will be inferred from
            ``data``.
        device: Desired device. If omitted and ``data`` is a
            :class:`torch.Tensor`, the device is taken from it. Otherwise, the video is constructed on the CPU.
        requires_grad: Whether autograd should record operations. If omitted and
            ``data`` is a :class:`torch.Tensor`, the value is taken from it. Otherwise, defaults to ``False``.
    """

    def __new__(
        cls,
        data: Any,
        *,
        dtype: Optional[dtype] = None,
        device: Optional[Union[device, str, int]] = None,
        requires_grad: Optional[bool] = None,
    ) -> BatchVideos:
        tensor = cls._to_tensor(data, dtype=dtype, device=device, requires_grad=requires_grad)
        if data.ndim < 5:
            raise ValueError
        return tensor.as_subclass(cls)



    @classmethod
    def cat(cls, video_batches: Sequence[BatchVideos]) -> BatchVideos:
        """Concatenates the sequence of images batches into a single batch."""
        attrs = [
            "requires_grad",
            "device",
            "dtype",
        ]

        for batch_videos in video_batches:
            if not isinstance(batch_videos, BatchVideos):
                raise ValueError("All batches must be of type BatchVideos.")
            for attr in attrs:
                if getattr(batch_videos, attr) != getattr(video_batches[0], attr):
                    raise ValueError(f"All batches of images must have the same {attr} attribute.")

        return cls(torch.cat([images.data for images in video_batches], 0))




    def to_samples(self) -> list[Video]:
        """Get the tensors."""
        return [Video(video.clone()) for video in self]


    def __repr__(self, *, tensor_contents: Any = None) -> str:  # type: ignore[override]
        return self._make_repr()