Spaces:
Running on Zero
Running on Zero
File size: 2,496 Bytes
41ff959 | 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 | from collections.abc import Sequence
import torch
from torch import nn
import torch.nn.functional as F
class ImplicitMLP(nn.Module):
def __init__(
self,
in_dim: int,
out_dim: int,
hidden_list: Sequence[int],
) -> None:
super().__init__()
layers = []
last_dim = in_dim
for hidden_dim in hidden_list:
layers.extend((nn.Linear(last_dim, hidden_dim), nn.ReLU()))
last_dim = hidden_dim
layers.append(nn.Linear(last_dim, out_dim))
self.layers = nn.Sequential(*layers)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
return self.layers(inputs)
class ImplicitGSHead(nn.Module):
"""Decode Gaussian parameters from DINOv3 and BasicEncoder features."""
def __init__(
self,
hidden_dim: int,
basic_dim: int,
hidden_list: Sequence[int],
) -> None:
super().__init__()
self.gate_proj = nn.Linear(basic_dim, hidden_dim)
self.gate = nn.Sequential(
nn.Linear(hidden_dim * 2, hidden_dim),
nn.Sigmoid(),
)
self.out_layer = ImplicitMLP(
in_dim=hidden_dim,
out_dim=14,
hidden_list=hidden_list,
)
def encode_feat(self, features, patch_h: int, patch_w: int) -> torch.Tensor:
"""Reshape the last DINOv3 layer into a spatial feature map."""
tokens = features[-1][0]
return tokens.permute(0, 2, 1).reshape(
tokens.shape[0],
tokens.shape[-1],
patch_h,
patch_w,
)
def decode_dpt(
self,
feat: torch.Tensor,
basic_feat: torch.Tensor,
coord: torch.Tensor,
) -> torch.Tensor:
coord = coord.clone()
coord.clamp_(-1 + 1e-6, 1 - 1e-6)
query_dino = F.grid_sample(
feat,
coord.flip(-1).unsqueeze(1),
mode="bilinear",
align_corners=False,
)[:, :, 0, :].permute(0, 2, 1)
query_basic = F.grid_sample(
basic_feat,
coord.flip(-1).unsqueeze(1),
mode="bilinear",
align_corners=False,
)[:, :, 0, :].permute(0, 2, 1)
query_basic = self.gate_proj(query_basic)
gate_weights = self.gate(torch.cat((query_dino, query_basic), dim=-1))
fused_features = gate_weights * query_dino + (1 - gate_weights) * query_basic
return self.out_layer(fused_features)
|