pose encoder model

data/pose_renderer.py

@@ -32,14 +32,14 @@ class PoseRenderer(Module):
         super().__init__()
         colors = np.load('custom_data/pose_colors.npy')
         self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
-        self.cam = torch.tensor([5.0, 0.0, 0.0], device=device).cpu().numpy()
+        self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
 
 
     def forward(self, verts: torch.Tensor, faces: torch.Tensor):
-        cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
+        cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512)
 
         raster_settings = RasterizationSettings(
-                            image_size=(480, 640), 
+                            image_size=(512, 512), 
                             blur_radius=0.0, 
                             faces_per_pixel=1, 
                             bin_size=0

docker/Dockerfile

@@ -1,6 +1,7 @@
 # Use an existing image as the base
-#FROM nvcr.io/nvidia/pytorch:22.06-py3
+# FROM nvcr.io/nvidia/pytorch:23.04-py3
-FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-devel
+FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
+# FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-devel
 
 # Add the current user to the image
 ARG USER_NAME=lhunor

docker/docker-compose.yml

@@ -38,3 +38,23 @@ services:
         reservations:
           devices:
             - capabilities: [gpu]
+
+  lhunor_torch_3:
+    image: lhunor_reconstruction:v2.0
+    container_name: lhunor_reconstruction_v2.0
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      #- '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

docker/requirements.txt

@@ -5,3 +5,4 @@ scipy
 pandas
 pytorch_msssim
 kornia
+wandb

models/convnext.py

@@ -94,6 +94,7 @@ class ConvNeXtBase(nn.Module):
         layer_scale: float = 1e-6,
         block: Optional[Callable[..., nn.Module]] = None,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
+        input_features: int = 3,
         device: str = 'cuda',
     ) -> None:
         super().__init__()
@@ -115,7 +116,7 @@ class ConvNeXtBase(nn.Module):
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
-                3,
+                input_features,
                 firstconv_output_channels,
                 kernel_size=4,
                 stride=4,
@@ -125,6 +126,7 @@ class ConvNeXtBase(nn.Module):
                 bias=True,
             ).to(device)
         
+        layers_list: List[nn.Module] = []
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
@@ -136,16 +138,17 @@ class ConvNeXtBase(nn.Module):
                 sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                 stage.append(block(cnf.input_channels, layer_scale, sd_prob))
                 stage_block_id += 1
-            self.stages.append(nn.Sequential(*stage).to(device))
+            layers_list.append(nn.Sequential(*stage))
             if cnf.out_channels is not None:
                 # Downsampling
-                self.downsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
                     ).to(device)
                 )
 
+        self.layers = nn.Sequential(*layers_list).to(device)
         self.act = nn.Tanh()
 
         for m in self.modules():
@@ -157,11 +160,7 @@ class ConvNeXtBase(nn.Module):
     def _forward_impl(self, x: Tensor) -> Tensor:
         x = self.stem(x)
         features: List[Tensor] = []
-        for i in range(4 - 1):
+        x = self.layers(x)
-            x = self.stages[i](x)
-            features.append(x)
-            x = self.downsamples[i](x)
-        x = self.stages[-1](x)
         x = self.act(x)
 
         return x, features
@@ -171,7 +170,7 @@ class ConvNeXtBase(nn.Module):
 
 
 
-class ConvNeXtReverseBase2(nn.Module):
+class ConvNeXtReverseDepthNormals(nn.Module):
     def __init__(
         self,
         block_setting: List[CNBlockConfig],
@@ -217,8 +216,8 @@ class ConvNeXtReverseBase2(nn.Module):
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        self.stages: List[nn.Module] = []
+        layers_list: List[nn.Module] = []
-        self.upsamples: List[nn.Module] = []
+
         for cnf in block_setting[1:]:
             if cnf.input_channels is not None:
                 # Bottlenecks
@@ -228,19 +227,21 @@ class ConvNeXtReverseBase2(nn.Module):
                     sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                     stage.append(block(cnf.input_channels, layer_scale, sd_prob))
                     stage_block_id += 1
-                self.stages.append(nn.Sequential(*stage).to(device))
+                layers_list.append(nn.Sequential(*stage).to(device))
             else:
-                self.stages.append(nn.Identity().to(device))
+                layers_list.append(nn.Identity().to(device))
 
             if cnf.out_channels is not None:
-                # Upnsampling
+                # Upsampling
-                self.upsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
                     ).to(device)
                 )
 
+        self.layers = nn.Sequential(*layers_list).to(device)
+
         for cnf in block_setting[-1:]:
             # Bottlenecks
             stage: List[nn.Module] = []
@@ -275,7 +276,10 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        input, features = x
+        if self.use_skip:
+            input, features = x
+        else:
+            input = x
         x = self.first(input)
         for i in range(5 - 2):
             if self.use_skip:
@@ -306,10 +310,14 @@ class ConvNeXtReverseBase(nn.Module):
         stochastic_depth_prob: float = 0.0,
         layer_scale: float = 1e-6,
         block: Optional[Callable[..., nn.Module]] = None,
-        norm_layer: Optional[Callable[..., nn.Module]] = None
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip=False,
     ) -> None:
         super().__init__()
 
+        self.use_skip = use_skip
+
         if not block_setting:
             raise ValueError("The block_setting should not be empty")
         elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
@@ -321,71 +329,44 @@ class ConvNeXtReverseBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        layers: List[nn.Module] = []
+        self.first = nn.Sequential(
-
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
-
+                        norm_layer(block_setting[0].out_channels),
-        # add bottleneck which will be executed at the end, to process after the final larger upsampling
+                    ).to(device)
-
-
-
-        # Stem
-        firstconv_output_channels = block_setting[0].input_channels
-
-        layers.append(
-            nn.Sequential(
-                nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
-
-            )
-        )
-
-        stage: List[nn.Module] = []
-        for _ in range(3):
-            stage.append(block(16, layer_scale, 0.0))
-        layers.append(nn.Sequential(*stage))
-        
-        layers.append(
-            # Convtr NormActivation(
-            #     firstconv_output_channels,
-            #     3,
-            #     kernel_size=4,
-            #     stride=4,
-            #     padding=0,
-            #     norm_layer=None,
-            #     activation_layer=None,
-            #     bias=True,
-            # )
-            nn.Sequential(
-                nn.ConvTranspose2d(
-                    firstconv_output_channels,
-                    16,
-                    kernel_size=4,
-                    stride=4,
-                    padding=0,
-                    ),
-            )
-        )
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        for cnf in block_setting:
+        layers_list: List[nn.Module] = []
-            # Bottlenecks
+
-            stage: List[nn.Module] = []
+        for cnf in block_setting[1:]:
-            for _ in range(cnf.num_layers):
+            if cnf.input_channels is not None:
-                # adjust stochastic depth probability based on the depth of the stage block
+                # Bottlenecks
-                sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                stage: List[nn.Module] = []
-                stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                for _ in range(cnf.num_layers):
-                stage_block_id += 1
+                    # adjust stochastic depth probability based on the depth of the stage block
-            layers.append(nn.Sequential(*stage))
+                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                    stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                    stage_block_id += 1
+                layers_list.append(nn.Sequential(*stage).to(device))
+            else:
+                layers_list.append(nn.Identity().to(device))
+
             if cnf.out_channels is not None:
-                # Upsampling
+                # Upnsampling
-                layers.append(
+                layers_list.append(
                     nn.Sequential(
-                        norm_layer(cnf.out_channels),
+                        norm_layer(cnf.input_channels),
-                        nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
-                    )
+                    ).to(device)
                 )
 
-        self.features = nn.Sequential(*reversed(layers))
+        self.layers = nn.Sequential(*layers_list).to(device)
+        self.last = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    1,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
 
         for m in self.modules():
             if isinstance(m, (nn.Conv2d, nn.Linear)):
@@ -393,12 +374,20 @@ class ConvNeXtReverseBase(nn.Module):
                 if m.bias is not None:
                     nn.init.zeros_(m.bias)
 
-    def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.features(x)
-        # x = torch.tanh(x)
-        return x
 
-    def forward(self, x: Tensor) -> Tensor:
+    def _forward_impl(self, x) -> Tensor:
+        if self.use_skip:
+            raise NotImplementedError
+            input, features = x
+        else:
+            input = x
+        x = self.first(input)
+        x = self.layers(x)
+        out = self.last(x)
+        
+        return out
+
+    def forward(self, x) -> Tensor:
         return self._forward_impl(x)
 
 
@@ -407,13 +396,14 @@ def _convnext(
     stochastic_depth_prob: float,
     weights: Optional[WeightsEnum],
     reversed: bool,
+    input_features: int,
     **kwargs: Any,
 ) -> ConvNeXtBase:
 
     if reversed:
-        model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob,**kwargs)
     else:
-        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, input_features=input_features, **kwargs)
 
     if weights is not None:
         model.load_state_dict(weights.get_state_dict())
@@ -422,10 +412,10 @@ def _convnext(
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
+    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, input_features=3, **kwargs) -> None:
         super().__init__()
 
-        self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, **kwargs)
+        self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, input_features, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.model(x)

models/pose_encoder.py

@@ -0,0 +1,205 @@
+import os
+import sys
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '7'
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+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 ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
+
+
+class ConvNeXtReverseMaskPose(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip=False,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        self.first = nn.Sequential(
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
+                        norm_layer(block_setting[0].out_channels),
+                    ).to(device)
+
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        layers_list: List[nn.Module] = []
+
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # Bottlenecks
+                stage: List[nn.Module] = []
+                for _ in range(cnf.num_layers):
+                    # adjust stochastic depth probability based on the depth of the stage block
+                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                    stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                    stage_block_id += 1
+                layers_list.append(nn.Sequential(*stage).to(device))
+            else:
+                layers_list.append(nn.Identity().to(device))
+
+            if cnf.out_channels is not None:
+                # Upsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    ).to(device)
+                )
+
+        self.layers = nn.Sequential(*layers_list[:-1]).to(device)
+        self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
+
+        for cnf in block_setting[-1:]:
+            # Bottlenecks
+            stage: List[nn.Module] = []
+            for _ in range(cnf.num_layers):
+                stage.append(block(cnf.input_channels, layer_scale, 0))
+            self.stage_mask = nn.Sequential(*stage).to(device)
+
+        self.last = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+        self.last_mask = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    1,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        if self.use_skip:
+            raise NotImplementedError
+            input, features = x
+        else:
+            input = x
+        x = self.first(input)
+        f = self.layers(x)
+        x = self.stage_pose(f)
+        pose = self.last(x)
+
+        x = self.stage_mask(f)
+        mask = self.last_mask(x)
+        
+        return pose, mask
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+class PoseEncoder(nn.Module):
+    def __init__(self, features=128, device='cuda'):
+        super().__init__()
+
+                
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, features, 9),
+            CNBlockConfig(features, None, 3),
+        ]
+
+        block_setting_reverse = [
+            CNBlockConfig(features, 384, 3),
+            CNBlockConfig(384, 192, 9),
+            CNBlockConfig(192, 96, 3),
+            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(27, 9, 3),
+            CNBlockConfig(9, None, 3),
+        ]
+
+        rotation_block_setting = [
+            CNBlockConfig(64, 128, 3),
+            CNBlockConfig(128, 128, 3),
+            CNBlockConfig(128, 256, 9),
+            CNBlockConfig(256, None, 3),
+        ]
+
+        rotation_block_setting_reverse = [
+            CNBlockConfig(256, 128, 3),
+            CNBlockConfig(128, 64, 9),
+            CNBlockConfig(64, 32, 3),
+            CNBlockConfig(32, 16, 3),
+            CNBlockConfig(16, None, 3),
+        ]
+
+        self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
+        self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
+        self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
+        self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)
+
+        self.flatten = nn.Flatten()
+
+
+    def forward(self, input) -> torch.Tensor:
+        mask, pose_in, pose_out = input
+        features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
+        x = features.view(features.size(0), features.size(1), -1)
+        y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
+        rotation = self.rotation_decoder(y)
+        rotation = torch.squeeze(rotation, dim=1)
+        rotation = torch.softmax(rotation, dim=1)
+
+        x = torch.einsum('bij,bkj->bik', x, rotation)
+        x = x.view(x.size(0), x.size(1), 16, 16)
+        pose, mask = self.decoder(x)
+
+        return pose, mask, rotation
+
+
+if __name__ == '__main__':
+    # test the model
+    model = PoseEncoder()
+    # model = torch.compile(model)
+    print(model)
+    print(sum(p.numel() for p in model.parameters() if p.requires_grad))
+    print(sum(p.numel() for p in model.parameters()))
+
+    x = torch.rand((4, 1, 512, 512)).to('cuda')
+    y = torch.rand((4, 3, 512, 512)).to('cuda')
+    z = torch.rand((4, 3, 512, 512)).to('cuda')
+    output = model((x, y, z))
+
+    from torchviz import make_dot
+    graph = make_dot(output, params=dict(model.named_parameters()))
+    graph.render("model_graph", format="pdf")
+
+    pass

utils/data_utils.py

@@ -113,10 +113,10 @@ def proj(camLoc):
 	return intrinsic().dot(extrinsic(camLoc))
 
 
-def get_camera(cam, img_size, device):
+def get_camera(cam, img_size, device, res_x=640, res_y=480):
     size = np.array(img_size, dtype=np.float32)
     ext = torch.tensor(extrinsic(cam).astype(np.float32), device=device)
-    K = torch.tensor(intrinsic()[None, :, :], device=device)
+    K = torch.tensor(intrinsic(res_x, res_y)[None, :, :], device=device)
     cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(size[None, :], device=device))
     return cameras, K
 

utils/pose_coloring.py

@@ -35,14 +35,14 @@ class PoseRenderer(Module):
         super().__init__()
         colors = np.load('custom_data/pose_colors.npy')
         self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
-        self.cam = torch.tensor([5.0, 0.0, 0.0], device=device).cpu().numpy()
+        self.cam = torch.tensor([3.0, 0.0, 0.0], device=device).cpu().numpy()
 
 
     def forward(self, verts: torch.Tensor, faces: torch.Tensor):
-        cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
+        cameras, K = get_camera(self.cam, [512, 512], verts.get_device(), res_x=512, res_y=512)
 
         raster_settings = RasterizationSettings(
-                            image_size=(480, 640), 
+                            image_size=(512, 512), 
                             blur_radius=0.0, 
                             faces_per_pixel=1, 
                             bin_size=0
@@ -63,7 +63,7 @@ if __name__ == '__main__':
     smpl = SMPL(path='superSMPL', device='cuda')
     pose = torch.zeros((1,72), device='cuda')
     # generate random pose
-    pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
+    # pose[:,3:] = torch.randn((1,69), device='cuda') * 0.1
     body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False)
     colors = np.zeros((body.shape[1], 3))