Add region_model.py for self-contained custom code

region_model.py

@@ -0,0 +1,43 @@
+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),
+        )