Spaces:
Running
Running
File size: 1,514 Bytes
4d4dd90 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import torch
import torch.nn.functional as F
def local_correlation(
feature0,
feature1,
local_radius,
padding_mode="zeros",
flow = None
):
device = feature0.device
b, c, h, w = feature0.size()
if flow is None:
# If flow is None, assume feature0 and feature1 are aligned
coords = torch.meshgrid(
(
torch.linspace(-1 + 1 / h, 1 - 1 / h, h, device=device),
torch.linspace(-1 + 1 / w, 1 - 1 / w, w, device=device),
))
coords = torch.stack((coords[1], coords[0]), dim=-1)[
None
].expand(b, h, w, 2)
else:
coords = flow.permute(0,2,3,1) # If using flow, sample around flow target.
r = local_radius
local_window = torch.meshgrid(
(
torch.linspace(-2*local_radius/h, 2*local_radius/h, 2*r+1, device=device),
torch.linspace(-2*local_radius/w, 2*local_radius/w, 2*r+1, device=device),
))
local_window = torch.stack((local_window[1], local_window[0]), dim=-1)[
None
].expand(b, 2*r+1, 2*r+1, 2).reshape(b, (2*r+1)**2, 2)
coords = (coords[:,:,:,None]+local_window[:,None,None]).reshape(b,h,w*(2*r+1)**2,2)
window_feature = F.grid_sample(
feature1, coords, padding_mode=padding_mode, align_corners=False
)[...,None].reshape(b,c,h,w,(2*r+1)**2)
corr = torch.einsum("bchw, bchwk -> bkhw", feature0, window_feature)/(c**.5)
return corr
|