File size: 1,482 Bytes
70d1188 |
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 |
import torch
import torch.nn as nn
import torch.nn.functional as F
from .postprocess import postprocess
class LinearPts3d (nn.Module):
"""
Linear head for dust3r
Each token outputs: - 16x16 3D points (+ confidence)
"""
def __init__(self, net, has_conf=False,mode='pts3d'):
super().__init__()
self.patch_size = net.patch_size
self.depth_mode = net.depth_mode
self.conf_mode = net.conf_mode
self.has_conf = has_conf
self.mode = mode
self.classifier_mode = None
if self.mode == 'pts3d':
self.proj = nn.Linear(net.dec_embed_dim, (3 + has_conf)*self.patch_size**2)
elif self.mode == 'depth':
self.proj = nn.Linear(net.dec_embed_dim, (1 + has_conf)*self.patch_size**2)
elif self.mode == 'classifier':
self.proj = nn.Linear(net.dec_embed_dim, (1 + has_conf)*self.patch_size**2)
self.classifier_mode = net.classifier_mode
def setup(self, croconet):
pass
def forward(self, decout, img_shape):
H, W = img_shape
tokens = decout[-1]
B, S, D = tokens.shape
# extract 3D points
feat = self.proj(tokens) # B,S,D
feat = feat.transpose(-1, -2).view(B, -1, H//self.patch_size, W//self.patch_size)
feat = F.pixel_shuffle(feat, self.patch_size) # B,3,H,W
# permute + norm depth
return postprocess(feat, self.depth_mode, self.conf_mode,self.classifier_mode)
|