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

This reverts commit 05341fd48f.

.gitignore

@@ -163,3 +163,8 @@ cython_debug/
 .vscode/
 temp/
 dataset/
+wandb_login.py
+custom_data/dijsktra*
+checkpoints/
+
+*.out

data/dataset.py

@@ -1,7 +1,7 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+# os.environ["CUDA_VISIBLE_DEVICES"] = '5'
 
 import math
 import time
@@ -34,7 +34,7 @@ class Dataset(Dataset):
         self.transforms = tv.transforms.ToTensor()
         self.frames = read_frames_file(os.path.join(folder, frames_file))
         self.render = DepthRenderer()
-        self.deodr_render : Render = Render(max_depth=10)
+        # self.deodr_render : Render = Render(max_depth=10)
         self.body_faces = np.load('custom_data/smpl_faces.npy')
         self.reader = DataReader(folder, 'superSMPL')
         
@@ -318,7 +318,7 @@ if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     # dataset = Dataset('/home/workspace/local-data/train')
-    dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True)
+    dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True, device=device)
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
     for batch in train_dataloader:

data/pose_dataset.py

@@ -0,0 +1,158 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+from tqdm import tqdm
+from time import time
+from cloth3d.read import DataReader
+from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
+from superSMPL.smpl_torch import SMPL
+from data.pose_renderer import PoseRenderer
+import trimesh
+import random
+
+
+class PoseDataset(Dataset):
+    def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True, ae_mode=False):
+        self.folder = folder
+        self.device = device
+        self.ae_mode = ae_mode
+        self.frames = read_frames_file(os.path.join(folder, frames_file))
+        random.shuffle(self.frames)
+        self.reader = DataReader(folder, 'superSMPL')
+        
+        self.cache = dict() if use_cache else None
+
+        self.poses = dict()
+        self.shapes = dict()
+
+        ids = set()
+        for frame in self.frames:
+            ids.add(frame.id)
+
+        for id in tqdm(ids):
+            info = self.reader.read_info(id)
+            self.poses[id] = info['poses'].astype(np.float32)
+            self.shapes[id] = np.concatenate((info['shape'], [info['gender']])).astype(np.float32)
+
+        self.smpl = SMPL(path='superSMPL', device=device)
+        self.renderer = PoseRenderer()
+        # self.colors = np.load('custom_data/pose_colors.npy')
+
+
+    def _render_body(self, body_to_render: torch.Tensor):
+        # swap the y,z axes of the body_to_render 
+        body_to_render = body_to_render[0,:,[0,2,1]]
+        # and rotate it 90 degress along the z axis
+        body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T
+
+        out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.device))
+
+        return out
+
+
+    def __len__(self):
+        return len(self.frames)
+    
+
+    def __getitem__(self, idx):
+        frame : Frame = self.frames[idx]
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                return self.cache[frame]
+
+        pose = torch.tensor(self.poses[frame.id][:,frame.frame_nr], device=self.device)
+        shape = torch.tensor(self.shapes[frame.id], device=self.device)
+        pose[:3] = 0
+        pose_out = pose.clone()
+
+        angle = torch.rand(1) * 2 * np.pi - np.pi
+        matrix = self.smpl.rodrigues(torch.tensor([0.0, angle.numpy()[0], 0.0], device=self.device)[None, None, :]).to(self.device)[0]
+        angle = angle.to(self.device)
+
+        pose_out[1] = angle
+        # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
+        # generate random number between -np.pi, np.pi
+        
+        # start = time()
+
+        if self.ae_mode:
+            body = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
+        else:
+            body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
+            body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
+
+        # mesh = trimesh.Trimesh(vertices=body[0][0,:,:].cpu().numpy(), faces=self.smpl.f_all, process=False, validate=False)
+        # mesh.apply_transform(trimesh.transformations.rotation_matrix(angle=3.14159, direction=[0, 1, 0], point=[0, 0, 0]))
+        # body_back = torch.tensor(mesh.vertices, device=body[0].device).unsqueeze(0).to(torch.float32)
+        # body_back[0,:,0] = -body_back[0,:,0] 
+
+        # print(f'smpl took {time() - start} seconds')
+        # start = time()
+
+        render = self._render_body(body[0])[0,:,:,:3]
+        # render_back = self._render_body(body_back)[0,:,:,:3]
+        if not self.ae_mode:
+            render_out = self._render_body(body_out[0])[0,:,:,:3]
+
+        render = render.permute(2,0,1)
+        # render_back = render_back.permute(2,0,1)
+        if not self.ae_mode:
+            render_out = render_out.permute(2,0,1)
+
+        # import plotly.express as ex
+        # ex.imshow(render[0,:,:,:3].cpu().numpy()).show()
+        # ex.imshow(render_out[0,:,:,:3].cpu().numpy()).show()
+
+        mask = torch.min(render != 1, dim=0)[0][None,:,:].to(torch.float32)
+        # mask_back = torch.min(render_back != 1, dim=0)[0][None,:,:].to(torch.float32)
+        if not self.ae_mode:
+            mask_out = torch.min(render_out != 1, dim=0)[0][None,:,:].to(torch.float32)
+
+        render = render * mask
+        # render_back = render_back * mask_back
+        if not self.ae_mode:
+            render_out = render_out * mask_out
+
+        # print(f'render took {time() - start} seconds')
+
+        data = { 
+                 'pose': render,
+                 'shape': shape,
+                #  'pose_back': render_back,
+                    # 'body': body,
+                    # 'body_out': body_out,
+                    # 'render': render,
+                    # 'render_out': render_out,
+                 'mask': mask,
+                 'angle': angle,
+                 'matrix': matrix,
+               }
+        
+        if not self.ae_mode:
+            data['pose_out'] = render_out
+            data['mask_out'] = mask_out
+
+        if self.cache is not None:
+            self.cache[frame] = data
+        return data
+    
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
+
+    for batch in train_dataloader:
+        batch
+        pass
+
+    pass 

data/pose_renderer.py

@@ -0,0 +1,88 @@
+import os
+import sys
+import numpy as np
+import torch
+from torch.nn import Module
+
+from utils.data_utils import intrinsic, extrinsic, get_camera
+from pytorch3d.renderer import (
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesVertex,
+    HardPhongShader,
+)
+from pytorch3d.structures import Meshes
+
+from pytorch3d.renderer import (
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesVertex,
+    HardPhongShader,
+    AmbientLights
+)
+from pytorch3d.structures import Meshes
+
+
+class PoseRenderer(Module):
+    def __init__(self, device='cuda') -> None:
+        super().__init__()
+        # colors = np.load('custom_data/pose_colors.npy')
+        colors_ = np.load('custom_data/body_colors_interpolated.npy').astype(np.float32) / 255.0
+        color_order = np.load('custom_data/color_ordered_indices.npy')
+        colors = np.zeros_like(colors_)
+        colors[color_order] = colors_[np.arange(27578)]
+
+        self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
+        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, [512, 512], verts.get_device(), res_x=512, res_y=512)
+
+        raster_settings = RasterizationSettings(
+                            image_size=(512, 512), 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=None,
+                            max_faces_per_bin=None,
+                        )
+        
+        renderer = MeshRenderer(
+            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
+            shader=HardPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device))
+            # shader=SoftPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device))
+        )
+
+        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
+
+        return output
+    
+if __name__ == '__main__':
+    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+    from superSMPL.smpl_torch import SMPL
+    import plotly.express as ex
+    renderer = PoseRenderer()
+    smpl = SMPL(path='superSMPL', device='cuda')
+    pose = torch.rand(72, dtype=torch.float32, device='cuda')
+    pose[:3] = 0
+    shape = torch.rand(11, dtype=torch.float32, device='cuda')
+    # pose_out[1] = 0.5 # np.random.uniform(-np.pi, np.pi)
+    # generate random number between -np.pi, np.pi
+    
+    # start = time()
+
+    body = smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
+    body_to_render = body[0,:,[0,2,1]]
+    # and rotate it 90 degress along the z axis
+    body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device=body_to_render.device), body_to_render.T).T
+
+    out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body_to_render.get_device()))
+
+    ex.imshow(out[0, ..., :3].cpu().numpy()).show()
+
+    print(out)

data/render.py

@@ -3,10 +3,10 @@ import sys
 import torch
 import numpy as np
 
-from deodr import Camera, Scene3D
+# from deodr import Camera, Scene3D
-from deodr.triangulated_mesh import TriMesh
+# from deodr.triangulated_mesh import TriMesh
-from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
+# from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
-from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
+# from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from cloth3d.util import proj
@@ -106,7 +106,8 @@ from pytorch3d.renderer import (
     MeshRasterizer,  
     SoftPhongShader,
     TexturesUV,
-    TexturesVertex
+    TexturesVertex,
+    HardPhongShader,
 )
 from pytorch3d.structures import Meshes
 from pytorch3d.utils import cameras_from_opencv_projection
@@ -162,3 +163,33 @@ class DepthRenderer(nn.Module):
         # ex.imshow(depth.cpu().numpy()).show()
 
         return depth, mask, v2i, K, cameras, pix_to_face
+
+
+class PoseRenderer(nn.Module):
+    def __init__(self, device='cuda') -> None:
+        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([0.0, 0.0, 1.0], device=device)
+
+
+    def forward(self, verts: torch.Tensor, faces: torch.Tensor):
+        cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
+
+        raster_settings = RasterizationSettings(
+                            image_size=(480, 640), 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+        
+        renderer = MeshRenderer(
+            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
+            shader=HardPhongShader(device=verts.get_device(), cameras=cameras)
+        )
+        # create a pytorch3d vertex texture from the colors
+
+        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
+
+
+        return output

data/view_dataset.py

@@ -0,0 +1,55 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+from tqdm import tqdm
+from time import time
+from cloth3d.read import DataReader
+from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
+from superSMPL.smpl_torch import SMPL
+from data.pose_renderer import PoseRenderer
+
+
+class Dataset(Dataset):
+    def __init__(self, device='cuda'):
+        self.device = device
+        self.smpl = SMPL(path='superSMPL', device=device)
+
+
+    def __len__(self):
+        return int(1e6)
+    
+
+    def __getitem__(self, idx):
+        # generate a random tenssor of size 3 with values betweeen -pi and pi
+        angles = torch.rand(3) * 2 * np.pi - np.pi
+        angles = angles.to(self.device)
+        matrix = self.smpl.rodrigues(angles[None, None, :]).to(self.device)[0]
+
+        data = { 
+                 'angles': angles,
+                 'matrix': matrix,
+               }
+
+
+        return data
+    
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader
+    
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dataset = Dataset(device=device)
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
+
+    for batch in train_dataloader:
+        batch
+        pass
+
+    pass 

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.1-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

@@ -12,6 +12,7 @@ services:
       - '/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/127-data2/lhunor/data:/data-local/127_2'
       - '/data-fast/128-data1/lhunor/data:/data-local/128'
     deploy:
       resources:
@@ -38,3 +39,23 @@ services:
         reservations:
           devices:
             - capabilities: [gpu]
+
+  lhunor_torch_3:
+    image: lhunor_reconstruction:v2.0.1
+    container_name: lhunor_reconstruction_v2.0.1
+    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,5 @@ scipy
 pandas
 pytorch_msssim
 kornia
+wandb
+trimesh

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():
@@ -156,22 +159,17 @@ class ConvNeXtBase(nn.Module):
 
     def _forward_impl(self, x: Tensor) -> Tensor:
         x = self.stem(x)
-        features: List[Tensor] = []
+        x = self.layers(x)
-        for i in range(4 - 1):
-            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
+        return x, None
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
 
-class ConvNeXtReverseBase2(nn.Module):
+class ConvNeXtReverseDepthNormals(nn.Module):
     def __init__(
         self,
         block_setting: List[CNBlockConfig],
@@ -217,8 +215,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 +226,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 +275,10 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
+        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 +309,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,53 +328,17 @@ 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] = []
+
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
                 # Bottlenecks
                 stage: List[nn.Module] = []
                 for _ in range(cnf.num_layers):
@@ -375,17 +346,26 @@ class ConvNeXtReverseBase(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
-            layers.append(nn.Sequential(*stage))
+                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 +373,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 +395,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 +411,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_ae.py

@@ -0,0 +1,183 @@
+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 ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
+
+
+class ConvNeXtReversePose(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)
+
+        stage: List[nn.Module] = []
+        for _ in range(3):
+            stage.append(block(block_setting[0].input_channels, layer_scale, 0))
+        self.stage_smooth = nn.Sequential(*stage).to(device)
+
+        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)
+
+        stage: List[nn.Module] = []
+        for _ in range(cnf.num_layers):
+            stage.append(block(block_setting[-1].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.stage_smooth(input)
+        x = self.first(x)
+        f = self.layers(x)
+        x = self.stage_pose(f)
+        pose = torch.sigmoid(self.last(x))
+
+        x = self.stage_mask(f)
+        mask = torch.sigmoid(self.last_mask(x))
+        pose = pose * (mask > 0.5).to(torch.float32)
+
+        return pose, mask
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+class PoseAE(nn.Module):
+    def __init__(self, features=256, 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),
+        ]
+
+        self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
+        self.decoder = ConvNeXtReversePose(block_setting_reverse, device=device)
+
+        self.flatten = nn.Flatten()
+
+
+    def forward(self, input) -> torch.Tensor:
+        mask, pose_in = input
+        x, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
+
+        pose_out, mask_out = self.decoder(x)
+
+        return pose_out, mask_out
+
+
+if __name__ == '__main__':
+    # test the model
+    model = PoseAE()
+    # 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')
+    output = model((x, y))
+
+
+
+    pass

models/pose_encoder.py

@@ -0,0 +1,260 @@
+import os
+import sys
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+# 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)
+
+        stage: List[nn.Module] = []
+        for _ in range(3):
+            stage.append(block(block_setting[0].input_channels, layer_scale, 0))
+        self.stage_smooth = nn.Sequential(*stage).to(device)
+
+        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.stage_smooth(input)
+        x = self.first(x)
+        f = self.layers(x)
+        x = self.stage_pose(f)
+        pose = torch.sigmoid(self.last(x))
+
+        # x = self.stage_mask(f)
+        # mask = torch.sigmoid(self.last_mask(x))
+        
+        # return pose, mask
+        return pose, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device)
+
+    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=7, 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_in_back, pose_out, mask_out = input
+        features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], 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 = torch.einsum('bij,bjk->bik', x, rotation)
+        x = x.view(x.size(0), x.size(1), 16, 16)
+        pose, mask_out_pred = self.decoder(x)
+        pose = pose * mask_out
+        # mask_out = torch.round(mask_out)
+        # mask_out = (mask_out > 0.5).to(torch.int64)
+        # pose = pose * (mask_out > 0.5).to(torch.float32)
+
+        return pose, mask_out_pred, rotation
+        # x = pose
+        # return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
+
+class DumbPoseEncoder(nn.Module):
+    def __init__(self, in_channels=7, out_channels=3, device='cuda'):
+        super(DumbPoseEncoder, self).__init__()
+
+        # Encoder
+        self.encoder_conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
+        self.encoder_conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.encoder_conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+        self.encoder_conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
+        self.encoder_pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Decoder
+        self.decoder_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
+        self.decoder_upsample_2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv2 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.decoder_upsample_3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.decoder_conv4 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        # Encoder
+        mask, pose_in, pose_out = x
+        x = torch.cat([mask, pose_in, pose_out], dim=1)
+        x1 = nn.functional.relu(self.encoder_conv1(x))
+        x = self.encoder_pool(x1)
+        x = nn.functional.relu(self.encoder_conv2(x))
+        x = nn.functional.relu(self.encoder_conv3(x))
+        x = nn.functional.relu(self.encoder_conv4(x))
+
+        # Decoder
+        x = self.decoder_upsample(x)
+        x = nn.functional.relu(self.decoder_conv1(x))
+        x = self.decoder_upsample_2(x)
+        x = nn.functional.relu(self.decoder_conv2(x))
+        x = self.decoder_upsample_3(x)
+        x = nn.functional.relu(self.decoder_conv3(x))
+        x = nn.functional.sigmoid(self.decoder_conv4(x))
+
+        return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
+
+
+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

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

models/view_encoder.py

@@ -0,0 +1,108 @@
+import os
+import sys
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+import torch
+import numpy as np
+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__)) + '/..')
+
+
+
+class ViewEncoder(nn.Module):
+    def __init__(self, bottleneck_size):
+        super(ViewEncoder, self).__init__()
+        self.bottleneck_size = bottleneck_size
+        
+        self.encoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(3*3, 64),
+                nn.BatchNorm1d(64),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(64, 128),
+                nn.BatchNorm1d(128),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(128, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(256, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(512, bottleneck_size),
+                nn.BatchNorm1d(bottleneck_size),)
+        ])
+
+        # Decoder
+        self.decoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(bottleneck_size, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(512*2, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(256*2, 128),
+                nn.BatchNorm1d(128),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(128*2, 64),
+                nn.BatchNorm1d(64),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(64*2, 3*3),
+                nn.Tanh(),),
+        ])
+    
+    def forward(self, x):
+        # Encode
+        f1 = self.encoder[0](x.view(-1, 3*3))
+        f2 = self.encoder[1](f1)
+        f3 = self.encoder[2](f2)
+        f4 = self.encoder[3](f3)
+        matrix = self.encoder[4](f4)
+
+        d1 = self.decoder[0](matrix)
+        d2 = self.decoder[1](torch.cat((d1, f4), dim=1))
+        d3 = self.decoder[2](torch.cat((d2, f3), dim=1))
+        d4 = self.decoder[3](torch.cat((d3, f2), dim=1))
+        x = self.decoder[4](torch.cat((d4, f1), dim=1))
+        
+        x = x * np.pi
+        return x.view(-1, 3, 3), matrix
+        
+        # # Decode
+        # x = self.decoder(matrix)
+        # x = x * np.pi
+        # return x.view(-1, 3, 3), matrix
+    
+
+if __name__ == '__main__':
+    # test the model
+    model = ViewEncoder(bottleneck_size=128).to('cuda')
+    # 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, 3, 3)).to('cuda')
+    output = model(x)
+
+    # from torchviz import make_dot
+    # graph = make_dot(output, params=dict(model.named_parameters()))
+    # graph.render("model_graph", format="pdf")
+
+    pass

scripts/pose_ae_train.py

@@ -0,0 +1,198 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+# config_module = importlib.import_module(args['config'])
+# config = config_module.config
+if args['debug']:
+    pass
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '7'
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+
+import wandb
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.profiler import profile, record_function, ProfilerActivity
+# from torchgeometry.losses.dice import DiceLoss
+# from torchmetrics import Dice
+
+from data.pose_dataset import PoseDataset
+from utils.utils import Accumulator
+from utils.checkpoints import CheckpointSaver, load_model
+from models.pose_ae import PoseAE
+from models.convnext import ConvNeXt
+import plotly.express as ex
+
+from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
+ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
+
+
+name = 'pose_ae_2'
+
+def init_wandb():
+    run = wandb.init(project='reconstruction', name=name, notes='now with random poses', dir='/data-local/127/temp', tags=['pose'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/pose_ae_train.py')
+    artifact.add_file('models/pose_ae.py')
+    artifact.add_file('data/pose_dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+    
+
+def train_step(batch):
+    model.train()
+
+    pred_pose, pred_mask = model([batch['mask'], batch['pose']])
+
+    mse = loss_fn(pred_pose, batch['pose'])
+    mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask'].squeeze(1))
+
+    loss: torch.Tensor = mse + mse_mask
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+
+    return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask }
+
+
+def valid_step(batch):
+    model.eval()
+    with torch.inference_mode():
+        pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
+
+        mse = loss_fn(pred_pose, batch['pose_out'])
+        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
+        symm_rotation = loss_fn_rotation(pred_rotation)
+        loss = mse
+
+    return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
+
+
+def plot_everything(batch, type: str):
+    model.eval()
+    with torch.inference_mode():
+        pred_pose, pred_mask = model([batch['mask'], batch['pose']])
+
+    mask = torch.permute(batch['mask'], (0, 2, 3, 1))
+    pose = torch.permute(batch['pose'], (0, 2, 3, 1))
+    pred_mask = torch.permute(pred_mask, (0, 2, 3, 1))
+    pred_pose = torch.permute(pred_pose, (0, 2, 3, 1))
+
+    masks = torch.concat((mask, pred_mask), axis=2).squeeze(-1).cpu().numpy()
+    poses = torch.concat((pose, pred_pose), axis=2).cpu().numpy()
+    masks_imgs = [wandb.Image((masks[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
+    poses_imgs = [wandb.Image((poses[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
+
+    return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } }
+
+
+def train():
+    iteration = 0
+    epoch = 0
+    while True:
+
+        for batch in train_dataloader:
+            train_loss = train_step(batch)
+            train_accumulator.update(train_loss)
+
+            if iteration % 1000 == 0 or (iteration < 500 and iteration % 100 == 0):
+                to_log = dict({ **train_loss,  **{ 'current_step': iteration } })
+                if iteration % 5000 == 0 or (iteration < 5000 and iteration % 1000 == 0):
+                    train_figs = plot_everything(batch, 'train')
+                    to_log = { **to_log, **train_figs }
+                # if iteration == 0:
+                #     to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+                train_accumulator.reset()
+
+                # figure out when to log when using validation set as well
+                if not args['debug']:
+                    wandb.log(to_log, step=iteration)
+            # print(f'Iteration {iteration} finished')
+
+            if iteration % 10000 == 0:
+                ckpt_saver_enc(model.encoder, iteration, optimizer, scaler=None)
+                ckpt_saver_dec(model.decoder, iteration, optimizer, scaler=None)
+
+            iteration += 1
+            
+        epoch += 1
+        print(f'Epoch {epoch} finished')
+
+        # if not iteration % 1000 == 0:
+        #     continue
+
+        # for batch in valid_dataloader:
+        #     valid_loss = valid_step(batch)
+        #     valid_accumulator.update(valid_loss)
+            
+        #     to_log = dict({ **valid_loss, **{ 'current_step': iteration } })
+        #     valid_figs = plot_everything(batch, 'train')
+        #     to_log = { **to_log, **valid_figs }
+        #     valid_accumulator.reset()
+
+        # if not args['debug']:
+        #     wandb.log(to_log, step=iteration)
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d')
+        init_wandb()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    batch_size = 4
+    use_cache = False
+
+    train_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=True)
+    valid_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=True)
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model = PoseAE(device=device).to(device)
+
+    loss_fn = nn.MSELoss()
+    loss_fn_mask = nn.BCELoss(reduction='mean')
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
+    # optimizer = optim.SGD(model.parameters(), lr=0.1)
+
+    ckpt_saver_enc = CheckpointSaver('checkpoints/reconstruction/' + name + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/reconstruction/' + name + '_dec', top_n=5)
+
+# Use the profiler to track memory usage
+    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
+    #     with record_function("model_inference"):
+    
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+    
+    train()
+
+    # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
+
+    pass 

scripts/pose_encoder_train.py

@@ -0,0 +1,208 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+# config_module = importlib.import_module(args['config'])
+# config = config_module.config
+if args['debug']:
+    pass
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+
+import wandb
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.profiler import profile, record_function, ProfilerActivity
+# from torchgeometry.losses.dice import DiceLoss
+# from torchmetrics import Dice
+
+from data.pose_dataset import Dataset
+from utils.utils import Accumulator
+from models.pose_encoder import PoseEncoder as PoseEncoder
+from models.convnext import ConvNeXt
+import plotly.express as ex
+
+from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
+ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
+
+
+def init_wandb():
+    run = wandb.init(project='reconstruction', name='pose_2', notes='have the back of the body also as the main input', dir='/data-local/127/temp', tags=['pose'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/pose_encoder_train.py')
+    artifact.add_file('models/pose_encoder.py')
+    artifact.add_file('data/pose_dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+
+class SymmetricLoss(nn.Module):
+    def __init__(self):
+        super(SymmetricLoss, self).__init__()
+
+    def forward(self, x):
+        # Check the difference between the matrix and its transpose
+        diff = x - x.transpose(1, 2)
+        # Compute the Frobenius norm of the difference
+        loss = torch.norm(diff, p='fro', dim=(1, 2))
+        # Return the average loss over the batch
+        return torch.mean(loss)
+    
+
+def train_step(batch):
+    model.train()
+
+    # pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
+    pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
+
+    mse = loss_fn(pred_pose, batch['pose_out'])
+
+    loss = mse #+ mse_mask + symm_rotation
+    # print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {symm_rotation}')
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    with torch.inference_mode():
+        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
+        symm_rotation = loss_fn_rotation(pred_rotation)
+
+    return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/symm_rotation': symm_rotation }
+
+
+def valid_step(batch):
+    model.eval()
+    with torch.inference_mode():
+        pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
+
+        mse = loss_fn(pred_pose, batch['pose_out'])
+        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
+        symm_rotation = loss_fn_rotation(pred_rotation)
+        loss = mse
+
+    return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
+
+
+def plot_everything(batch, type: str):
+    model.eval()
+    with torch.inference_mode():
+        pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
+
+    mask = torch.permute(batch['mask'], (0, 2, 3, 1))
+    mask_out = torch.permute(batch['mask_out'], (0, 2, 3, 1))
+    pose = torch.permute(batch['pose'], (0, 2, 3, 1))
+    pose_back = torch.permute(batch['pose_back'], (0, 2, 3, 1))
+    pose_out = torch.permute(batch['pose_out'], (0, 2, 3, 1))
+    pred_mask = torch.permute(pred_mask, (0, 2, 3, 1))
+    pred_pose = torch.permute(pred_pose, (0, 2, 3, 1))
+
+    masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy()
+    poses = torch.concat((pose, pose_back, pose_out, pred_pose), axis=2).cpu().numpy()
+    masks_imgs = [wandb.Image((masks[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
+    poses_imgs = [wandb.Image((poses[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
+
+    # do a batchwise min-max normalization of pred_rotation
+    pred_rotation_img = pred_rotation.cpu().numpy()
+    pred_rotation_img = (pred_rotation_img - pred_rotation_img.min(axis=(1, 2,), keepdims=True)) / (pred_rotation_img.max(axis=(1, 2), keepdims=True) - pred_rotation_img.min(axis=(1, 2), keepdims=True))
+    pred_rotation_imgs = [wandb.Image(pred_rotation_img[i]) for i in range(pred_pose.shape[0])]
+
+    return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, f'{type}/rotation': pred_rotation_imgs} }
+
+
+def train():
+    iteration = 0
+    epoch = 0
+    while True:
+
+        for batch in train_dataloader:
+            train_loss = train_step(batch)
+            train_accumulator.update(train_loss)
+
+            if iteration % 1000 == 0 or (iteration < 500 and iteration % 100 == 0):
+                to_log = dict({ **train_loss,  **{ 'current_step': iteration } })
+                if iteration % 5000 == 0 or (iteration < 5000 and iteration % 1000 == 0):
+                    train_figs = plot_everything(batch, 'train')
+                    to_log = { **to_log, **train_figs }
+                # if iteration == 0:
+                #     to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+                train_accumulator.reset()
+
+                # figure out when to log when using validation set as well
+                if not args['debug']:
+                    wandb.log(to_log, step=iteration)
+            iteration += 1
+            # print(f'Iteration {iteration} finished')
+
+        epoch += 1
+        print(f'Epoch {epoch} finished')
+
+        # if not iteration % 1000 == 0:
+        #     continue
+
+        # for batch in valid_dataloader:
+        #     valid_loss = valid_step(batch)
+        #     valid_accumulator.update(valid_loss)
+            
+        #     to_log = dict({ **valid_loss, **{ 'current_step': iteration } })
+        #     valid_figs = plot_everything(batch, 'train')
+        #     to_log = { **to_log, **valid_figs }
+        #     valid_accumulator.reset()
+
+        # if not args['debug']:
+        #     wandb.log(to_log, step=iteration)
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d')
+        init_wandb()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    batch_size = 4
+    use_cache = False
+
+    train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
+    valid_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model = PoseEncoder(device=device).to(device)
+
+    loss_fn = nn.MSELoss()
+    loss_fn_mask = nn.BCELoss(reduction='mean')
+    loss_fn_rotation = SymmetricLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
+    # optimizer = optim.SGD(model.parameters(), lr=0.1)
+
+# Use the profiler to track memory usage
+    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
+    #     with record_function("model_inference"):
+    train()
+
+    # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
+
+    pass 

scripts/pose_manipulator_train.py

@@ -0,0 +1,199 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+# config_module = importlib.import_module(args['config'])
+# config = config_module.config
+if args['debug']:
+    pass
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+
+import wandb
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torch.profiler import profile, record_function, ProfilerActivity
+# from torchgeometry.losses.dice import DiceLoss
+# from torchmetrics import Dice
+
+from data.pose_dataset import PoseDataset
+from utils.utils import Accumulator
+from utils.checkpoints import CheckpointSaver, load_model
+from models.pose_manipulator import PoseManipulator
+import plotly.express as ex
+
+from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
+ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
+
+
+name = 'pose_manipulator_1'
+
+def init_wandb():
+    run = wandb.init(project='reconstruction', name=name, notes='now with random poses, 10k ae training', dir='/data-local/127/temp', tags=['pose'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/pose_manipulator_train.py')
+    artifact.add_file('models/pose_manipulator.py')
+    artifact.add_file('data/pose_dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+    
+
+def train_step(batch):
+    model.train()
+
+    pred_pose, pred_mask, expected_features, modified_features = model([batch['mask'], batch['pose'], batch['mask_out'], batch['pose_out'], batch['angle']])
+
+    mse_features = loss_fn_features(expected_features, modified_features)
+    loss: torch.Tensor = mse_features
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    with torch.inference_mode():
+        mse = loss_fn(pred_pose, batch['pose_out'])
+        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
+
+    return { 'train/loss': loss, 'train/mse': mse, 'train/mse_mask': mse_mask, 'train/mse_features': mse_features }
+
+
+def valid_step(batch):
+    model.eval()
+    with torch.inference_mode():
+        pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_back'], batch['pose_out'], batch['mask_out']])
+
+        mse = loss_fn(pred_pose, batch['pose_out'])
+        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
+        symm_rotation = loss_fn_rotation(pred_rotation)
+        loss = mse
+
+    return { 'valid/loss': loss, 'valid/mse': mse, 'valid/mse_mask': mse_mask, 'valid/symm_rotation': symm_rotation }
+
+
+def plot_everything(batch, type: str):
+    model.eval()
+    with torch.inference_mode():
+        pred_pose, pred_mask, _, _ = model([batch['mask'], batch['pose'], batch['mask_out'], batch['pose_out'], batch['angle']])
+
+    mask = torch.permute(batch['mask'], (0, 2, 3, 1))
+    mask_out = torch.permute(batch['mask_out'], (0, 2, 3, 1))
+    pose = torch.permute(batch['pose'], (0, 2, 3, 1))
+    pose_out = torch.permute(batch['pose_out'], (0, 2, 3, 1))
+    pred_mask = torch.permute(pred_mask, (0, 2, 3, 1))
+    pred_pose = torch.permute(pred_pose, (0, 2, 3, 1))
+
+    masks = torch.concat((mask, mask_out, pred_mask), axis=2).squeeze(-1).cpu().numpy()
+    poses = torch.concat((pose, pose_out, pred_pose), axis=2).cpu().numpy()
+    masks_imgs = [wandb.Image((masks[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
+    poses_imgs = [wandb.Image((poses[i] + 1) / 2.0) for i in range(pred_pose.shape[0])]
+
+    return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } }
+
+
+def train():
+    iteration = 0
+    epoch = 0
+    while True:
+
+        for batch in train_dataloader:
+            train_loss = train_step(batch)
+            train_accumulator.update(train_loss)
+
+            if iteration % 1000 == 0 or (iteration < 500 and iteration % 100 == 0):
+                to_log = dict({ **train_loss,  **{ 'current_step': iteration } })
+                if iteration % 5000 == 0 or (iteration < 5000 and iteration % 1000 == 0):
+                    train_figs = plot_everything(batch, 'train')
+                    to_log = { **to_log, **train_figs }
+                # if iteration == 0:
+                #     to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+                train_accumulator.reset()
+
+                # figure out when to log when using validation set as well
+                if not args['debug']:
+                    wandb.log(to_log, step=iteration)
+            # print(f'Iteration {iteration} finished')
+
+            if iteration % 10000 == 0:
+                ckpt_saver(model.manipulator, iteration, optimizer, scaler=None)
+
+            iteration += 1
+            
+        epoch += 1
+        print(f'Epoch {epoch} finished')
+
+        # if not iteration % 1000 == 0:
+        #     continue
+
+        # for batch in valid_dataloader:
+        #     valid_loss = valid_step(batch)
+        #     valid_accumulator.update(valid_loss)
+            
+        #     to_log = dict({ **valid_loss, **{ 'current_step': iteration } })
+        #     valid_figs = plot_everything(batch, 'train')
+        #     to_log = { **to_log, **valid_figs }
+        #     valid_accumulator.reset()
+
+        # if not args['debug']:
+        #     wandb.log(to_log, step=iteration)
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d')
+        init_wandb()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    batch_size = 4
+    use_cache = False
+
+    train_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=False)
+    valid_dataset = PoseDataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache, ae_mode=False)
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model: PoseManipulator = PoseManipulator(device=device, ae_checpoint='checkpoints/reconstruction/pose_ae_2').to(device)
+
+    loss_fn = nn.MSELoss()
+    loss_fn_features = nn.MSELoss()
+    loss_fn_mask = nn.BCELoss(reduction='mean')
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
+    # optimizer = optim.SGD(model.parameters(), lr=0.1)
+
+    ckpt_saver = CheckpointSaver('checkpoints/reconstruction/' + name, top_n=5)
+
+    # Use the profiler to track memory usage
+    # with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
+    #     with record_function("model_inference"):
+    
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+    
+    train()
+
+    # print(prof.key_averages().table(sort_by="cuda_memory_usage", row_limit=10))
+
+    pass 

scripts/view_encoder_train.py

@@ -0,0 +1,218 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import argparse
+import importlib
+
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+# config_module = importlib.import_module(args['config'])
+# config = config_module.config
+if args['debug']:
+    pass
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+
+import wandb
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from torch.utils.data import DataLoader
+# from torchgeometry.losses.dice import DiceLoss
+# from torchmetrics import Dice
+
+from data.view_dataset import Dataset
+from utils.utils import Accumulator
+from models.view_encoder import ViewEncoder
+import plotly.express as ex
+
+
+def init_wandb():
+    run = wandb.init(project='reconstruction', name='view_7', notes='add skip connections. no activation as bottleneck. different ortholoss', dir='/data-local/127/temp', tags=['view'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/view_encoder_train.py')
+    artifact.add_file('models/view_encoder.py')
+    artifact.add_file('data/view_dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+
+class OrthogonalLoss(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, tensor):
+        # Compute the Gram-Schmidt process to orthogonalize the columns of the matrix
+        matrix_size = np.sqrt(tensor.size(-1)).astype(np.int64)
+        matrix = tensor.view(-1, matrix_size, matrix_size)
+        q, _ = torch.linalg.qr(matrix)
+
+        # Compute the Frobenius norm of the difference between the original matrix and the orthogonalized matrix
+        loss = torch.norm(matrix - q, p='fro')
+
+        return loss
+    
+# class OrthogonalLoss(nn.Module):
+#     def __init__(self):
+#         super(OrthogonalLoss, self).__init__()
+
+#     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+#         # Reshape the input tensor into a square matrix
+#         matrix_size = np.sqrt(tensor.size(-1)).astype(np.int64)
+#         matrix = tensor.view(-1, matrix_size, matrix_size)
+#         # matrix = matrix / torch.norm(matrix, dim=1, keepdim=True)  # Normalize the rows
+
+#         # Compute the Gram matrix
+#         gram_matrix = torch.bmm(matrix, matrix.transpose(1, 2))
+
+#         # Compute the Frobenius norm of the difference between the Gram matrix and the identity matrix
+#         identity_matrix = torch.eye(matrix_size, device=tensor.device).unsqueeze(0).expand(batch_size, -1, -1)
+#         loss = torch.norm(gram_matrix - identity_matrix, p='fro')
+
+#         return loss
+    
+
+def train_step(batch):
+    model.train()
+
+    pred, matrix = model(batch['matrix'])
+
+    mse = loss_fn(pred, batch['matrix'])
+    ortho_loss = ortho_loss_fn(matrix)
+    loss: torch.Tensor = mse + ortho_loss
+
+    # TODO: calculate max of ortho_loss
+
+    # print(f'Loss: {loss} MSE: {mse}')
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+
+
+    return { 'train/loss': loss, 'train/mse': mse, 'train/ortho_loss': ortho_loss }
+
+
+def valid_step(batch):
+    model.eval()
+    with torch.inference_mode():
+        pred, matrix = model(batch['matrix'])
+
+        mse = loss_fn(pred, batch['matrix'])
+        ortho_loss = ortho_loss_fn(matrix)
+        loss = mse #+ ortho_loss
+
+
+    return { 'valid/loss': loss, 'valid/mse': mse, 'valid/ortho_loss': ortho_loss }
+
+
+def plot_everything(batch, type):
+    model.eval()
+    with torch.inference_mode():
+        pred, pred_matrix = model(batch['matrix'])
+
+    matrix = batch['matrix'].cpu().numpy()
+    pred = pred.detach().cpu().numpy()
+    matrix_size = np.sqrt(pred_matrix.size(-1)).astype(np.int64)
+    pred_matrix = pred_matrix.view(-1, matrix_size, matrix_size)
+    pred_matrix = pred_matrix.detach().cpu().numpy()
+    
+    matrix_str = [np.array2string(matrix[i], precision=2, suppress_small=True) for i in range(4)]
+    pred_str = [np.array2string(pred[i], precision=2, suppress_small=True) for i in range(4)]
+
+    # plot pred_matrix as image where each matrix element is a 8x8 pixel square
+    
+
+
+
+    matrix_img = [wandb.Image((np.resize(matrix[i], (matrix_size*8,matrix_size*8)))) for i in range(4)]
+
+    logs = { f'{type}/matrix': wandb.Table(data=[[matrix_str, pred_str]], columns=['matrix', 'pred']), f'{type}/matrix_img': matrix_img }
+
+    return logs
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d')
+        init_wandb()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    batch_size = 128
+    use_cache = False
+
+    train_dataset = Dataset(device=device)
+    valid_dataset = Dataset(device=device,)
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model = ViewEncoder(bottleneck_size=32**2).to(device)
+
+    loss_fn = nn.MSELoss()
+    ortho_loss_fn = OrthogonalLoss()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    # lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=False)
+    # optimizer = optim.SGD(model.parameters(), lr=0.1)
+
+    iteration = 0
+    epoch = 0
+    while True:
+
+        for batch in train_dataloader:
+            train_loss = train_step(batch)
+            train_accumulator.update(train_loss)
+
+            if iteration % 1000 == 0 or (iteration < 1000 and iteration % 100 == 0):
+                to_log = dict({ **train_loss,  **{ 'current_step': iteration } })
+                if iteration % 1000 == 0 or (iteration < 1000 and iteration % 1000 == 0):
+                    train_figs = plot_everything(batch, 'train')
+                    to_log = { **to_log, **train_figs }
+                # if iteration == 0:
+                #     to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+                if not args['debug']:
+                    wandb.log(to_log, step=iteration)
+                train_accumulator.reset()
+            iteration += 1
+            # print(f'Iteration {iteration} finished')
+
+        epoch += 1
+        print(f'Epoch {epoch} finished')
+
+        # if not iteration % 1000 == 0:
+        #     continue
+
+        # for batch in valid_dataloader:
+        #     valid_loss = valid_step(batch)
+        #     valid_accumulator.update(valid_loss)
+            
+
+        # to_log = dict({ **valid_loss, **{ 'current_step': iteration } })
+        # valid_figs = plot_everything(batch, 'train')
+        # to_log = { **to_log, **valid_figs }
+        # valid_accumulator.reset()
+
+        # if not args['debug']:
+        #     wandb.log(to_log, step=iteration)
+
+
+
+
+    pass 

speed.py

@@ -2,12 +2,11 @@ import numpy as np
 import time
 
 if __name__ == '__main__':
+
+    l = []
+    while True:
         start = time.perf_counter()
-    a, b = np.random.rand(1024, 1024), np.random.rand(1024, 1024)
+        a, b = np.random.rand(10240, 1024), np.random.rand(10240, 1024)
+        l.append([a,b])
+        c = a * b
         print(f'init time: {time.perf_counter() - start}')
-
-    start = time.perf_counter()
-
-    c = np.dot(a, b)
-
-    print(f'dot time: {time.perf_counter() - start}')

speed_2.py

@@ -0,0 +1,56 @@
+# import sys
+# import numpy as np
+# from timeit import Timer
+# print(np.__version__, sys.version)
+# # results sensitive to hardware cache size
+# n = 5
+# sizes = [2**20, 2**21, 2**22, 2**23]
+# stmts = [
+#     'np.zeros(({},))',
+#     'np.random.rand({})',
+#     'np.linspace(0,1,{})',
+#     'np.exp(np.zeros(({},)))',
+# ]
+# for stmt in stmts:
+#     for size in sizes:
+#         s = stmt.format(size)
+#         print(s+':')
+#         t = Timer(s, globals=globals()).timeit(n)
+#         print(f'\t{size/n/t} elements/second')
+#!/usr/bin/env python3
+
+
+
+import sys
+import time
+import concurrent.futures
+import numpy as np
+
+SIZE = 16 * 1024 * 1024
+PASSES = 100
+THREADS = 64
+
+print(np.__version__, sys.version)
+with open('/sys/kernel/mm/transparent_hugepage/enabled') as f:
+    print(f.read())
+
+def work(thread_id):
+    arrs = []
+    for i in range(PASSES):
+        arrs.append(np.ones(SIZE, np.uint8))
+    np.stack(arrs)
+
+with concurrent.futures.ThreadPoolExecutor(THREADS) as pool:
+    start = time.monotonic()
+    list(pool.map(work, range(THREADS)))
+    stop = time.monotonic()
+    print(stop - start)
+
+
+# import time
+# import numpy as np
+# start = time.monotonic()
+# N_GB = 10
+# a = np.ones(N_GB * 1024 * 1024 * 1024, 'ubyte')
+# stop = time.monotonic()
+# print(stop - start)

superSMPL/smpl_torch.py

@@ -50,6 +50,10 @@ class SMPL(Module):
     self.f_skirt = np.array(F['faces_skirt'])
     self.nn = np.array(sio.loadmat(os.path.join(path, 'NN.mat')))
     self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
+    self.f_all = np.asarray(female['f'])
+    # convert quad faces to triangles
+    self.f_all = np.concatenate([self.f_all[:,:3], self.f_all[:,[0,2,3]]], axis=0)
+    
 
     self.weights_all = self.weights
     self.weights = self.weights[:,self.to_keep,:]
@@ -206,7 +210,8 @@ class SMPL(Module):
     
     joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
     
-    T = torch.stack([torch.tensordot(results[i], self.weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    weights = self.weights if remove else self.weights_all
+    T = torch.stack([torch.tensordot(results[i], weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
     )

test.py

@@ -1,5 +1,6 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = '6'
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 
 import sys
 
@@ -199,7 +200,7 @@ def normals_reconst_loss(y_pred, normals_gt, mask):
 
 
 def train_step(batch):
-    # plot3d_with_mesh(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'], batch['joints'], batch['original'], batch['original_mask'])
+    plot3d_with_mesh(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'], batch['joints'], batch['original'], batch['original_mask'])
     # plot3d(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'])
     ### Training
     model.train() # train mode is on by default after construction
@@ -275,8 +276,8 @@ if __name__ == '__main__':
     batch_size = 1
     use_cache = True
 
-    train_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
+    train_dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
-    valid_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
+    valid_dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
     train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
     valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
     train_accumulator = Accumulator()

test/perspective.py

@@ -0,0 +1,149 @@
+import os
+import sys
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '5'
+
+import math
+import time
+import torch
+import kornia
+import pytorch3d
+import torchvision as tv
+import numpy as np
+import scipy.io as sio
+import plotly.express as ex
+
+from torch.utils.data import Dataset
+
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.render import Render, DepthRenderer
+from superSMPL.IO import readPC2Frame, writePC2
+from utils.data_utils import Frame, read_frames_file, mask_faces, get_normals_from_depth, get_camera, filter_info, resize_image
+from cloth3d.read import DataReader
+from cloth3d.util import quads2tris, zRotMatrix
+import kornia
+import plotly.graph_objects as go
+from data.render import intrinsic
+from simpleicp import SimpleICP, PointCloud
+
+
+def depth_to_3d(depth, depth_mask):
+    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+    pc_full = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
+    pc = [pc_full[i,:,depth_mask[i]] for i in range(len(depth_mask))]
+    # x = torch.sum(torch.stack(depth_mask), dim=0).cpu().numpy() == len(depth_mask)
+    # choice = [0,0]
+    # while x[choice[0], choice[1]] == False:
+    #     choice = (np.random.random(2) * [480, 640]).astype(int)
+    # for i in range(len(pc)):
+    #     pc[i] = pc[i] - pc_full[i][:,choice[0], choice[1]][:,None]
+    return pc
+
+
+def perspective_test(folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, fix_camera=True):
+    transforms = tv.transforms.ToTensor()
+    frames = read_frames_file(os.path.join(folder, frames_file))
+    render = DepthRenderer()
+    body_faces = np.load('custom_data/smpl_faces.npy')
+    reader = DataReader(folder, 'superSMPL')
+
+    frame : Frame = frames[0]
+
+    path = os.path.join(folder, frame.id)
+    info = reader.read_info(frame.id)
+    # info['zrot'] = torch.pi / 2 + torch.pi/4
+
+    zrot = info['zrot'] # if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
+    trans = info['trans'][:,frame.frame_nr].astype(np.float32)
+    # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
+
+    mesh_posed = reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot)
+    mesh_posed = torch.tensor(mesh_posed, device=device)
+    # mesh_posed_no_trans = reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, None, zrot)
+    # mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=device)
+
+    # raw_mask = sio.loadmat(os.path.join(path, frame.garment + '_mask.mat'))['mask'].reshape((-1))
+    # mesh_mask = raw_mask > 0 #np.load(os.path.join(path, frame.garment + '_mask.npy'))
+    # TODO: what's the difference between masking and reading .obj file. is the obj file the not registered version?
+    # faces = mask_faces(self.body_faces, mesh_mask)
+    faces = quads2tris(reader.read_garment_topology(frame.id, frame.garment))
+    # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
+    body, body_faces, root_joint, joints, joints_no_trans = reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=True)
+    body_no_trans, _, _, _, _ = reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
+    body = torch.tensor(body.astype(np.float32), device=device)
+    body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=device)
+    body_faces = quads2tris(body_faces)
+
+    joints = torch.tensor(joints[0], device=device).to(torch.float32)
+
+    offsets = [[1, -0.9, 0], [1, 0.9, 0], [1, 0, 0.4], [1, 0, -0.4], [1,-0.9,0.4], [1, 0.9, -0.4]]
+    # x: - = closer, + = farther (+-0.5)
+    # y: - = to the right, + = to the left (+-1)
+    # z: - = up, + = down (+-0.5)
+
+    combineds = []
+    for offset in offsets:
+        combineds.append(torch.concat((body, mesh_posed)) - torch.tensor(offset, device=body.get_device())[None,:])
+    # combined_no_trans = torch.concat((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))
+    combined_faces = torch.Tensor(np.concatenate((body_faces, faces + body.shape[0]))).to(device)
+
+    depth_images = []
+    depth_masks = []
+    for i, combined in enumerate(combineds):
+        start = time.perf_counter()
+        cam = info['camLoc'].copy()
+        if not fix_camera:
+            cam = cam - offsets[i]
+        depth_image, depth_mask, _, _, _, _ = render(combined, combined_faces, cam)
+        print(f'pytorch3d time: {time.perf_counter() - start}')
+        depth_images.append(depth_image)
+        depth_masks.append(depth_mask)
+
+    # depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
+    # depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask
+
+    # rgb_image = tv.io.read_image(os.path.join(folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32).to(device)
+
+    all_depth = torch.stack(depth_images)
+    pointclouds = depth_to_3d(all_depth[:,None,:,:], depth_masks)
+
+    all_depth_merged = torch.sum(all_depth, dim=0)
+    ex.imshow(all_depth_merged.cpu()).show()
+
+    mesh_data = []
+    for pc in pointclouds:
+        pc = pc.cpu().numpy()
+        mesh_data.append(go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),)
+
+    go.Figure(data = mesh_data).show()
+
+    pc_0 = pointclouds[0] - torch.mean(pointclouds[0], dim=1)[:,None]
+    pc_fix = PointCloud(torch.permute(pc_0,[1,0]).cpu().numpy(), columns=["x", "y", "z"])
+    pc_registered = [np.transpose(pc_0.cpu().numpy(), (1,0))]
+    for pc in pointclouds[1:]:
+        pc = pc - torch.mean(pc, dim=1)[:,None]
+        # pc_reg = torch.permute(pc,[1,0]).cpu().numpy()
+        pc_mov = PointCloud(torch.permute(pc,[1,0]).cpu().numpy(), columns=["x", "y", "z"])
+        icp = SimpleICP()
+        icp.add_point_clouds(pc_fix, pc_mov)
+        _, pc_reg, _, _ = icp.run(max_iterations=50, correspondences=10000, min_change=0.5, neighbors=50)
+        pc_registered.append(pc_reg)
+        # pass
+
+    mesh_data = []
+    for pc in pc_registered:
+        pc = np.transpose(pc, (1,0))
+        mesh_data.append(go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),)
+
+    go.Figure(data = mesh_data).show()
+    pass
+
+
+if __name__ == '__main__':
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    fix_camera = True
+    perspective_test('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', fix_camera=fix_camera, device=device)
+
+
+    pass 

utils/checkpoints.py

@@ -0,0 +1,43 @@
+import torch
+import numpy as np 
+import os
+import logging
+
+
+class CheckpointSaver:
+    def __init__(self, dirpath, top_n=5):
+        """
+        dirpath: Directory path where to store all model weights 
+        decreasing: If decreasing is `True`, then lower metric is better
+        top_n: Total number of models to track based on validation metric value
+        """
+        if not os.path.exists(dirpath): os.makedirs(dirpath)
+        self.dirpath = dirpath
+        self.top_n = top_n
+        self.top_model_paths = []
+        
+    def __call__(self, model, iteration, optimizer, scaler):
+        model_path = os.path.join(self.dirpath, f'iteration_{iteration}.pt')
+        torch.save({ 'model': model.state_dict(), 'optimizer': optimizer, 'scaler': scaler }, model_path)
+        self.top_model_paths.append({'path': model_path})
+        if len(self.top_model_paths)>self.top_n: 
+            self.cleanup()
+    
+    def cleanup(self):
+        to_remove = self.top_model_paths[:len(self.top_model_paths)-self.top_n]
+        logging.info(f"Removing extra models.. {to_remove}")
+        for o in to_remove:
+            os.remove(o['path'])
+        self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:]
+
+
+def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''):
+    # list of checkpoint files in the folder checkpoints
+    checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')]
+    checkpoint = torch.load(checkpoints[-1])
+    iteration = checkpoints[-1].split('_')[-1].split('.')[0]
+    model.load_state_dict(checkpoint['model'])
+    if optimizer is not None: optimizer.load_state_dict(checkpoint['optimizer'].state_dict())
+    if scaler is not None and checkpoint['scaler'] is not None: scaler.load_state_dict(checkpoint['scaler'])
+    return int(iteration) if iteration != '' else 0
+

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

@@ -0,0 +1,484 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import torch
+from torch.nn import Module
+import numpy as np
+from superSMPL.smpl_torch import SMPL
+import trimesh
+import time
+import plotly.express as ex
+import plotly.graph_objects as go
+
+
+from data_utils import intrinsic, extrinsic, get_camera
+
+from pytorch3d.renderer import (
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesVertex,
+    HardPhongShader,
+    AmbientLights
+)
+from pytorch3d.structures import Meshes
+
+
+class PoseRenderer(Module):
+    def __init__(self, device='cuda') -> None:
+        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([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, [512, 512], verts.get_device(), res_x=512, res_y=512)
+
+        raster_settings = RasterizationSettings(
+                            image_size=(512, 512), 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+        
+        renderer = MeshRenderer(
+            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
+            shader=HardPhongShader(device=verts.device, cameras=cameras, lights=AmbientLights(device=verts.device))
+        )
+
+        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
+
+        return output
+    
+
+def get_dijsktra(m: trimesh.Trimesh) -> np.ndarray:
+    edge_lengths = m.edges_unique_length
+    # create and adjacency matrix based on the edge_lengths
+    adjacency_matrix = np.zeros((m.vertices.shape[0], m.vertices.shape[0]), dtype=np.float32)
+    adjacency_matrix[m.edges_unique[:,0], m.edges_unique[:,1]] = edge_lengths
+    adjacency_matrix[m.edges_unique[:,1], m.edges_unique[:,0]] = edge_lengths
+    # calculate the distance between all vertices using dijsktra algorithm
+    from scipy.sparse import csr_matrix
+    from scipy.sparse.csgraph import dijkstra
+    dists = csr_matrix(adjacency_matrix)
+    dists = dijkstra(dists, directed=False)    
+
+    return dists
+
+
+def get_distances(vertices, edges):
+    from scipy.sparse import csr_matrix
+    from scipy.sparse.csgraph import dijkstra
+    num_vertices = vertices.shape[0]
+
+    # edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:,1]], axis=-1)
+    edge_lengths = np.ones(edges.shape[0])
+    adjacency_matrix = csr_matrix((edge_lengths.flatten(), (edges[:, 0], edges[:, 1])), shape=(num_vertices, num_vertices))
+
+    # Compute the distances between all vertices using Dijkstra's algorithm
+    # distances = dijkstra(adjacency_matrix, directed=False)
+    distances = dijkstra(adjacency_matrix, directed=False, unweighted=True)
+
+    return distances
+
+
+def get_distances2(vertices, faces):
+    from scipy.sparse import coo_matrix, csgraph
+
+    num_vertices = vertices.shape[0]
+    num_faces = faces.shape[0]
+
+    # Calculate the edge lengths
+    edge_lengths = np.linalg.norm(vertices[faces[:, np.newaxis]] - vertices[faces], axis=2)
+
+    # Create the adjacency matrix as a COO matrix
+    row_indices = np.repeat(faces.flatten(), 3)
+    col_indices = np.tile(faces.flatten(), 3)
+    data = np.tile(edge_lengths.flatten(), 3)
+    adjacency_matrix = coo_matrix((data, (row_indices, col_indices)), shape=(num_vertices, num_vertices))
+
+    # Compute the distances between all vertices using Dijkstra's algorithm
+    distances = csgraph.dijkstra(adjacency_matrix, directed=False)
+
+    return distances
+
+
+import networkx as nx
+from scipy.spatial.distance import cdist
+ 
+# def order_mesh_vertices(vertices, distances, edges):
+#     # Create a graph with vertices as nodes and distances as edge weights
+#     graph = nx.Graph()
+#     graph.add_nodes_from(range(len(vertices)))
+#     # calcualate the edge lengths
+#     edge_lengths = np.array([np.linalg.norm(vertices[edges[i, 0]] - vertices[edges[i, 1]]) for i in range(edges.shape[0])])
+
+#     # Add edges to the graph based on the mesh edges
+#     for edge, edge_length in zip(edges, edge_lengths):
+#         graph.add_edge(edge[0], edge[1], weight=edge_length)
+ 
+#     ordered_vertices = []
+#     ordered_indices = []
+ 
+#     start_vertex = 0
+#     # Perform breadth-first search starting from the selected vertex
+#     visited = set()
+#     queue = [(1221, 1221)]  # Tuple of (vertex, index)
+
+#     while queue:
+#         vertex, index = queue.pop(0)
+#         if vertex not in visited:
+#             visited.add(vertex)
+#             ordered_vertices.append(vertices[vertex])
+#             ordered_indices.append(vertex)
+#             # neighbors = sorted(graph.neighbors(vertex), key=lambda n: graph.edges[vertex, n]['weight'])
+#             neighbors = graph.neighbors(vertex)
+#             queue.extend([(n, n) for n in neighbors])
+ 
+#     return np.array(ordered_vertices), np.array(ordered_indices)
+
+from collections import deque
+def order_mesh_vertices(vertices, distances, edges):
+    # Create a graph with vertices as nodes and distances as edge weights
+    graph = nx.Graph()
+    graph.add_nodes_from(range(len(vertices)))
+    
+    # Calculate the edge lengths
+    edge_lengths = np.linalg.norm(vertices[edges[:, 0]] - vertices[edges[:, 1]], axis=1)
+    
+    # Add edges to the graph based on the mesh edges
+    for edge, edge_length in zip(edges, edge_lengths):
+        graph.add_edge(edge[0], edge[1], weight=edge_length)
+ 
+    ordered_indices = []
+ 
+    start_vertex = 13833 
+    # Perform breadth-first search starting from the selected vertex
+    visited = set()
+    queue = deque([start_vertex])
+
+    while queue:
+        vertex = queue.popleft()
+        if vertex not in visited:
+            visited.add(vertex)
+            ordered_indices.append(vertex)
+            neighbors = sorted(graph.neighbors(vertex), key=lambda n: graph.edges[vertex, n]['weight'])
+            # neighbors = graph.neighbors(vertex)
+            queue.extend(neighbors)
+ 
+    return np.array(ordered_indices)
+
+
+def interpolate_colors(distance_matrix, colored_indices, colored_colors):
+    from scipy.spatial import distance, cKDTree
+    all_colors = np.array(np.meshgrid(np.arange(0,256), np.arange(0,256), np.arange(0,256))).T.reshape(-1,3)
+    color_map = dict()
+    for idx, color in enumerate(all_colors):
+        color_map[tuple(color)] = color
+    color_map_array = np.array(list(color_map.values()))
+
+    num_vertices = distance_matrix.shape[0]
+    kd_tree = cKDTree(color_map_array)
+
+    # Create an array to store the interpolated colors
+    interpolated_colors = np.zeros((num_vertices, 3), dtype=np.uint8)
+
+    # Assign the given colors to the corresponding vertices
+    for index, color in zip(colored_indices, colored_colors):
+        interpolated_colors[index] = color
+
+    used_colors = set([tuple(color) for color in colored_colors])
+
+    # Interpolate colors for the remaining vertices
+    for vertex in range(num_vertices):
+        print(vertex)
+        if vertex not in colored_indices:
+            # start = time.time()
+            # Find the distances from the current vertex to the colored vertices
+            distances = distance_matrix[vertex, colored_indices]
+            # take the three closest one only
+            closest_indices = np.argsort(distances)[:3]
+            distances = distances[closest_indices]
+            # print(time.time() - start)
+
+            # Calculate the weights for color interpolation based on inverse distances
+            weights = 1.0 / distances
+            # Normalize the weights to sum up to 1
+            weights /= np.sum(weights)
+            # Interpolate the colors based on the weights
+            color = np.dot(weights, np.array(colored_colors)[closest_indices]).astype(np.uint8)
+
+            # get the closest color to this from color_map if this color is not available
+            if tuple(color) in used_colors:
+                start = time.time()
+                color_idx = kd_tree.query(color, workers=48)[1]
+                color = color_map_array[color_idx]
+                # if color is already in used_colors then get the next closest color until find one that is not in used_colors
+                k = 2
+                while tuple(color) in used_colors:
+                    neighbour_dist, color_idx = kd_tree.query(color, k=k, workers=48)
+                    color = color_map_array[color_idx[-1]]
+                    k += 1
+                # assert neighbour_dist[-1] < 5
+                # color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
+                print(time.time() - start)
+
+            # start = time.time()
+            interpolated_colors[vertex] = color
+            used_colors.add(tuple(color))
+            # color_map.pop(tuple(color))
+            # color_map_array = np.delete(color_map_array, np.where(np.all(color_map_array == color, axis=1))[0], axis=0)
+            # np.delete(color_map_array, np.where , axis=0)
+            # print(time.time() - start)
+
+    return interpolated_colors
+
+
+def interpolate_colors_2(distance_matrix, colored_indices, colored_colors):
+    from scipy.spatial import distance, cKDTree
+    all_colors = np.array(np.meshgrid(np.arange(0,256), np.arange(0,256), np.arange(0,256))).T.reshape(-1,3)
+    color_map = dict()
+    for idx, color in enumerate(all_colors):
+        color_map[tuple(color)] = color
+    color_map_array = np.array(list(color_map.values()))
+
+    num_vertices = distance_matrix.shape[0]
+    kd_tree = cKDTree(color_map_array)
+
+    # Create an array to store the interpolated colors
+    interpolated_colors = np.zeros((num_vertices, 3), dtype=np.uint8)
+
+    # Assign the given colors to the corresponding vertices
+    for index, color in zip(colored_indices, colored_colors):
+        interpolated_colors[index] = color
+
+    # Interpolate colors for the remaining vertices
+    for vertex in range(num_vertices):
+        print(vertex)
+        if vertex not in colored_indices:
+            # Find the distances from the current vertex to the colored vertices
+            distances = distance_matrix[vertex, colored_indices]
+
+            # Calculate the weights for color interpolation based on inverse distances
+            weights = 1.0 / distances
+
+            # Normalize the weights to sum up to 1
+            weights /= np.sum(weights)
+
+            # Interpolate the colors based on the weights
+            color = np.dot(weights, colored_colors).astype(np.uint8)
+
+            # get the closest color to this from color_map if this color is not available
+            if tuple(color) not in color_map:
+                kd_tree = cKDTree(color_map_array)
+                color_idx = kd_tree.query(color, workers=48)[1]
+                color = color_map_array[color_idx]
+                # color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
+
+            interpolated_colors[vertex] = color
+            color_map.pop(tuple(color))
+            color_map_array = np.delete(color_map_array, np.where(np.all(color_map_array == color, axis=1))[0], axis=0)
+            # np.delete(color_map_array, np.where , axis=0)
+
+    return interpolated_colors
+
+
+def get_edges(faces):
+    edges = []
+    unique_edges = set()
+
+    # Iterate over each face
+    for face in faces:
+        # Get the vertices of the current face
+        v0, v1, v2 = face
+
+        # Define the three edges of the current face
+        # and add them to the unique_edges set
+        edge0 = tuple(sorted((v0, v1)))
+        edge1 = tuple(sorted((v1, v2)))
+        edge2 = tuple(sorted((v2, v0)))
+
+        unique_edges.add(edge0)
+        unique_edges.add(edge1)
+        unique_edges.add(edge2)
+
+    # Convert the unique_edges set to a list
+    edges = list(unique_edges)
+
+    # Convert the list of edges to a numpy array
+    edges = np.array(edges)
+
+    return edges
+
+
+def get_data():
+    from plyfile import PlyData
+    # Read the .ply file
+    plydata = PlyData.read('custom_data/body_color_init.ply')
+
+    # Get the vertex coordinates
+    vertices = plydata['vertex']
+    coordinates = np.vstack((vertices['x'], vertices['y'], vertices['z'])).T
+
+    # Get the vertex colors
+    colors = np.vstack((vertices['red'], vertices['green'], vertices['blue'])).T
+
+    return coordinates, colors
+
+
+
+def find_corresponding_vertices(mesh1_vertices, mesh2_vertices, indices_mesh1):
+    corresponding_vertices_mesh2 = []
+    # mesh1_vertices[:, [1, 2]] = mesh1_vertices[:, [2, 1]]
+    for index_mesh1 in indices_mesh1:
+        vertex_mesh1 = mesh1_vertices[index_mesh1]
+        for i, vertex_mesh2 in enumerate(mesh2_vertices):
+            if np.allclose(vertex_mesh1, vertex_mesh2):
+                corresponding_vertices_mesh2.append(i)
+                break
+    return np.array(corresponding_vertices_mesh2)
+
+def transform_selected_indices(selected_indices, new_order):
+    mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(new_order)}
+    transformed_indices = np.array([mapping[idx] for idx in selected_indices])
+    return transformed_indices
+
+def transform_faces(faces, new_order):
+    mapping = {old_idx: new_idx for new_idx, old_idx in enumerate(new_order)}
+    transformed_faces = np.array([[mapping[idx] for idx in face] for face in faces])
+    return transformed_faces
+
+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
+    body, _ = smpl(pose, torch.zeros((1,11), device='cuda'), remove=False)
+    colors = np.zeros((body.shape[1], 3))
+
+
+    renderer = PoseRenderer()
+    body_to_render = body[0]
+    # swap the y,z axes of the body_to_render 
+    body_to_render = body_to_render[:,[0,2,1]]
+    # and rotate it 90 degress along the z axis
+    body_to_render = torch.matmul(torch.tensor([[0,1,0],[-1,0,0],[0,0,1]], dtype=torch.float32, device='cuda'), body_to_render.T).T
+
+    # out = renderer(body_to_render, torch.from_numpy(smpl.f_all).to(body.get_device()))
+    # ex.imshow(out[0,:,:,:3].cpu().numpy()).show()
+
+    body = body.cpu().numpy()[0]
+    faces = smpl.f_all
+
+    # m = trimesh.Trimesh(vertices=body, faces=faces, process=False, validate=False)
+    # # export the mesh to an obj file
+    # m.export('custom_data/body.obj')
+    # imoprt the mesh from the obj file
+    m = trimesh.load('custom_data/body_color_init.ply', process=False, validate=False)
+    color_init = np.array(m.visual.vertex_colors)[:,:3]
+    colored_vertices = {}
+    for idx, color in enumerate(color_init):
+        if np.all(color == 255):
+            continue
+        colored_vertices[idx] = color
+
+    # create a 3d plotly scatter plot fo body vertices with vertex 26157 highlited and limits between -1 and 1
+    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,1], z=body[:,2], mode='markers',  marker=dict(color='blue', size=1))])
+    # fig.add_trace(go.Scatter3d(x=[body[26157,0]], y=[body[26157,1]], z=[body[26157,2]], mode='markers', marker=dict(color='red', size=5)))
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    # fig.show()
+
+
+
+
+    # dist = get_distances(body, get_edges(faces))
+    # # # dist = compute_distances(body, faces)
+    # np.save('custom_data/dijsktra_int.npy', dist)
+
+    body_ply, colors_ply = get_data()
+    # rotate body_ply -90 degress along x axis
+    body_ply = body_ply[:,[0,2,1]]
+    # rotate body_ply 180 degress along y axis
+    body_ply[:,2] = -body_ply[:,2]
+
+    # colored_vertices.pop(13833)
+    # colored_vertices[1221] = np.array([0, 128, 0]).astype(np.uint8)
+    c = find_corresponding_vertices(body_ply, body, list(colored_vertices.keys()))
+    # c = np.delete(c, np.where(c == 180))
+
+    dist = np.load('custom_data/dijsktra2.npy')
+    ordered_indices = order_mesh_vertices(body, dist, get_edges(faces))
+    # reverse ordered_indices
+    # ordered_indices = ordered_indices[::-1]
+
+    c = transform_selected_indices(c, ordered_indices)
+    body = body[ordered_indices]
+    dist = dist[ordered_indices]
+    dist = dist[:,ordered_indices]
+    faces = transform_faces(faces, ordered_indices)
+    colors = interpolate_colors(dist, list(c), list(colored_vertices.values()))
+    np.save('custom_data/body_colors_interpolated_unweighted.npy', colors)
+ 
+    #c = [14110, 3192, 25113, 10822, 27513, 25872]
+    # generate a gradient grayscale colors for the vertice based ont heir indices
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = np.arange(body.shape[0]) / body.shape[0]
+    # colors[:,1] = np.arange(body.shape[0]) / body.shape[0]
+    # colors[:,2] = np.arange(body.shape[0]) / body.shape[0]
+    # colors = colors * 255
+    # colors = colors.astype(np.uint8)
+
+
+    # weights = dist[0,:]
+    # # define a gradient rgb colors for body based on weights that is not just gray
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = weights
+    # colors[:,1] = weights
+    # colors[:,2] = weights
+    # colors = colors * 255
+    # colors = colors.astype(np.uint8)
+
+    # # define a gradient rgb colors for body based on their distance from vertex 1221
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = dist[1221,:]
+    # colors[:,1] = dist[1221,:]
+    # colors[:,2] = dist[1221,:]
+    # norm = colors / np.max(colors)
+    # colors = colors * 255
+
+    # get the distances between all vertices using dijsktra algorithm without get_dijsktra
+    # from scipy.sparse import csr_matrix
+    # from scipy.sparse.csgraph import dijkstra
+
+
+
+    # define a gradinet color cube for the body that goes from green to red based on the coordinates of the vertices and not their indices
+    # colors = np.zeros((body.shape[0], 3))
+    # colors[:,0] = (body[:,0] - np.min(body[:,0])) / (np.max(body[:,0]) - np.min(body[:,0]))
+    # colors[:,1] = (body[:,1] - np.min(body[:,1])) / (np.max(body[:,1]) - np.min(body[:,1]))
+    # colors[:,2] = (body[:,2] - np.min(body[:,2])) / (np.max(body[:,2]) - np.min(body[:,2]))
+    fig = go.Figure(data=[go.Mesh3d(x=body[:,0], y=body[:,2], z=body[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], vertexcolor=colors)])
+    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    fig.show()  
+    # np.save('custom_data/pose_colors.npy', colors.astype(np.float32))
+
+    # same as a scatter plot
+    # fig = go.Figure(data=[go.Scatter3d(x=body[:,0], y=body[:,2], z=body[:,1], mode='markers', marker=dict(size=2, color=colors, opacity=0.8))])
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+    # fig.show()
+
+    # save body with colors into a ply file
+    m = trimesh.Trimesh(vertices=body, faces=faces, vertex_colors=colors)
+    m.export('custom_data/body_colors.ply')
+
+
+    
+    pass

utils/utils.py

@@ -11,9 +11,9 @@ class Accumulator():
     def update(self, dict):
         for key in dict.keys():
             if key in self.dict.keys():
-                self.dict[key].append(dict[key])
+                self.dict[key].append(dict[key].detach().cpu().numpy())
             else:
-                self.dict[key] = [dict[key]]
+                self.dict[key] = [dict[key].detach().cpu().numpy()]
 
     def get(self):
         result = dict()