Source code for torchaug.ta_tensors.nested._bounding_boxes

# ==================================
# Copyright: CEA-LIST/DIASI/SIALV/
# Author : Torchaug Developers
# License: CECILL-C
# ==================================

from __future__ import annotations

from typing import (
    List,
)

from torchaug.ta_tensors import (
    BatchBoundingBoxes,
    BoundingBoxes,
)
from torchaug.ta_tensors._batch_bounding_boxes import convert_bboxes_to_batch_bboxes

from ._ta_nested_tensors import TANestedTensors



[docs]
class BoundingBoxesNestedTensors(TANestedTensors[BoundingBoxes, BatchBoundingBoxes]):
    """Implement BoundingBoxes Nested Tensor for PyTorch."""

    tensors_type = BoundingBoxes
    batch_tensors_type = BatchBoundingBoxes
    tensors: List[BoundingBoxes]


[docs]
    def to_batch(self) -> BatchBoundingBoxes:
        """Return the batched tensor and the shapes of the nested tensors."""
        return convert_bboxes_to_batch_bboxes(self.tensors)