| | """Adapted from https://github.com/cvg/GeoCalib""" |
| |
|
| | import collections.abc as collections |
| | import functools |
| | import inspect |
| | from typing import Callable, List, Tuple |
| |
|
| | import numpy as np |
| | import torch |
| |
|
| | |
| | |
| |
|
| |
|
| | string_classes = (str, bytes) |
| |
|
| |
|
| | def autocast(func: Callable) -> Callable: |
| | """Cast the inputs of a TensorWrapper method to PyTorch tensors if they are numpy arrays. |
| | |
| | Use the device and dtype of the wrapper. |
| | |
| | Args: |
| | func (Callable): Method of a TensorWrapper class. |
| | |
| | Returns: |
| | Callable: Wrapped method. |
| | """ |
| |
|
| | @functools.wraps(func) |
| | def wrap(self, *args): |
| | device = torch.device("cpu") |
| | dtype = None |
| | if isinstance(self, TensorWrapper): |
| | if self._data is not None: |
| | device = self.device |
| | dtype = self.dtype |
| | elif not inspect.isclass(self) or not issubclass(self, TensorWrapper): |
| | raise ValueError(self) |
| |
|
| | cast_args = [] |
| | for arg in args: |
| | if isinstance(arg, np.ndarray): |
| | arg = torch.from_numpy(arg) |
| | arg = arg.to(device=device, dtype=dtype) |
| | cast_args.append(arg) |
| | return func(self, *cast_args) |
| |
|
| | return wrap |
| |
|
| |
|
| | class TensorWrapper: |
| | """Wrapper for PyTorch tensors.""" |
| |
|
| | _data = None |
| |
|
| | @autocast |
| | def __init__(self, data: torch.Tensor): |
| | """Wrapper for PyTorch tensors.""" |
| | self._data = data |
| |
|
| | @property |
| | def shape(self) -> torch.Size: |
| | """Shape of the underlying tensor.""" |
| | return self._data.shape[:-1] |
| |
|
| | @property |
| | def device(self) -> torch.device: |
| | """Get the device of the underlying tensor.""" |
| | return self._data.device |
| |
|
| | @property |
| | def dtype(self) -> torch.dtype: |
| | """Get the dtype of the underlying tensor.""" |
| | return self._data.dtype |
| |
|
| | def __getitem__(self, index) -> torch.Tensor: |
| | """Get the underlying tensor.""" |
| | return self.__class__(self._data[index]) |
| |
|
| | def __setitem__(self, index, item): |
| | """Set the underlying tensor.""" |
| | self._data[index] = item.data |
| |
|
| | def to(self, *args, **kwargs): |
| | """Move the underlying tensor to a new device.""" |
| | return self.__class__(self._data.to(*args, **kwargs)) |
| |
|
| | def cpu(self): |
| | """Move the underlying tensor to the CPU.""" |
| | return self.__class__(self._data.cpu()) |
| |
|
| | def cuda(self): |
| | """Move the underlying tensor to the GPU.""" |
| | return self.__class__(self._data.cuda()) |
| |
|
| | def pin_memory(self): |
| | """Pin the underlying tensor to memory.""" |
| | return self.__class__(self._data.pin_memory()) |
| |
|
| | def float(self): |
| | """Cast the underlying tensor to float.""" |
| | return self.__class__(self._data.float()) |
| |
|
| | def double(self): |
| | """Cast the underlying tensor to double.""" |
| | return self.__class__(self._data.double()) |
| |
|
| | def detach(self): |
| | """Detach the underlying tensor.""" |
| | return self.__class__(self._data.detach()) |
| |
|
| | def numpy(self): |
| | """Convert the underlying tensor to a numpy array.""" |
| | return self._data.detach().cpu().numpy() |
| |
|
| | def new_tensor(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self._data.new_tensor(*args, **kwargs) |
| |
|
| | def new_zeros(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self._data.new_zeros(*args, **kwargs) |
| |
|
| | def new_ones(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self._data.new_ones(*args, **kwargs) |
| |
|
| | def new_full(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self._data.new_full(*args, **kwargs) |
| |
|
| | def new_empty(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self._data.new_empty(*args, **kwargs) |
| |
|
| | def unsqueeze(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self.__class__(self._data.unsqueeze(*args, **kwargs)) |
| |
|
| | def squeeze(self, *args, **kwargs): |
| | """Create a new tensor of the same type and device.""" |
| | return self.__class__(self._data.squeeze(*args, **kwargs)) |
| |
|
| | @classmethod |
| | def stack(cls, objects: List, dim=0, *, out=None): |
| | """Stack a list of objects with the same type and shape.""" |
| | data = torch.stack([obj._data for obj in objects], dim=dim, out=out) |
| | return cls(data) |
| |
|
| | @classmethod |
| | def __torch_function__(cls, func, types, args=(), kwargs=None): |
| | """Support torch functions.""" |
| | if kwargs is None: |
| | kwargs = {} |
| | return cls.stack(*args, **kwargs) if func is torch.stack else NotImplemented |
| |
|
| |
|
| | def map_tensor(input_, func): |
| | if isinstance(input_, string_classes): |
| | return input_ |
| | elif isinstance(input_, collections.Mapping): |
| | return {k: map_tensor(sample, func) for k, sample in input_.items()} |
| | elif isinstance(input_, collections.Sequence): |
| | return [map_tensor(sample, func) for sample in input_] |
| | elif input_ is None: |
| | return None |
| | else: |
| | return func(input_) |
| |
|
| |
|
| | def batch_to_numpy(batch): |
| | return map_tensor(batch, lambda tensor: tensor.cpu().numpy()) |
| |
|
| |
|
| | def batch_to_device(batch, device, non_blocking=True, detach=False): |
| | def _func(tensor): |
| | t = tensor.to(device=device, non_blocking=non_blocking, dtype=torch.float32) |
| | return t.detach() if detach else t |
| |
|
| | return map_tensor(batch, _func) |
| |
|
| |
|
| | def remove_batch_dim(data: dict) -> dict: |
| | """Remove batch dimension from elements in data""" |
| | return { |
| | k: v[0] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v for k, v in data.items() |
| | } |
| |
|
| |
|
| | def add_batch_dim(data: dict) -> dict: |
| | """Add batch dimension to elements in data""" |
| | return { |
| | k: v[None] if isinstance(v, (torch.Tensor, np.ndarray, list)) else v |
| | for k, v in data.items() |
| | } |
| |
|
| |
|
| | def fit_to_multiple(x: torch.Tensor, multiple: int, mode: str = "center", crop: bool = False): |
| | """Get padding to make the image size a multiple of the given number. |
| | |
| | Args: |
| | x (torch.Tensor): Input tensor. |
| | multiple (int, optional): Multiple. |
| | crop (bool, optional): Whether to crop or pad. Defaults to False. |
| | |
| | Returns: |
| | torch.Tensor: Padding. |
| | """ |
| | h, w = x.shape[-2:] |
| |
|
| | if crop: |
| | pad_w = (w // multiple) * multiple - w |
| | pad_h = (h // multiple) * multiple - h |
| | else: |
| | pad_w = (multiple - w % multiple) % multiple |
| | pad_h = (multiple - h % multiple) % multiple |
| |
|
| | if mode == "center": |
| | pad_l = pad_w // 2 |
| | pad_r = pad_w - pad_l |
| | pad_t = pad_h // 2 |
| | pad_b = pad_h - pad_t |
| | elif mode == "left": |
| | pad_l = 0 |
| | pad_r = pad_w |
| | pad_t = 0 |
| | pad_b = pad_h |
| | else: |
| | raise ValueError(f"Unknown mode {mode}") |
| |
|
| | return (pad_l, pad_r, pad_t, pad_b) |
| |
|
| |
|
| | def fit_features_to_multiple( |
| | features: torch.Tensor, multiple: int = 32, crop: bool = False |
| | ) -> Tuple[torch.Tensor, Tuple[int, int]]: |
| | """Pad image to a multiple of the given number. |
| | |
| | Args: |
| | features (torch.Tensor): Input features. |
| | multiple (int, optional): Multiple. Defaults to 32. |
| | crop (bool, optional): Whether to crop or pad. Defaults to False. |
| | |
| | Returns: |
| | Tuple[torch.Tensor, Tuple[int, int]]: Padded features and padding. |
| | """ |
| | pad = fit_to_multiple(features, multiple, crop=crop) |
| | return torch.nn.functional.pad(features, pad, mode="reflect"), pad |
| |
|
