Spaces:
Running
Running
File size: 3,805 Bytes
7f0f123 |
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 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
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 |