general save

.gitignore

@@ -163,3 +163,5 @@ cython_debug/
 .vscode/
 temp/
 dataset/
+wandb_login.py
+custom_data/dijsktra.npy

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,125 @@
+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, folder, frames_file='frames_all.txt', device='cuda', use_cache=True):
+        self.folder = folder
+        self.device = device
+        self.frames = read_frames_file(os.path.join(folder, frames_file))
+        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.get_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()
+        pose_out[1] = np.random.uniform(-np.pi, np.pi)
+        # 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)
+
+        render = self._render_body(body[0])[0,:,:,:3]
+        render_out = self._render_body(body_out[0])[0,:,:,:3]
+
+        render = render.permute(2,0,1)
+        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)
+
+        render = render * mask
+        render_out = render_out * mask_out
+
+        # print(f'one frame took {time() - start} seconds')
+
+        data = { 
+                 'pose': render,
+                 'shape': shape,
+                 'pose_out': render_out,
+                    # 'body': body,
+                    # 'body_out': body_out,
+                    # 'render': render,
+                    # 'render_out': render_out,
+                    'mask': mask,
+                    '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 = Dataset('/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/render.py

@@ -3,10 +3,10 @@ import sys
 import torch
 import numpy as np
 
-from deodr import Camera, Scene3D
-from deodr.triangulated_mesh import TriMesh
-from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
-from deodr.pytorch.differentiable_renderer_pytorch import Scene3DPytorch, CameraPytorch
+# from deodr import Camera, Scene3D
+# from deodr.triangulated_mesh import TriMesh
+# from deodr.pytorch.triangulated_mesh_pytorch import ColoredTriMeshPytorch
+# 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

docker/requirements.txt

@@ -6,3 +6,4 @@ pandas
 pytorch_msssim
 kornia
 wandb
+trimesh

models/pose_encoder.py

@@ -1,8 +1,8 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '7'
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 
 import torch
 from torch import nn, Tensor
@@ -77,12 +77,12 @@ class ConvNeXtReverseMaskPose(nn.Module):
         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)
+        # 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,
@@ -90,12 +90,12 @@ class ConvNeXtReverseMaskPose(nn.Module):
                     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)
+        # self.last_mask = nn.Conv2d(
+        #             block_setting[-1].input_channels,
+        #             1,
+        #             kernel_size=3,
+        #             padding='same',
+        #             ).to(device)
 
 
         for m in self.modules():
@@ -114,12 +114,13 @@ class ConvNeXtReverseMaskPose(nn.Module):
         x = self.first(input)
         f = self.layers(x)
         x = self.stage_pose(f)
-        pose = self.last(x)
+        pose = torch.sigmoid(self.last(x))
 
-        x = self.stage_mask(f)
-        mask = self.last_mask(x)
+        # x = self.stage_mask(f)
+        # mask = torch.sigmoid(self.last_mask(x))
         
-        return pose, mask
+        # 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)
@@ -178,11 +179,58 @@ class PoseEncoder(nn.Module):
         rotation = torch.squeeze(rotation, dim=1)
         rotation = torch.softmax(rotation, dim=1)
 
-        x = torch.einsum('bij,bkj->bik', x, rotation)
+        # 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 = self.decoder(x)
+        pose, mask_out = self.decoder(x)
+        # 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, rotation
+        # 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)
+
+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__':

scripts/pose_encoder_train.py

@@ -0,0 +1,188 @@
+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
+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
+
+
+import wandb
+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.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='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    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')
+    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']])
+
+    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}')
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    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']])
+
+        mse = loss_fn(pred_pose, batch['pose_out'])
+        mse_mask = loss_fn_mask(pred_mask, batch['mask_out'])
+        symm_rotation = loss_fn_rotation(pred_rotation)
+        loss = mse + mse_mask + symm_rotation
+
+
+    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']])
+
+    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)).detach()
+    pred_pose = torch.permute(pred_pose, (0, 2, 3, 1)).detach()
+
+    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, } }
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    batch_size = 1
+    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)
+
+    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)
+
+
+
+
+    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)

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)
-    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)
+    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('/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/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"] = '0'
+os.environ["CUDA_VISIBLE_DEVICES"] = '3'
 # os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import torch
@@ -58,6 +58,98 @@ class PoseRenderer(Module):
         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)
+    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)
+
+    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
+
+
+def interpolate_colors(distance_matrix, colored_indices, colored_colors):
+    from scipy.spatial import distance
+    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
+
+    num_vertices = distance_matrix.shape[0]
+
+    # 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):
+        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:
+                color = color_map[min(color_map.keys(), key=lambda x: distance.euclidean(x, color))]
+
+            interpolated_colors[vertex] = color
+
+    return interpolated_colors
+
+
+
+
 if __name__ == '__main__':
 
     smpl = SMPL(path='superSMPL', device='cuda')
@@ -75,20 +167,54 @@ if __name__ == '__main__':
     # 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()
+    # 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
+
+    # dist = get_distances(body, m.edges)
+    # # dist = compute_distances(body, faces)
+    # np.save('custom_data/dijsktra.npy', dist)
+
+    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)
+
+    # 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()  
+    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.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