| import torch | |
| import torch.nn as nn | |
| from .fourier_features import FourierFeatures | |
| class RegionModel(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.position_features = FourierFeatures(2, 256) | |
| self.position_encoder = nn.Linear(256, 2048) | |
| self.size_features = FourierFeatures(2, 256) | |
| self.size_encoder = nn.Linear(256, 2048) | |
| self.position_decoder = nn.Linear(2048, 2) | |
| self.size_decoder = nn.Linear(2048, 2) | |
| self.confidence_decoder = nn.Linear(2048, 1) | |
| def encode_position(self, position): | |
| return self.position_encoder(self.position_features(position)) | |
| def encode_size(self, size): | |
| return self.size_encoder(self.size_features(size)) | |
| def decode_position(self, x): | |
| return self.position_decoder(x) | |
| def decode_size(self, x): | |
| return self.size_decoder(x) | |
| def decode_confidence(self, x): | |
| return self.confidence_decoder(x) | |
| def encode(self, position, size): | |
| return torch.stack( | |
| [self.encode_position(position), self.encode_size(size)], dim=0 | |
| ) | |
| def decode(self, position_logits, size_logits): | |
| return ( | |
| self.decode_position(position_logits), | |
| self.decode_size(size_logits), | |
| self.decode_confidence(size_logits), | |
| ) | |