Merge branch 'main' of https://gitea.lnrweb.com/luni/reconstruction into main

models/pose_manipulator.py

@@ -0,0 +1,79 @@
+import os
+import sys
+
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+torch.set_float32_matmul_precision('high')
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from models.convnext import CNBlock, LayerNorm2d
+from models.pose_ae import PoseAE
+from utils.checkpoints import load_model
+
+
+class FeatureManipulator(nn.Module):
+    def __init__(
+        self,
+        features: int = 256,
+        device: str = 'cuda',
+    ) -> None:
+        super().__init__()
+
+        norm_layer = partial(LayerNorm2d, eps=1e-6)
+        self.last = nn.Sequential(
+                        norm_layer(features + 1),
+                        nn.Conv2d(features + 1, features, kernel_size=3, padding='same'),
+                    ).to(device)
+        
+        stage: List[nn.Module] = []
+        for _ in range(10):
+            stage.append(CNBlock(features + 1, 1.0, 0))
+        self.stage = nn.Sequential(*stage).to(device)
+
+    def _forward_impl(self, x) -> Tensor:
+        x = self.stage(x)
+        x = self.last(x)
+        return x
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+class PoseManipulator(nn.Module):
+    def __init__(self, features=256, device='cuda', ae_checpoint=None):
+        super().__init__()
+
+        self.pose_ae: PoseAE = PoseAE(features=features, device=device)
+        iteration = load_model(self.pose_ae.encoder, ae_checpoint + '_enc')
+        load_model(self.pose_ae.decoder, ae_checpoint + '_dec')
+        for param in self.pose_ae.encoder.parameters():
+            param.requires_grad = False
+        for param in self.pose_ae.decoder.parameters():
+            param.requires_grad = False
+
+        print(f'Loaded pose_ae checkpoint at iteration: {iteration}')
+
+        self.manipulator = FeatureManipulator(features=features, device=device)
+
+
+    def forward(self, input) -> torch.Tensor:
+        mask, pose_in, mask_out, pose_out, angle = input
+        original_features, _ = self.pose_ae.encoder(torch.cat([mask, pose_in], dim=1))
+        expected_features, _ = self.pose_ae.encoder(torch.cat([mask_out, pose_out], dim=1))
+
+        tiled_angle = torch.ones_like(original_features[:, :1, :, :]) * angle[:,:,None,None]
+        modified_features = self.manipulator(torch.cat((original_features, tiled_angle), dim=1))
+        pose_pred, mask_pred = self.pose_ae.decoder(modified_features)
+
+        return pose_pred, mask_pred, expected_features, modified_features
+
+
+if __name__ == '__main__':
+    pass