File size: 2,721 Bytes
9882c88 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from anycalib.manifolds import Unit3
def cvx_upsample(x: Tensor, mask: Tensor, up_factor: int = 7) -> Tensor:
"""Upsample [H/k, W/k, C] -> [H, W, C] using convex combination of 3x3 patches.
Code adapted from RAFT (Teed and Deng, 2020):
https://github.com/princeton-vl/RAFT/blob/3fa0bb0a9c633ea0a9bb8a79c576b6785d4e6a02/core/raft.py#L72
Args:
x: (N, C, H, W) input tensor
mask: (N, 9, 1, 1, H, W) already softmaxed mask tensor
up_factor: upsample factor
"""
N, C, H, W = x.shape
up_x = F.unfold(x, (3, 3), padding=1)
up_x = up_x.view(N, C, 9, 1, 1, H, W)
up_x = torch.sum(mask * up_x, dim=2)
up_x = up_x.permute(0, 1, 4, 2, 5, 3)
return up_x.reshape(N, C, up_factor * H, up_factor * W)
class ConvexTangentDecoder(nn.Module):
"""Convex Tangent Coordinates Decoder.
This decoder predicts 2D coordinates in the tangent space of the unit sphere at the
optical axis: z_1 = [0, 0, 1]. These coordinates are subsequently mapped to unit
rays using the exponential map.
Args:
in_channels: number of input channels
up_factor: upsampling factor
"""
def __init__(self, in_channels: int = 256, up_factor: int = 7):
super().__init__()
self.in_channels = in_channels
self.up_factor = up_factor
# tangent head
self.tangent_head = nn.Sequential(
nn.Conv2d(in_channels, in_channels // 2, 3, padding=1),
nn.ReLU(True),
# tangent coords (2)
nn.Conv2d(in_channels // 2, 2, 1),
)
# weights head for convex upsampling to input resolution
self.upsampling_weights_head = nn.Sequential(
# convex combination of 3x3 patches
nn.Conv2d(in_channels, in_channels // 2, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(in_channels // 2, up_factor**2 * 9, 1, padding=0),
nn.Unflatten(1, (1, 9, up_factor, up_factor)),
nn.Softmax(dim=2),
)
def forward(self, x: Tensor) -> dict[str, Tensor | float]:
# head
tangent_pred = self.tangent_head(x) # (B, 5, H/7, W/7)
weights = self.upsampling_weights_head(x)
# upsample
tangent_pred = cvx_upsample(tangent_pred, weights, self.up_factor)
# postprocess
tangent_coords = tangent_pred[:, :2]
rays = Unit3.expmap_at_z1(tangent_coords.permute(0, 2, 3, 1)).permute(
0, 3, 1, 2
)
out: dict[str, Tensor | float] = {
"rays": rays,
"tangent_coords": tangent_coords,
}
return out
|