in occant_utils/common.py [0:0]
def spatial_transform_map(p, x, invert=True, mode="bilinear"):
"""
Inputs:
p - (bs, f, H, W) Tensor
x - (bs, 3) Tensor (x, y, theta) transforms to perform
Outputs:
p_trans - (bs, f, H, W) Tensor
Conventions:
Shift in X is rightward, and shift in Y is downward. Rotation is clockwise.
Note: These denote transforms in an agent's position. Not the image directly.
For example, if an agent is moving upward, then the map will be moving downward.
To disable this behavior, set invert=False.
"""
device = p.device
H, W = p.shape[2:]
trans_x = x[:, 0]
trans_y = x[:, 1]
# Convert translations to -1.0 to 1.0 range
Hby2 = (H - 1) / 2 if H % 2 == 1 else H / 2
Wby2 = (W - 1) / 2 if W % 2 == 1 else W / 2
trans_x = trans_x / Wby2
trans_y = trans_y / Hby2
rot_t = x[:, 2]
sin_t = torch.sin(rot_t)
cos_t = torch.cos(rot_t)
# This R convention means Y axis is downwards.
A = torch.zeros(p.size(0), 3, 3).to(device)
A[:, 0, 0] = cos_t
A[:, 0, 1] = -sin_t
A[:, 1, 0] = sin_t
A[:, 1, 1] = cos_t
A[:, 0, 2] = trans_x
A[:, 1, 2] = trans_y
A[:, 2, 2] = 1
# Since this is a source to target mapping, and F.affine_grid expects
# target to source mapping, we have to invert this for normal behavior.
Ainv = torch.inverse(A)
# If target to source mapping is required, invert is enabled and we invert
# it again.
if invert:
Ainv = torch.inverse(Ainv)
Ainv = Ainv[:, :2]
grid = F.affine_grid(Ainv, p.size())
p_trans = F.grid_sample(p, grid, mode=mode)
return p_trans