116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
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
|