| """PoseConv3D PyTorch architecture for inference.""" |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class PoseConv3DEncoder(nn.Module): |
| def __init__(self, in_channels: int = 1, embed_dim: int = 64): |
| super().__init__() |
| self.c1 = nn.Conv3d(in_channels, 32, kernel_size=(3, 5, 5), stride=(1, 2, 2), padding=(1, 2, 2)) |
| self.c2 = nn.Conv3d(32, 64, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)) |
| self.c3 = nn.Conv3d(64, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)) |
| self.c4 = nn.Conv3d(128, 128, kernel_size=(3, 3, 3), stride=(2, 2, 2), padding=(1, 1, 1)) |
| self.fc = nn.Linear(256, embed_dim) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """(B, C, T, H, W) -> (B, embed_dim) normalized embedding.""" |
| h = F.relu(self.c1(x)) |
| h = F.relu(self.c2(h)) |
| h = F.relu(self.c3(h)) |
| h = F.relu(self.c4(h)) |
| h = torch.cat([h.mean(dim=(2, 3, 4)), h.amax(dim=(2, 3, 4))], dim=1) |
| z = self.fc(h) |
| return z / torch.sqrt((z * z).sum(dim=1, keepdim=True) + 1e-12) |
|
|
