"vscode:/vscode.git/clone" did not exist on "88f9c347b288733c87b3056b6e1a8d2717036ce6"
_meta.py 2.6 KB
Newer Older
1
from typing import Any, Dict, Union
Philip Meier's avatar
Philip Meier committed
2

3
import torch
4

5
6
from torchvision import datapoints, transforms as _transforms
from torchvision.transforms.v2 import functional as F, Transform
7

8
9
from .utils import is_simple_tensor

10
11

class ConvertBoundingBoxFormat(Transform):
12
    _transformed_types = (datapoints.BoundingBox,)
13

14
    def __init__(self, format: Union[str, datapoints.BoundingBoxFormat]) -> None:
15
16
        super().__init__()
        if isinstance(format, str):
17
            format = datapoints.BoundingBoxFormat[format]
18
19
        self.format = format

20
    def _transform(self, inpt: datapoints.BoundingBox, params: Dict[str, Any]) -> datapoints.BoundingBox:
21
        return F.convert_format_bounding_box(inpt, new_format=self.format)  # type: ignore[return-value]
22
23


24
class ConvertDtype(Transform):
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
    """[BETA] Convert a tensor image/box/mask to the given ``dtype`` and scale the values accordingly

    .. betastatus:: ConvertDtype transform

    This function does not support PIL Image.

    Args:
        dtype (torch.dtype): Desired data type of the output

    .. note::

        When converting from a smaller to a larger integer ``dtype`` the maximum values are **not** mapped exactly.
        If converted back and forth, this mismatch has no effect.

    Raises:
        RuntimeError: When trying to cast :class:`torch.float32` to :class:`torch.int32` or :class:`torch.int64` as
            well as for trying to cast :class:`torch.float64` to :class:`torch.int64`. These conversions might lead to
            overflow errors since the floating point ``dtype`` cannot store consecutive integers over the whole range
            of the integer ``dtype``.
    """

46
47
    _v1_transform_cls = _transforms.ConvertImageDtype

48
    _transformed_types = (is_simple_tensor, datapoints.Image, datapoints.Video)
49

50
51
52
53
    def __init__(self, dtype: torch.dtype = torch.float32) -> None:
        super().__init__()
        self.dtype = dtype

54
    def _transform(
Philip Meier's avatar
Philip Meier committed
55
56
        self, inpt: Union[datapoints._TensorImageType, datapoints._TensorVideoType], params: Dict[str, Any]
    ) -> Union[datapoints._TensorImageType, datapoints._TensorVideoType]:
57
58
59
60
61
62
        return F.convert_dtype(inpt, self.dtype)


# We changed the name to align it with the new naming scheme. Still, `ConvertImageDtype` is
# prevalent and well understood. Thus, we just alias it without deprecating the old name.
ConvertImageDtype = ConvertDtype
63
64


65
class ClampBoundingBox(Transform):
66
    _transformed_types = (datapoints.BoundingBox,)
67

68
    def _transform(self, inpt: datapoints.BoundingBox, params: Dict[str, Any]) -> datapoints.BoundingBox:
69
        return F.clamp_bounding_box(inpt)  # type: ignore[return-value]