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

This reverts commit 05341fd48f.

.gitignore

@@ -164,4 +164,7 @@ cython_debug/
 temp/
 dataset/
 wandb_login.py
-custom_data/dijsktra.npy
+custom_data/dijsktra*
+checkpoints/
+
+*.out

data/pose_dataset.py

@@ -14,13 +14,17 @@ 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 Dataset(Dataset):
-    def __init__(self, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
+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
@@ -48,7 +52,7 @@ class Dataset(Dataset):
         # 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.get_device()))
+        out = self.renderer(body_to_render, torch.from_numpy(self.smpl.f_all).to(body_to_render.device))
 
         return out
 
@@ -67,44 +71,73 @@ class Dataset(Dataset):
         shape = torch.tensor(self.shapes[frame.id], device=self.device)
         pose[:3] = 0
         pose_out = pose.clone()
-        pose_out[1] = np.random.uniform(-np.pi, np.pi)
+
+        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()
 
-        body = self.smpl(pose.unsqueeze(0), shape.unsqueeze(0), remove=False)
-        body_out = self.smpl(pose_out.unsqueeze(0), shape.unsqueeze(0), remove=False)
+        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_out = self._render_body(body_out[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_out = render_out.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_out = torch.min(render_out != 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_out = render_out * mask_out
+        # render_back = render_back * mask_back
+        if not self.ae_mode:
+            render_out = render_out * mask_out
 
-        # print(f'one frame took {time() - start} seconds')
+        # print(f'render took {time() - start} seconds')
 
         data = { 
                  'pose': render,
                  'shape': shape,
-                 'pose_out': render_out,
+                #  'pose_back': render_back,
                     # 'body': body,
                     # 'body_out': body_out,
                     # 'render': render,
                     # 'render_out': render_out,
-                    'mask': mask,
-                    'mask_out': mask_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
@@ -115,7 +148,7 @@ if __name__ == '__main__':
     from torch.utils.data import DataLoader
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005.txt', device=device)
+    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:

data/pose_renderer.py

@@ -30,7 +30,12 @@ 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/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()
 
@@ -42,14 +47,42 @@ class PoseRenderer(Module):
                             image_size=(512, 512), 
                             blur_radius=0.0, 
                             faces_per_pixel=1, 
-                            bin_size=0
+                            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
+        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/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,6 @@
 # Use an existing image as the base
 # FROM nvcr.io/nvidia/pytorch:23.04-py3
-FROM pytorch/pytorch:2.0.0-cuda11.7-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

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:
@@ -40,8 +41,8 @@ services:
             - capabilities: [gpu]
 
   lhunor_torch_3:
-    image: lhunor_reconstruction:v2.0
-    container_name: lhunor_reconstruction_v2.0
+    image: lhunor_reconstruction:v2.0.1
+    container_name: lhunor_reconstruction_v2.0.1
     ipc: host
     tty: true
     ulimits:

models/convnext.py

@@ -159,11 +159,10 @@ class ConvNeXtBase(nn.Module):
 
     def _forward_impl(self, x: Tensor) -> Tensor:
         x = self.stem(x)
-        features: List[Tensor] = []
         x = self.layers(x)
         x = self.act(x)
 
-        return x, features
+        return x, None
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(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

@@ -42,6 +42,11 @@ class ConvNeXtReverseMaskPose(nn.Module):
         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),
@@ -111,7 +116,8 @@ class ConvNeXtReverseMaskPose(nn.Module):
             input, features = x
         else:
             input = x
-        x = self.first(input)
+        x = self.stage_smooth(input)
+        x = self.first(x)
         f = self.layers(x)
         x = self.stage_pose(f)
         pose = torch.sigmoid(self.last(x))
@@ -162,7 +168,7 @@ class PoseEncoder(nn.Module):
             CNBlockConfig(16, None, 3),
         ]
 
-        self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
+        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)
@@ -171,8 +177,8 @@ class PoseEncoder(nn.Module):
 
 
     def forward(self, input) -> torch.Tensor:
-        mask, pose_in, pose_out = input
-        features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
+        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)
@@ -182,14 +188,15 @@ class PoseEncoder(nn.Module):
         # 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 = self.decoder(x)
+        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, 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)
+        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'):

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

@@ -1,14 +1,10 @@
 import os
 import sys
-
-os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '4' #config['gpu']
-import torch
-
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-import importlib
 import argparse
+import importlib
+
 parser = argparse.ArgumentParser(description='')
 parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
@@ -17,12 +13,20 @@ 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
 
@@ -37,15 +41,11 @@ ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    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/cvae.py')
-    artifact.add_file('losses/losses.py')
-    artifact.add_file('models/base.py')
-    artifact.add_file('models/cond_encoder.py')
-    artifact.add_file('models/convnext.py')
-    artifact.add_file('models/cvae_model.py')
-    artifact.add_file('data/dataset.py')
+    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'
@@ -70,59 +70,116 @@ class SymmetricLoss(nn.Module):
 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_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'])
-    mse_mask = 0 #loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
-    symm_rotation = 0 #loss_fn_rotation(pred_rotation)
+
     loss = mse #+ mse_mask + symm_rotation
-    print(f'Loss: {loss} MSE: {mse} MSE Mask: {mse_mask} Symm Rotation: {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_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'])
-        mse_mask = loss_fn_mask(pred_mask, batch['mask_out'])
+        mse_mask = loss_fn_mask(pred_mask.squeeze(1), batch['mask_out'].squeeze(1))
         symm_rotation = loss_fn_rotation(pred_rotation)
-        loss = mse + mse_mask + symm_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):
-    pred_pose, pred_mask, pred_rotation = model([batch['mask'], batch['pose'], batch['pose_out']])
+    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)).detach()
-    pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)).detach()
+    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)
-    poses = torch.concat((pose, pose_out, pred_pose), axis=2)
-    masks_imgs = [wandb.Image((masks[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
-    poses_imgs = [wandb.Image((poses[i].cpu().numpy() + 1) / 2.0) for i in range(pred_pose.shape[0])]
-    return { ** {f'{type}/masks': masks_imgs, f'{type}/poses': poses_imgs, } }
+    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 = 1
+    batch_size = 4
     use_cache = False
 
     train_dataset = Dataset('/data-net/dataset/cloth3d/train', 'frames_00005_single.txt', device=device, use_cache=use_cache)
@@ -141,48 +198,11 @@ if __name__ == '__main__':
     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 % 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)) } }
-
-                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)
-
-
+# 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 

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/pose_coloring.py

@@ -3,7 +3,7 @@ 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'
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 # os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import torch
@@ -11,7 +11,7 @@ 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
 
@@ -78,11 +78,13 @@ def get_distances(vertices, edges):
     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.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)
+    distances = dijkstra(adjacency_matrix, directed=False, unweighted=True)
 
     return distances
 
@@ -108,14 +110,147 @@ def get_distances2(vertices, faces):
     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
+    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)
@@ -126,6 +261,7 @@ def interpolate_colors(distance_matrix, colored_indices, colored_colors):
 
     # 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]
@@ -141,14 +277,83 @@ def interpolate_colors(distance_matrix, colored_indices, colored_colors):
 
             # get the closest color to this from color_map if this color is not available
             if tuple(color) not in color_map:
-                color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
+                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__':
 
@@ -185,21 +390,69 @@ if __name__ == '__main__':
             continue
         colored_vertices[idx] = color
 
-    # dist = get_distances(body, m.edges)
-    # # dist = compute_distances(body, faces)
-    # np.save('custom_data/dijsktra.npy', dist)
+    # 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 = np.load('custom_data/dijsktra.npy')
-    colors = interpolate_colors(dist, list(colored_vertices.keys()), list(colored_vertices.values()))
 
-    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)
+
+
+    # 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
@@ -212,7 +465,7 @@ if __name__ == '__main__':
     # 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[:,2], k=faces[:,1], vertexcolor=colors)])
+    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))
@@ -222,6 +475,10 @@ if __name__ == '__main__':
     # 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()