Spaces:
Running
Running
| 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**0.5) | |
| return corr | |