Multi-Robot-application / code /baselines /lhd /src /utils /diffusion_utils.py
tarzanagh's picture
Add files using upload-large-folder tool
c35fc6b verified
Raw
History Blame Contribute Delete
830 Bytes
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