File size: 830 Bytes
c35fc6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | import torch.nn.functional as F
def bilinear_interpolate(fmap, pos):
assert fmap.size(0) == pos.size(0), "Batch size of fmap and pos must be the same"
grid = pos.unsqueeze(-2).unsqueeze(-2) # Shape will be (N, 1, 1, 2)
grid = grid.to(fmap.device)
interpolated_value = F.grid_sample(fmap, grid, mode="bilinear", padding_mode="zeros", align_corners=True)
return interpolated_value.squeeze(-1).squeeze(-1)
def bilinear_interpolate_samples(fmap, pos):
assert fmap.size(0) == pos.size(0), 'Batch size of fmap and pos must be the same'
grid = pos.unsqueeze(-2)
grid = grid.to(fmap.device)
interpolated_value = F.grid_sample(fmap, grid, mode="bilinear", padding_mode="zeros", align_corners=True)
scores = interpolated_value.squeeze(-1).permute(0,2,1)
return scores
|