|
|
|
|
|
|
|
|
|
|
|
|
|
import torch |
|
import torch.nn.functional as F |
|
from typing import Optional, Tuple |
|
|
|
EPS = 1e-6 |
|
|
|
|
|
def smart_cat(tensor1, tensor2, dim): |
|
if tensor1 is None: |
|
return tensor2 |
|
return torch.cat([tensor1, tensor2], dim=dim) |
|
|
|
|
|
def get_points_on_a_grid( |
|
size: int, |
|
extent: Tuple[float, ...], |
|
center: Optional[Tuple[float, ...]] = None, |
|
device: Optional[torch.device] = torch.device("cpu"), |
|
shift_grid: bool = False, |
|
): |
|
r"""Get a grid of points covering a rectangular region |
|
|
|
`get_points_on_a_grid(size, extent)` generates a :attr:`size` by |
|
:attr:`size` grid fo points distributed to cover a rectangular area |
|
specified by `extent`. |
|
|
|
The `extent` is a pair of integer :math:`(H,W)` specifying the height |
|
and width of the rectangle. |
|
|
|
Optionally, the :attr:`center` can be specified as a pair :math:`(c_y,c_x)` |
|
specifying the vertical and horizontal center coordinates. The center |
|
defaults to the middle of the extent. |
|
|
|
Points are distributed uniformly within the rectangle leaving a margin |
|
:math:`m=W/64` from the border. |
|
|
|
It returns a :math:`(1, \text{size} \times \text{size}, 2)` tensor of |
|
points :math:`P_{ij}=(x_i, y_i)` where |
|
|
|
.. math:: |
|
P_{ij} = \left( |
|
c_x + m -\frac{W}{2} + \frac{W - 2m}{\text{size} - 1}\, j,~ |
|
c_y + m -\frac{H}{2} + \frac{H - 2m}{\text{size} - 1}\, i |
|
\right) |
|
|
|
Points are returned in row-major order. |
|
|
|
Args: |
|
size (int): grid size. |
|
extent (tuple): height and with of the grid extent. |
|
center (tuple, optional): grid center. |
|
device (str, optional): Defaults to `"cpu"`. |
|
|
|
Returns: |
|
Tensor: grid. |
|
""" |
|
if size == 1: |
|
return torch.tensor([extent[1] / 2, extent[0] / 2], device=device)[None, None] |
|
|
|
if center is None: |
|
center = [extent[0] / 2, extent[1] / 2] |
|
|
|
margin = extent[1] / 64 |
|
range_y = (margin - extent[0] / 2 + center[0], extent[0] / 2 + center[0] - margin) |
|
range_x = (margin - extent[1] / 2 + center[1], extent[1] / 2 + center[1] - margin) |
|
grid_y, grid_x = torch.meshgrid( |
|
torch.linspace(*range_y, size, device=device), |
|
torch.linspace(*range_x, size, device=device), |
|
indexing="ij", |
|
) |
|
|
|
if shift_grid: |
|
|
|
|
|
|
|
shift_x = (range_x[1] - range_x[0]) / (size - 1) |
|
shift_y = (range_y[1] - range_y[0]) / (size - 1) |
|
grid_x = grid_x + torch.randn_like(grid_x) / 3 * shift_x / 2 |
|
grid_y = grid_y + torch.randn_like(grid_y) / 3 * shift_y / 2 |
|
|
|
|
|
grid_x = torch.clamp(grid_x, range_x[0], range_x[1]) |
|
grid_y = torch.clamp(grid_y, range_y[0], range_y[1]) |
|
|
|
return torch.stack([grid_x, grid_y], dim=-1).reshape(1, -1, 2) |
|
|
|
|
|
def reduce_masked_mean(input, mask, dim=None, keepdim=False): |
|
r"""Masked mean |
|
|
|
`reduce_masked_mean(x, mask)` computes the mean of a tensor :attr:`input` |
|
over a mask :attr:`mask`, returning |
|
|
|
.. math:: |
|
\text{output} = |
|
\frac |
|
{\sum_{i=1}^N \text{input}_i \cdot \text{mask}_i} |
|
{\epsilon + \sum_{i=1}^N \text{mask}_i} |
|
|
|
where :math:`N` is the number of elements in :attr:`input` and |
|
:attr:`mask`, and :math:`\epsilon` is a small constant to avoid |
|
division by zero. |
|
|
|
`reduced_masked_mean(x, mask, dim)` computes the mean of a tensor |
|
:attr:`input` over a mask :attr:`mask` along a dimension :attr:`dim`. |
|
Optionally, the dimension can be kept in the output by setting |
|
:attr:`keepdim` to `True`. Tensor :attr:`mask` must be broadcastable to |
|
the same dimension as :attr:`input`. |
|
|
|
The interface is similar to `torch.mean()`. |
|
|
|
Args: |
|
inout (Tensor): input tensor. |
|
mask (Tensor): mask. |
|
dim (int, optional): Dimension to sum over. Defaults to None. |
|
keepdim (bool, optional): Keep the summed dimension. Defaults to False. |
|
|
|
Returns: |
|
Tensor: mean tensor. |
|
""" |
|
|
|
mask = mask.expand_as(input) |
|
|
|
prod = input * mask |
|
|
|
if dim is None: |
|
numer = torch.sum(prod) |
|
denom = torch.sum(mask) |
|
else: |
|
numer = torch.sum(prod, dim=dim, keepdim=keepdim) |
|
denom = torch.sum(mask, dim=dim, keepdim=keepdim) |
|
|
|
mean = numer / (EPS + denom) |
|
return mean |
|
|
|
|
|
def bilinear_sampler(input, coords, align_corners=True, padding_mode="border"): |
|
r"""Sample a tensor using bilinear interpolation |
|
|
|
`bilinear_sampler(input, coords)` samples a tensor :attr:`input` at |
|
coordinates :attr:`coords` using bilinear interpolation. It is the same |
|
as `torch.nn.functional.grid_sample()` but with a different coordinate |
|
convention. |
|
|
|
The input tensor is assumed to be of shape :math:`(B, C, H, W)`, where |
|
:math:`B` is the batch size, :math:`C` is the number of channels, |
|
:math:`H` is the height of the image, and :math:`W` is the width of the |
|
image. The tensor :attr:`coords` of shape :math:`(B, H_o, W_o, 2)` is |
|
interpreted as an array of 2D point coordinates :math:`(x_i,y_i)`. |
|
|
|
Alternatively, the input tensor can be of size :math:`(B, C, T, H, W)`, |
|
in which case sample points are triplets :math:`(t_i,x_i,y_i)`. Note |
|
that in this case the order of the components is slightly different |
|
from `grid_sample()`, which would expect :math:`(x_i,y_i,t_i)`. |
|
|
|
If `align_corners` is `True`, the coordinate :math:`x` is assumed to be |
|
in the range :math:`[0,W-1]`, with 0 corresponding to the center of the |
|
left-most image pixel :math:`W-1` to the center of the right-most |
|
pixel. |
|
|
|
If `align_corners` is `False`, the coordinate :math:`x` is assumed to |
|
be in the range :math:`[0,W]`, with 0 corresponding to the left edge of |
|
the left-most pixel :math:`W` to the right edge of the right-most |
|
pixel. |
|
|
|
Similar conventions apply to the :math:`y` for the range |
|
:math:`[0,H-1]` and :math:`[0,H]` and to :math:`t` for the range |
|
:math:`[0,T-1]` and :math:`[0,T]`. |
|
|
|
Args: |
|
input (Tensor): batch of input images. |
|
coords (Tensor): batch of coordinates. |
|
align_corners (bool, optional): Coordinate convention. Defaults to `True`. |
|
padding_mode (str, optional): Padding mode. Defaults to `"border"`. |
|
|
|
Returns: |
|
Tensor: sampled points. |
|
""" |
|
|
|
sizes = input.shape[2:] |
|
|
|
assert len(sizes) in [2, 3] |
|
|
|
if len(sizes) == 3: |
|
|
|
coords = coords[..., [1, 2, 0]] |
|
|
|
if align_corners: |
|
coords = coords * torch.tensor( |
|
[2 / max(size - 1, 1) for size in reversed(sizes)], device=coords.device |
|
) |
|
else: |
|
coords = coords * torch.tensor([2 / size for size in reversed(sizes)], device=coords.device) |
|
|
|
coords -= 1 |
|
|
|
return F.grid_sample(input, coords, align_corners=align_corners, padding_mode=padding_mode) |
|
|
|
|
|
def sample_features4d(input, coords): |
|
r"""Sample spatial features |
|
|
|
`sample_features4d(input, coords)` samples the spatial features |
|
:attr:`input` represented by a 4D tensor :math:`(B, C, H, W)`. |
|
|
|
The field is sampled at coordinates :attr:`coords` using bilinear |
|
interpolation. :attr:`coords` is assumed to be of shape :math:`(B, R, |
|
3)`, where each sample has the format :math:`(x_i, y_i)`. This uses the |
|
same convention as :func:`bilinear_sampler` with `align_corners=True`. |
|
|
|
The output tensor has one feature per point, and has shape :math:`(B, |
|
R, C)`. |
|
|
|
Args: |
|
input (Tensor): spatial features. |
|
coords (Tensor): points. |
|
|
|
Returns: |
|
Tensor: sampled features. |
|
""" |
|
|
|
B, _, _, _ = input.shape |
|
|
|
|
|
coords = coords.unsqueeze(2) |
|
|
|
|
|
feats = bilinear_sampler(input, coords) |
|
|
|
return feats.permute(0, 2, 1, 3).view( |
|
B, -1, feats.shape[1] * feats.shape[3] |
|
) |
|
|
|
|
|
def sample_features5d(input, coords): |
|
r"""Sample spatio-temporal features |
|
|
|
`sample_features5d(input, coords)` works in the same way as |
|
:func:`sample_features4d` but for spatio-temporal features and points: |
|
:attr:`input` is a 5D tensor :math:`(B, T, C, H, W)`, :attr:`coords` is |
|
a :math:`(B, R1, R2, 3)` tensor of spatio-temporal point :math:`(t_i, |
|
x_i, y_i)`. The output tensor has shape :math:`(B, R1, R2, C)`. |
|
|
|
Args: |
|
input (Tensor): spatio-temporal features. |
|
coords (Tensor): spatio-temporal points. |
|
|
|
Returns: |
|
Tensor: sampled features. |
|
""" |
|
|
|
B, T, _, _, _ = input.shape |
|
|
|
|
|
input = input.permute(0, 2, 1, 3, 4) |
|
|
|
|
|
coords = coords.unsqueeze(3) |
|
|
|
|
|
feats = bilinear_sampler(input, coords) |
|
|
|
return feats.permute(0, 2, 3, 1, 4).view( |
|
B, feats.shape[2], feats.shape[3], feats.shape[1] |
|
) |
|
|
