Spaces:
Running
on
Zero
Running
on
Zero
| import torch | |
| import torch.nn.functional as F | |
| import torchvision | |
| class BilinearDownsampler(torch.nn.Module): | |
| def __init__( | |
| self, | |
| patch_size, | |
| ): | |
| super().__init__() | |
| if isinstance(patch_size, int): | |
| self.patch_size = (patch_size, patch_size) | |
| elif isinstance(patch_size, tuple): | |
| self.patch_size = patch_size | |
| def forward(self, x, mode): | |
| n, v, h, w, _, c = x.shape | |
| assert h % self.patch_size[0] == 0 | |
| target_h = h // self.patch_size[0] | |
| assert w % self.patch_size[1] == 0 | |
| target_w = w // self.patch_size[1] | |
| x = x.permute(0, 1, 4, 5, 2, 3).flatten(0, 2) | |
| x = F.interpolate(x, size=(target_h, target_w), mode="bilinear") | |
| x = x.reshape(n, v, -1, c, target_h, target_w).permute(0, 1, 4, 5, 2, 3) | |
| return x.squeeze(2, 3) | |
| class PatchSalienceDownsampler(torch.nn.Module): | |
| def __init__( | |
| self, | |
| channels, | |
| patch_size, | |
| normalize_features, | |
| ): | |
| super().__init__() | |
| if isinstance(patch_size, int): | |
| self.patch_size = (patch_size, patch_size) | |
| elif isinstance(patch_size, tuple): | |
| self.patch_size = patch_size | |
| self.conv = torch.nn.Conv2d(channels, 1, kernel_size=1) | |
| self.patch_weight = torch.nn.Parameter(torch.ones(self.patch_size)) | |
| self.patch_bias = torch.nn.Parameter(torch.zeros(self.patch_size)) | |
| self.normalize_features = normalize_features | |
| torch.nn.init.kaiming_normal_(self.conv.weight, a=0, mode="fan_in") | |
| torch.nn.init.zeros_(self.conv.bias) | |
| torch.nn.init.normal_(self.patch_weight, mean=1.0, std=0.01) | |
| torch.nn.init.normal_(self.patch_bias, mean=0.0, std=0.01) | |
| def forward(self, x, mode): | |
| if mode == "patch": | |
| return self.forward_patches(x) | |
| elif mode == "image": | |
| n, v, h, w, _, c = x.shape | |
| patch_h, patch_w = self.patch_size[0], self.patch_size[1] | |
| no_patches_h, no_patches_w = h // patch_h, w // patch_w | |
| patches = x.reshape(n, v, no_patches_h, patch_h, no_patches_w, patch_w, 1, c) | |
| patches = patches.swapaxes(3, 4).flatten(1, 3) | |
| patched_result, salience_map, weight_map, patch_weight_bias = self.forward_patches(patches) | |
| patched_result = patched_result.reshape(n, v, no_patches_h, no_patches_w, 1, c) | |
| salience_map = salience_map.reshape(n, v, no_patches_h, no_patches_w, patch_h, patch_w, 1, 1) | |
| salience_map = salience_map.swapaxes(3, 4).reshape(n, v, h, w, 1, 1) | |
| weight_map = weight_map.reshape(n, v, no_patches_h, no_patches_w, patch_h, patch_w, 1, 1) | |
| weight_map = weight_map.swapaxes(3, 4).reshape(n, v, h, w, 1, 1) | |
| return patched_result, salience_map, weight_map, patch_weight_bias | |
| else: | |
| return None | |
| def forward_patches(self, x): | |
| n, p, patch_h, patch_w, _, c = x.shape | |
| x_flat = x.reshape(-1, patch_h, patch_w, c).permute(0, 3, 1, 2) | |
| salience_map = self.conv(x_flat).squeeze(1) | |
| weight_map = salience_map * self.patch_weight + self.patch_bias | |
| weight_map = torch.nn.functional.softmax(weight_map.reshape(-1, patch_h * patch_w), dim=1) | |
| weight_map = weight_map.reshape(n, p, patch_h, patch_w, 1, 1) | |
| patched_features = torch.sum(weight_map * x, dim=(2, 3)) | |
| if self.normalize_features: | |
| patched_features = patched_features / torch.linalg.norm(patched_features, dim=-1, keepdim=True) | |
| return (patched_features, | |
| salience_map.reshape(n, p, patch_h, patch_w, 1, 1), | |
| weight_map, | |
| torch.cat([self.patch_weight, self.patch_bias], dim=1)) | |