Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| from .fourier_features import FourierFeatures | |
| class MLP(nn.Module): | |
| def __init__( | |
| self, | |
| in_features: int, | |
| hidden_features: int = None, | |
| out_features: int = None, | |
| ) -> None: | |
| super().__init__() | |
| out_features = out_features or in_features | |
| hidden_features = hidden_features or in_features * 4 | |
| self.fc1 = nn.Linear(in_features, hidden_features) | |
| self.act = nn.GELU(approximate="tanh") | |
| self.fc2 = nn.Linear(hidden_features, out_features) | |
| torch.nn.init.kaiming_normal_( | |
| self.fc1.weight, mode="fan_in", nonlinearity="relu" | |
| ) | |
| torch.nn.init.kaiming_normal_( | |
| self.fc2.weight, mode="fan_in", nonlinearity="relu" | |
| ) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.fc1(x) | |
| x = self.act(x) | |
| x = self.fc2(x) | |
| return x | |
| class RegionModel(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.coordinate_features = FourierFeatures(1, 256) | |
| self.coordinate_encoder = nn.Linear(256, 2048) | |
| self.size_features = FourierFeatures(2, 512) | |
| self.size_encoder = nn.Linear(512, 2048) | |
| self.coordinate_decoder = MLP(2048, 8192, 1024) | |
| self.size_decoder = MLP(2048, 8192, 2048) | |
| def encode_coordinate(self, coordinate): | |
| return self.coordinate_encoder(self.coordinate_features(coordinate)) | |
| def encode_size(self, size): | |
| return self.size_encoder(self.size_features(size)) | |
| def decode_coordinate(self, logit): | |
| return self.coordinate_decoder(logit) | |
| def decode_size(self, logit): | |
| o = self.size_decoder(logit) | |
| return o.view(-1, 2, 1024) | |
| def encode(self, position, size): | |
| c = self.encode_coordinate(position.view(2, 1)).view(2, 2048) | |
| return torch.stack([c[0], c[1], self.encode_size(size)], dim=0) | |
| def decode(self, position_logits, size_logits): | |
| return ( | |
| self.decode_coordinate(position_logits), | |
| self.decode_size(size_logits), | |
| ) | |