_meta.py 10.4 KB
Newer Older
1
from typing import List, Optional, Tuple, Union
2

3
4
import PIL.Image
import torch
5
6
from torchvision import datapoints
from torchvision.datapoints import BoundingBoxFormat
7
from torchvision.transforms import _functional_pil as _FP
8

9
10
from torchvision.utils import _log_api_usage_once

11
12
from ._utils import is_simple_tensor

13

14
15
16
17
18
19
20
21
22
23
24
25
def get_dimensions_image_tensor(image: torch.Tensor) -> List[int]:
    chw = list(image.shape[-3:])
    ndims = len(chw)
    if ndims == 3:
        return chw
    elif ndims == 2:
        chw.insert(0, 1)
        return chw
    else:
        raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}")


26
27
28
get_dimensions_image_pil = _FP.get_dimensions


Philip Meier's avatar
Philip Meier committed
29
def get_dimensions(inpt: Union[datapoints._ImageTypeJIT, datapoints._VideoTypeJIT]) -> List[int]:
30
31
32
    if not torch.jit.is_scripting():
        _log_api_usage_once(get_dimensions)

33
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
34
        return get_dimensions_image_tensor(inpt)
35
    elif isinstance(inpt, (datapoints.Image, datapoints.Video)):
36
37
        channels = inpt.num_channels
        height, width = inpt.spatial_size
38
        return [channels, height, width]
39
40
    elif isinstance(inpt, PIL.Image.Image):
        return get_dimensions_image_pil(inpt)
41
    else:
42
        raise TypeError(
43
            f"Input can either be a plain tensor, an `Image` or `Video` datapoint, or a PIL image, "
44
45
            f"but got {type(inpt)} instead."
        )
46
47


48
49
50
51
52
53
54
55
56
57
58
def get_num_channels_image_tensor(image: torch.Tensor) -> int:
    chw = image.shape[-3:]
    ndims = len(chw)
    if ndims == 3:
        return chw[0]
    elif ndims == 2:
        return 1
    else:
        raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}")


59
get_num_channels_image_pil = _FP.get_image_num_channels
60
61


62
63
64
65
def get_num_channels_video(video: torch.Tensor) -> int:
    return get_num_channels_image_tensor(video)


Philip Meier's avatar
Philip Meier committed
66
def get_num_channels(inpt: Union[datapoints._ImageTypeJIT, datapoints._VideoTypeJIT]) -> int:
67
68
69
    if not torch.jit.is_scripting():
        _log_api_usage_once(get_num_channels)

70
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
71
        return get_num_channels_image_tensor(inpt)
72
    elif isinstance(inpt, (datapoints.Image, datapoints.Video)):
73
74
75
        return inpt.num_channels
    elif isinstance(inpt, PIL.Image.Image):
        return get_num_channels_image_pil(inpt)
76
    else:
77
        raise TypeError(
78
            f"Input can either be a plain tensor, an `Image` or `Video` datapoint, or a PIL image, "
79
80
            f"but got {type(inpt)} instead."
        )
81
82


83
84
85
86
87
# We changed the names to ensure it can be used not only for images but also videos. Thus, we just alias it without
# deprecating the old names.
get_image_num_channels = get_num_channels


88
def get_spatial_size_image_tensor(image: torch.Tensor) -> List[int]:
89
90
91
92
93
94
    hw = list(image.shape[-2:])
    ndims = len(hw)
    if ndims == 2:
        return hw
    else:
        raise TypeError(f"Input tensor should have at least two dimensions, but got {ndims}")
95
96
97
98
99
100
101
102


@torch.jit.unused
def get_spatial_size_image_pil(image: PIL.Image.Image) -> List[int]:
    width, height = _FP.get_image_size(image)
    return [height, width]


103
104
105
106
107
108
109
110
111
def get_spatial_size_video(video: torch.Tensor) -> List[int]:
    return get_spatial_size_image_tensor(video)


def get_spatial_size_mask(mask: torch.Tensor) -> List[int]:
    return get_spatial_size_image_tensor(mask)


@torch.jit.unused
112
def get_spatial_size_bounding_box(bounding_box: datapoints.BoundingBox) -> List[int]:
113
    return list(bounding_box.spatial_size)
114
115


Philip Meier's avatar
Philip Meier committed
116
def get_spatial_size(inpt: datapoints._InputTypeJIT) -> List[int]:
117
118
119
    if not torch.jit.is_scripting():
        _log_api_usage_once(get_spatial_size)

120
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
121
        return get_spatial_size_image_tensor(inpt)
122
    elif isinstance(inpt, (datapoints.Image, datapoints.Video, datapoints.BoundingBox, datapoints.Mask)):
123
        return list(inpt.spatial_size)
124
    elif isinstance(inpt, PIL.Image.Image):
125
        return get_spatial_size_image_pil(inpt)
126
127
    else:
        raise TypeError(
128
            f"Input can either be a plain tensor, any TorchVision datapoint, or a PIL image, "
129
130
            f"but got {type(inpt)} instead."
        )
131
132
133
134
135
136


def get_num_frames_video(video: torch.Tensor) -> int:
    return video.shape[-4]


Philip Meier's avatar
Philip Meier committed
137
def get_num_frames(inpt: datapoints._VideoTypeJIT) -> int:
138
139
140
    if not torch.jit.is_scripting():
        _log_api_usage_once(get_num_frames)

141
    if torch.jit.is_scripting() or is_simple_tensor(inpt):
142
        return get_num_frames_video(inpt)
143
    elif isinstance(inpt, datapoints.Video):
144
        return inpt.num_frames
145
    else:
146
        raise TypeError(f"Input can either be a plain tensor or a `Video` datapoint, but got {type(inpt)} instead.")
147
148


149
150
def _xywh_to_xyxy(xywh: torch.Tensor, inplace: bool) -> torch.Tensor:
    xyxy = xywh if inplace else xywh.clone()
151
152
153
154
    xyxy[..., 2:] += xyxy[..., :2]
    return xyxy


155
156
def _xyxy_to_xywh(xyxy: torch.Tensor, inplace: bool) -> torch.Tensor:
    xywh = xyxy if inplace else xyxy.clone()
157
158
159
160
    xywh[..., 2:] -= xywh[..., :2]
    return xywh


161
162
163
def _cxcywh_to_xyxy(cxcywh: torch.Tensor, inplace: bool) -> torch.Tensor:
    if not inplace:
        cxcywh = cxcywh.clone()
164

165
166
167
168
169
170
171
    # Trick to do fast division by 2 and ceil, without casting. It produces the same result as
    # `torchvision.ops._box_convert._box_cxcywh_to_xyxy`.
    half_wh = cxcywh[..., 2:].div(-2, rounding_mode=None if cxcywh.is_floating_point() else "floor").abs_()
    # (cx - width / 2) = x1, same for y1
    cxcywh[..., :2].sub_(half_wh)
    # (x1 + width) = x2, same for y2
    cxcywh[..., 2:].add_(cxcywh[..., :2])
172

173
174
175
176
177
178
179
180
181
182
183
184
185
    return cxcywh


def _xyxy_to_cxcywh(xyxy: torch.Tensor, inplace: bool) -> torch.Tensor:
    if not inplace:
        xyxy = xyxy.clone()

    # (x2 - x1) = width, same for height
    xyxy[..., 2:].sub_(xyxy[..., :2])
    # (x1 * 2 + width) / 2 = x1 + width / 2 = x1 + (x2-x1)/2 = (x1 + x2)/2 = cx, same for cy
    xyxy[..., :2].mul_(2).add_(xyxy[..., 2:]).div_(2, rounding_mode=None if xyxy.is_floating_point() else "floor")

    return xyxy
186
187


188
def _convert_format_bounding_box(
189
    bounding_box: torch.Tensor, old_format: BoundingBoxFormat, new_format: BoundingBoxFormat, inplace: bool = False
190
) -> torch.Tensor:
191

192
    if new_format == old_format:
193
        return bounding_box
194

195
    # TODO: Add _xywh_to_cxcywh and _cxcywh_to_xywh to improve performance
196
    if old_format == BoundingBoxFormat.XYWH:
197
        bounding_box = _xywh_to_xyxy(bounding_box, inplace)
198
    elif old_format == BoundingBoxFormat.CXCYWH:
199
        bounding_box = _cxcywh_to_xyxy(bounding_box, inplace)
200
201

    if new_format == BoundingBoxFormat.XYWH:
202
        bounding_box = _xyxy_to_xywh(bounding_box, inplace)
203
    elif new_format == BoundingBoxFormat.CXCYWH:
204
        bounding_box = _xyxy_to_cxcywh(bounding_box, inplace)
205
206
207
208

    return bounding_box


209
def convert_format_bounding_box(
Philip Meier's avatar
Philip Meier committed
210
    inpt: datapoints._InputTypeJIT,
211
212
213
    old_format: Optional[BoundingBoxFormat] = None,
    new_format: Optional[BoundingBoxFormat] = None,
    inplace: bool = False,
Philip Meier's avatar
Philip Meier committed
214
) -> datapoints._InputTypeJIT:
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
    # This being a kernel / dispatcher hybrid, we need an option to pass `old_format` explicitly for simple tensor
    # inputs as well as extract it from `datapoints.BoundingBox` inputs. However, putting a default value on
    # `old_format` means we also need to put one on `new_format` to have syntactically correct Python. Here we mimic the
    # default error that would be thrown if `new_format` had no default value.
    if new_format is None:
        raise TypeError("convert_format_bounding_box() missing 1 required argument: 'new_format'")

    if not torch.jit.is_scripting():
        _log_api_usage_once(convert_format_bounding_box)

    if torch.jit.is_scripting() or is_simple_tensor(inpt):
        if old_format is None:
            raise ValueError("For simple tensor inputs, `old_format` has to be passed.")
        return _convert_format_bounding_box(inpt, old_format=old_format, new_format=new_format, inplace=inplace)
    elif isinstance(inpt, datapoints.BoundingBox):
        if old_format is not None:
            raise ValueError("For bounding box datapoint inputs, `old_format` must not be passed.")
232
233
234
235
        output = _convert_format_bounding_box(
            inpt.as_subclass(torch.Tensor), old_format=inpt.format, new_format=new_format, inplace=inplace
        )
        return datapoints.BoundingBox.wrap_like(inpt, output, format=new_format)
236
237
238
239
240
241
    else:
        raise TypeError(
            f"Input can either be a plain tensor or a bounding box datapoint, but got {type(inpt)} instead."
        )


242
def _clamp_bounding_box(
243
    bounding_box: torch.Tensor, format: BoundingBoxFormat, spatial_size: Tuple[int, int]
244
) -> torch.Tensor:
245
246
    # TODO: Investigate if it makes sense from a performance perspective to have an implementation for every
    #  BoundingBoxFormat instead of converting back and forth
247
248
    in_dtype = bounding_box.dtype
    bounding_box = bounding_box.clone() if bounding_box.is_floating_point() else bounding_box.float()
249
    xyxy_boxes = convert_format_bounding_box(
250
        bounding_box, old_format=format, new_format=datapoints.BoundingBoxFormat.XYXY, inplace=True
251
    )
252
253
    xyxy_boxes[..., 0::2].clamp_(min=0, max=spatial_size[1])
    xyxy_boxes[..., 1::2].clamp_(min=0, max=spatial_size[0])
254
255
256
257
    out_boxes = convert_format_bounding_box(
        xyxy_boxes, old_format=BoundingBoxFormat.XYXY, new_format=format, inplace=True
    )
    return out_boxes.to(in_dtype)
258
259


260
def clamp_bounding_box(
Philip Meier's avatar
Philip Meier committed
261
    inpt: datapoints._InputTypeJIT,
262
263
    format: Optional[BoundingBoxFormat] = None,
    spatial_size: Optional[Tuple[int, int]] = None,
Philip Meier's avatar
Philip Meier committed
264
) -> datapoints._InputTypeJIT:
265
266
267
268
269
270
271
272
273
274
    if not torch.jit.is_scripting():
        _log_api_usage_once(clamp_bounding_box)

    if torch.jit.is_scripting() or is_simple_tensor(inpt):
        if format is None or spatial_size is None:
            raise ValueError("For simple tensor inputs, `format` and `spatial_size` has to be passed.")
        return _clamp_bounding_box(inpt, format=format, spatial_size=spatial_size)
    elif isinstance(inpt, datapoints.BoundingBox):
        if format is not None or spatial_size is not None:
            raise ValueError("For bounding box datapoint inputs, `format` and `spatial_size` must not be passed.")
275
        output = _clamp_bounding_box(inpt.as_subclass(torch.Tensor), format=inpt.format, spatial_size=inpt.spatial_size)
276
277
278
279
280
        return datapoints.BoundingBox.wrap_like(inpt, output)
    else:
        raise TypeError(
            f"Input can either be a plain tensor or a bounding box datapoint, but got {type(inpt)} instead."
        )