Fsrt node (#9)

module_fsrt/encoder.py

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+
+from .layers import FSRTPosEncoder, Transformer
+
+
+class SRTConvBlock(nn.Module):
+    def __init__(self, idim, hdim=None, odim=None):
+        super().__init__()
+        if hdim is None:
+            hdim = idim
+
+        if odim is None:
+            odim = 2 * hdim
+
+        conv_kwargs = {"bias": False, "kernel_size": 3, "padding": 1}
+        self.layers = nn.Sequential(
+            nn.Conv2d(idim, hdim, stride=1, **conv_kwargs),
+            nn.ReLU(),
+            nn.Conv2d(hdim, odim, stride=2, **conv_kwargs),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class ImprovedFSRTEncoder(nn.Module):
+    """
+    Scene Representation Transformer Encoder with the improvements from Appendix A.4 in the OSRT paper.
+    """
+
+    def __init__(
+        self,
+        num_conv_blocks=3,
+        num_att_blocks=5,
+        pix_octaves=16,
+        pix_start_octave=-1,
+        num_kp=10,
+        expression_size=256,
+        encode_with_expression=True,
+        kp_octaves=4,
+        kp_start_octave=-1,
+    ):
+        super().__init__()
+        self.positional_encoder = FSRTPosEncoder(
+            kp_octaves=kp_octaves,
+            kp_start_octave=kp_start_octave,
+            pix_octaves=pix_octaves,
+            pix_start_octave=pix_start_octave,
+        )
+
+        self.encode_with_expression = encode_with_expression
+        if self.encode_with_expression:
+            self.expression_size = expression_size
+        else:
+            self.expression_size = 0
+        conv_blocks = [
+            SRTConvBlock(
+                idim=3
+                + pix_octaves * 4
+                + num_kp * kp_octaves * 4
+                + self.expression_size,
+                hdim=96,
+            )
+        ]
+        cur_hdim = 192
+        for i in range(1, num_conv_blocks):
+            conv_blocks.append(SRTConvBlock(idim=cur_hdim, odim=None))
+            cur_hdim *= 2
+
+        self.conv_blocks = nn.Sequential(*conv_blocks)
+
+        self.per_patch_linear = nn.Conv2d(cur_hdim, 768, kernel_size=1)
+
+        self.transformer = Transformer(
+            768, depth=num_att_blocks, heads=12, dim_head=64, mlp_dim=1536, selfatt=True
+        )
+        self.num_kp = num_kp
+
+    def forward(self, images, keypoints, pixels, expression_vector=None):
+        """
+        Args:
+            images: [batch_size, num_images, 3, height, width]
+            keypoints: [batch_size, num_images, num_kp, 2]
+            pixels: [batch_size, num_images, height, width, 2]
+            expression_vector: [batch_size, num_images, expression_size]
+        Returns:
+            scene representation: [batch_size, num_patches, channels_per_patch]
+        """
+
+        batch_size, num_images = images.shape[:2]
+
+        x = images.flatten(0, 1)
+        keypoints = keypoints.flatten(-2, -1).flatten(0, 1)
+        pixels = pixels.flatten(0, 1)
+
+        pos_enc = self.positional_encoder(pixels, keypoints)
+        if expression_vector is not None and self.encode_with_expression:
+            expression_vector = expression_vector.flatten(0, 1)[
+                :, :, None, None
+            ].repeat(1, 1, images.shape[-2], images.shape[-1])
+            x = torch.cat([x, pos_enc, expression_vector], 1)
+        else:
+            x = torch.cat([x, pos_enc], 1)
+        x = self.conv_blocks(x)
+        x = self.per_patch_linear(x)
+        x = x.flatten(2, 3).permute(0, 2, 1)
+
+        patches_per_image, channels_per_patch = x.shape[1:]
+        x = x.reshape(batch_size, num_images * patches_per_image, channels_per_patch)
+
+        x = self.transformer(x)
+
+        return x