File size: 9,768 Bytes
36c95ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 | from functools import wraps
from typing import Callable, cast, List, Optional, Tuple, Union
import torch
from torch.distributions import Beta, Uniform
from kornia.utils import _extract_device_dtype
def _validate_input(f: Callable) -> Callable:
r"""Validate the 2D input of the wrapped function.
Args:
f: a function that takes the first argument as tensor.
Returns:
the wrapped function after input is validated.
"""
@wraps(f)
def wrapper(input: torch.Tensor, *args, **kwargs):
if not torch.is_tensor(input):
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
_validate_shape(input.shape, required_shapes=('BCHW',))
_validate_input_dtype(input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
return f(input, *args, **kwargs)
return wrapper
def _validate_input3d(f: Callable) -> Callable:
r"""Validate the 3D input of the wrapped function.
Args:
f: a function that takes the first argument as tensor.
Returns:
the wrapped function after input is validated.
"""
@wraps(f)
def wrapper(input: torch.Tensor, *args, **kwargs):
if not torch.is_tensor(input):
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
input_shape = len(input.shape)
if input_shape != 5:
raise AssertionError(f'Expect input of 5 dimensions, got {input_shape} instead')
_validate_input_dtype(input, accepted_dtypes=[torch.float16, torch.float32, torch.float64])
return f(input, *args, **kwargs)
return wrapper
def _infer_batch_shape(input: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) -> torch.Size:
r"""Infer input shape. Input may be either (tensor,) or (tensor, transform_matrix)"""
if isinstance(input, tuple):
tensor = _transform_input(input[0])
else:
tensor = _transform_input(input)
return tensor.shape
def _infer_batch_shape3d(input: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]) -> torch.Size:
r"""Infer input shape. Input may be either (tensor,) or (tensor, transform_matrix)"""
if isinstance(input, tuple):
tensor = _transform_input3d(input[0])
else:
tensor = _transform_input3d(input)
return tensor.shape
def _transform_input(input: torch.Tensor) -> torch.Tensor:
r"""Reshape an input tensor to be (*, C, H, W). Accept either (H, W), (C, H, W) or (*, C, H, W).
Args:
input: torch.Tensor
Returns:
torch.Tensor
"""
if not torch.is_tensor(input):
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
if len(input.shape) not in [2, 3, 4]:
raise ValueError(f"Input size must have a shape of either (H, W), (C, H, W) or (*, C, H, W). Got {input.shape}")
if len(input.shape) == 2:
input = input.unsqueeze(0)
if len(input.shape) == 3:
input = input.unsqueeze(0)
return input
def _transform_input3d(input: torch.Tensor) -> torch.Tensor:
r"""Reshape an input tensor to be (*, C, D, H, W). Accept either (D, H, W), (C, D, H, W) or (*, C, D, H, W).
Args:
input: torch.Tensor
Returns:
torch.Tensor
"""
if not torch.is_tensor(input):
raise TypeError(f"Input type is not a torch.Tensor. Got {type(input)}")
if len(input.shape) not in [3, 4, 5]:
raise ValueError(
f"Input size must have a shape of either (D, H, W), (C, D, H, W) or (*, C, D, H, W). Got {input.shape}"
)
if len(input.shape) == 3:
input = input.unsqueeze(0)
if len(input.shape) == 4:
input = input.unsqueeze(0)
return input
def _validate_input_dtype(input: torch.Tensor, accepted_dtypes: List) -> None:
r"""Check if the dtype of the input tensor is in the range of accepted_dtypes
Args:
input: torch.Tensor
accepted_dtypes: List. e.g. [torch.float32, torch.float64]
"""
if input.dtype not in accepted_dtypes:
raise TypeError(f"Expected input of {accepted_dtypes}. Got {input.dtype}")
def _transform_output_shape(
output: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], shape: Tuple
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
r"""Collapse the broadcasted batch dimensions an input tensor to be the specified shape.
Args:
input: torch.Tensor
shape: List/tuple of int
Returns:
torch.Tensor
"""
is_tuple = isinstance(output, tuple)
out_tensor: torch.Tensor
trans_matrix: Optional[torch.Tensor]
if is_tuple:
out_tensor, trans_matrix = cast(Tuple[torch.Tensor, torch.Tensor], output)
else:
out_tensor = cast(torch.Tensor, output)
trans_matrix = None
if trans_matrix is not None:
if len(out_tensor.shape) > len(shape) and trans_matrix.shape[0] != 1:
raise AssertionError(
f'Dimension 0 of transformation matrix is ' f'expected to be 1, got {trans_matrix.shape[0]}'
)
trans_matrix = trans_matrix.squeeze(0)
for dim in range(len(out_tensor.shape) - len(shape)):
if out_tensor.shape[0] != 1:
raise AssertionError(f'Dimension {dim} of input is ' f'expected to be 1, got {out_tensor.shape[0]}')
out_tensor = out_tensor.squeeze(0)
return (out_tensor, trans_matrix) if is_tuple else out_tensor # type: ignore
def _validate_shape(shape: Union[Tuple, torch.Size], required_shapes: Tuple[str, ...] = ("BCHW",)) -> None:
r"""Check if the dtype of the input tensor is in the range of accepted_dtypes
Args:
shape: tensor shape
required_shapes: List. e.g. ["BCHW", "BCDHW"]
"""
passed = False
for required_shape in required_shapes:
if len(shape) == len(required_shape):
passed = True
break
if not passed:
raise TypeError(f"Expected input shape in {required_shape}. Got {shape}.")
def _validate_input_shape(input: torch.Tensor, channel_index: int, number: int) -> bool:
r"""Validate if an input has the right shape. e.g. to check if an input is channel first.
If channel first, the second channel of an RGB input shall be fixed to 3. To verify using:
_validate_input_shape(input, 1, 3)
Args:
input: torch.Tensor
channel_index: int
number: int
Returns:
bool
"""
return input.shape[channel_index] == number
def _adapted_rsampling(
shape: Union[Tuple, torch.Size], dist: torch.distributions.Distribution, same_on_batch=False
) -> torch.Tensor:
r"""The uniform reparameterized sampling function that accepts 'same_on_batch'.
If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]).
By default, same_on_batch is set to False.
"""
if same_on_batch:
return dist.rsample((1, *shape[1:])).repeat(shape[0], *[1] * (len(shape) - 1))
return dist.rsample(shape)
def _adapted_sampling(
shape: Union[Tuple, torch.Size], dist: torch.distributions.Distribution, same_on_batch=False
) -> torch.Tensor:
r"""The uniform sampling function that accepts 'same_on_batch'.
If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]).
By default, same_on_batch is set to False.
"""
if same_on_batch:
return dist.sample((1, *shape[1:])).repeat(shape[0], *[1] * (len(shape) - 1))
return dist.sample(shape)
def _adapted_uniform(
shape: Union[Tuple, torch.Size],
low: Union[float, int, torch.Tensor],
high: Union[float, int, torch.Tensor],
same_on_batch: bool = False,
) -> torch.Tensor:
r"""The uniform sampling function that accepts 'same_on_batch'.
If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]).
By default, same_on_batch is set to False.
By default, sampling happens on the default device and dtype. If low/high is a tensor, sampling will happen
in the same device/dtype as low/high tensor.
"""
device, dtype = _extract_device_dtype(
[low if isinstance(low, torch.Tensor) else None, high if isinstance(high, torch.Tensor) else None]
)
low = torch.as_tensor(low, device=device, dtype=dtype)
high = torch.as_tensor(high, device=device, dtype=dtype)
# validate_args=False to fix pytorch 1.7.1 error:
# ValueError: Uniform is not defined when low>= high.
dist = Uniform(low, high, validate_args=False)
return _adapted_rsampling(shape, dist, same_on_batch)
def _adapted_beta(
shape: Union[Tuple, torch.Size],
a: Union[float, int, torch.Tensor],
b: Union[float, int, torch.Tensor],
same_on_batch: bool = False,
) -> torch.Tensor:
r"""The beta sampling function that accepts 'same_on_batch'.
If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]).
By default, same_on_batch is set to False.
By default, sampling happens on the default device and dtype. If a/b is a tensor, sampling will happen
in the same device/dtype as a/b tensor.
"""
device, dtype = _extract_device_dtype(
[a if isinstance(a, torch.Tensor) else None, b if isinstance(b, torch.Tensor) else None]
)
a = torch.as_tensor(a, device=device, dtype=dtype)
b = torch.as_tensor(b, device=device, dtype=dtype)
dist = Beta(a, b, validate_args=False)
return _adapted_rsampling(shape, dist, same_on_batch)
def _shape_validation(param: torch.Tensor, shape: Union[tuple, list], name: str) -> None:
if param.shape != torch.Size(shape):
raise AssertionError(f"Invalid shape for {name}. Expected {shape}. Got {param.shape}")
|