Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from typing import Tuple | |
| class GeoPriorGen(nn.Module): | |
| def __init__(self,weight=[0.5,0.5]): | |
| super().__init__() | |
| self.weight = weight | |
| def generate_depth_decay(self, H: int, W: int, depth_grid): | |
| ''' | |
| generate 2d decay mask, the result is (HW)*(HW) | |
| H, W are the numbers of patches at each column and row | |
| ''' | |
| B,_,H,W = depth_grid.shape | |
| grid_d = depth_grid.reshape(B, H*W, 1) | |
| mask_d = grid_d[:, :, None, :] - grid_d[:, None, :, :] | |
| mask_d = (mask_d.abs()).sum(dim=-1) | |
| mask_d = mask_d.unsqueeze(1) * self.decay[None, :, None, None] | |
| return mask_d | |
| def generate_pos_decay(self, H: int, W: int): | |
| ''' | |
| generate 2d decay mask, the result is (HW)*(HW) | |
| H, W are the numbers of patches at each column and row | |
| ''' | |
| index_h = torch.arange(H).to(self.decay) | |
| index_w = torch.arange(W).to(self.decay) | |
| grid = torch.meshgrid([index_h, index_w]) | |
| grid = torch.stack(grid, dim=-1).reshape(H*W, 2) | |
| mask = grid[:, None, :] - grid[None, :, :] | |
| mask = (mask.abs()).sum(dim=-1) | |
| mask = mask * self.decay[:, None, None] | |
| return mask | |
| def generate_1d_depth_decay(self, H, W, depth_grid): | |
| ''' | |
| generate 1d depth decay mask, the result is l*l | |
| ''' | |
| mask = depth_grid[:, :, :, :, None] - depth_grid[:, :, :, None, :] | |
| mask = mask.abs() | |
| mask = mask * self.decay[:, None, None, None] | |
| assert mask.shape[2:] == (W,H,H) | |
| return mask | |
| def generate_1d_decay(self, l: int): | |
| ''' | |
| generate 1d decay mask, the result is l*l | |
| ''' | |
| index = torch.arange(l).to(self.decay) | |
| mask = index[:, None] - index[None, :] | |
| mask = mask.abs() | |
| mask = mask * self.decay[:, None, None] | |
| return mask | |
| def forward(self, HW_tuple: Tuple[int], depth_map, split_or_not=False): | |
| ''' | |
| depth_map: depth patches | |
| HW_tuple: (H, W) | |
| H * W == l | |
| ''' | |
| depth_map = F.interpolate(depth_map, size=HW_tuple,mode='bilinear',align_corners=False) | |
| if split_or_not: | |
| index = torch.arange(HW_tuple[0]*HW_tuple[1]).to(self.decay) | |
| sin = torch.sin(index[:, None] * self.angle[None, :]) | |
| sin = sin.reshape(HW_tuple[0], HW_tuple[1], -1) | |
| cos = torch.cos(index[:, None] * self.angle[None, :]) | |
| cos = cos.reshape(HW_tuple[0], HW_tuple[1], -1) | |
| mask_d_h = self.generate_1d_depth_decay(HW_tuple[0], HW_tuple[1], depth_map.transpose(-2,-1)) | |
| mask_d_w = self.generate_1d_depth_decay(HW_tuple[1], HW_tuple[0], depth_map) | |
| mask_h = self.generate_1d_decay(HW_tuple[0]) | |
| mask_w = self.generate_1d_decay(HW_tuple[1]) | |
| mask_h = self.weight[0]*mask_h.unsqueeze(0).unsqueeze(2) + self.weight[1]*mask_d_h | |
| mask_w = self.weight[0]*mask_w.unsqueeze(0).unsqueeze(2) + self.weight[1]*mask_d_w | |
| geo_prior = ((sin, cos), (mask_h, mask_w)) | |
| else: | |
| index = torch.arange(HW_tuple[0]*HW_tuple[1]).to(self.decay) | |
| sin = torch.sin(index[:, None] * self.angle[None, :]) | |
| sin = sin.reshape(HW_tuple[0], HW_tuple[1], -1) | |
| cos = torch.cos(index[:, None] * self.angle[None, :]) | |
| cos = cos.reshape(HW_tuple[0], HW_tuple[1], -1) | |
| mask = self.generate_pos_decay(HW_tuple[0], HW_tuple[1]) | |
| mask_d = self.generate_depth_decay(HW_tuple[0], HW_tuple[1], depth_map) | |
| mask = (self.weight[0]*mask+self.weight[1]*mask_d) | |
| geo_prior = ((sin, cos), mask) | |
| return geo_prior |