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

This reverts commit 05341fd48f.

.gitignore

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

Dockerfile

@@ -1,15 +0,0 @@
-FROM  pytorch/pytorch:1.12.0-cuda11.3-cudnn8-devel
-
-RUN apt update && apt upgrade -y
-RUN apt install git -y
-
-RUN git clone https://gitea.lnrweb.com/luni/reconstruction.git
-
-RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
-RUN conda install -c bottler nvidiacub -y
-RUN conda install pytorch3d -c pytorch3d -y
-
-COPY requirements.txt /tmp/
-RUN pip install --requirement /tmp/requirements.txt
-
-WORKDIR /home/workspace

cloth3d/read.py

@@ -57,11 +57,11 @@ class DataReader:
     - V: human mesh vertices
     - F: mesh faces
     """
-    def read_human(self, sample, frame, trans, zrot):
+    def read_human(self, sample, frame, zrot, use_trans=True):
         # SMPL parameters
         gender, pose, shape, trans = self.read_smpl_params(sample, frame)
         # Compute SMPL
-        V, J = self.smpl[gender].set_params(pose=pose, beta=shape, trans=trans)
+        V, J = self.smpl[gender].set_params(pose=pose, beta=shape, trans=trans if use_trans else None)
         V -= J[0:1]
         # Apply rotation on z-axis
         zRot = zRotMatrix(zrot)
@@ -80,7 +80,7 @@ class DataReader:
     """
     def read_garment_vertices(self, sample, garment, frame, trans, zrot):
         # Read garment vertices (relative to root joint)
-        pc16_path = os.path.join(self.SRC, sample, garment + '_registered.pc16')
+        pc16_path = os.path.join(self.SRC, sample, garment + '.pc16')
         V = readPC2Frame(pc16_path, frame, True)
         # Read sample data
         if trans is not None:

data/dataset.py

@@ -1,12 +1,13 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+# 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
@@ -18,21 +19,22 @@ 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
+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
+from cloth3d.util import quads2tris, zRotMatrix
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, random_view=True):
         self.folder = folder
         self.folder_rgb = folder_rgb
         self.device = device
+        self.random_view = random_view
 
         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')
         
@@ -52,45 +54,129 @@ class Dataset(Dataset):
 
         path = os.path.join(self.folder, frame.id)
         info = self.reader.read_info(frame.id)
-        # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
-        mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot'])
-        # mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False)
+        info['zrot'] = torch.pi / 2 + torch.pi/4
 
-        mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
+        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 = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, trans, zrot)
+        mesh_posed = torch.tensor(mesh_posed, device=self.device)
+        mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, None, zrot)
+        mesh_posed_no_trans = torch.tensor(mesh_posed_no_trans, device=self.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(self.reader.read_garment_topology(frame.id, frame.garment))
+        # faces = mask_faces(self.body_faces, mesh_mask)
+        faces = quads2tris(self.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 = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot'])
-        # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
+        body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=True)
+        body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, zrot, use_trans=False)
+        body = torch.tensor(body.astype(np.float32), device=self.device)
+        body_no_trans = torch.tensor(body_no_trans.astype(np.float32), device=self.device)
         body_faces = quads2tris(body_faces)
 
-        combined = torch.Tensor(np.concatenate((body, mesh_posed[mesh_mask, :]))).to(self.device)
-        # combined_no_trans = torch.Tensor(np.concatenate((body_no_trans, mesh_posed_no_trans[mesh_mask, :]))).to('cuda')
+        joints = torch.tensor(joints[0], device=self.device).to(torch.float32)
+
+        if self.random_view:
+            x_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
+            y_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi) / 10
+            z_angle = (torch.rand(1, device=self.device) * 2 * torch.pi - torch.pi)
+            view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((x_angle, y_angle, z_angle)))
+
+            # zRot = torch.tensor(zRotMatrix(-zrot), device=self.device)
+            # mesh_posed = torch.matmul(zRot, mesh_posed.T).T
+            # mesh_posed = mesh_posed - torch.tensor(trans, device=self.device)
+            # body = torch.matmul(zRot, body.T).T
+            # body = body - torch.tensor(trans, device=self.device)
+            mesh_posed = mesh_posed - joints[0,:][None,:]
+            body = body - joints[0,:][None,:]
+
+
+            transform = pytorch3d.transforms.Rotate(view)
+            mesh_posed = transform.transform_points(mesh_posed)
+            # mesh_posed_no_trans = transform.transform_points(mesh_posed_no_trans)
+            body = transform.transform_points(body)
+            # body_no_trans = transform.transform_points(body_no_trans)
+
+
+            # zRot = torch.tensor(zRotMatrix(zrot), device=self.device)
+            # mesh_posed = mesh_posed + torch.tensor(trans, device=self.device)
+            # mesh_posed = torch.matmul(zRot, mesh_posed.T).T
+            # body = body + torch.tensor(trans, device=self.device)
+            # body = torch.matmul(zRot, body.T).T
+            mesh_posed = mesh_posed + joints[0,:][None,:]
+            body = body + joints[0,:][None,:]
+
+        else:
+            view = torch.eye(3, device=self.device)
+
+        combined = torch.concat((body, mesh_posed))
+        # 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(self.device)
 
+
+
         start = time.perf_counter()
         # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
         # depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
-        deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32))
+        # deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined.cpu().numpy(), faces=combined_faces.cpu().numpy().astype(np.int32))
         # deodr_depth = self.deodr_render.render(cam=info['camLoc'].astype(np.float32), mesh=combined, faces=combined_faces.cpu().numpy().astype(np.int32))
         print(f'deodr time: {time.perf_counter() - start}')
 
         start = time.perf_counter()
         depth_image, depth_mask, v2i, K, camera, pix_to_face = self.render(combined, combined_faces, info['camLoc'])
-        # cam_no_trans = info['camLoc'].copy()
-        # cam_no_trans[1:] = 0
+        cam_no_trans = info['camLoc'].copy()
+        cam_no_trans[1:] = 0
         # depth_image_no_trans, depth_mask_no_trans, _, _, camera_no_trans, _ = self.render(combined_no_trans, combined_faces, cam_no_trans)
+        # camera_no_trans, _ = get_camera(cam_no_trans, [480, 640], self.device)
         print(f'pytorch3d time: {time.perf_counter() - start}')
 
+        
         start = time.perf_counter()
         # TODO: is it sure I need this K, and not some transformed one?
-        normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
+        normal_image = get_normals_from_depth(depth_image[None,None,:,:], depth_mask[None,:,:], K)
+        normal_image = torch.permute(normal_image[0], [2,0,1])
         print(f'normals time: {time.perf_counter() - start}')
 
         depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
         depth_mask_garment = (pix_to_face > body_faces.shape[0]) * depth_mask
 
+        top_left = root_joint.astype(np.float32)
+        top_left[:,1] -= 1.25; top_left[:,2] += 1.25
+        bottom_right = root_joint.astype(np.float32)
+        bottom_right[:,1] += 1.25; bottom_right[:,2] -= 1.25
+        corners = torch.tensor(np.concatenate((top_left, bottom_right), axis=0), device=self.device)
+        corners_pixel = camera.transform_points_screen(corners)[:,:2].round().to(torch.int64)
+        root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=self.device))[0,:2].round().to(torch.int64)
+        
+        cropper = lambda x: tv.transforms.functional.crop(x, corners_pixel[0,1], corners_pixel[0,0], corners_pixel[1,1] - corners_pixel[0,1], corners_pixel[1,0] - corners_pixel[0,0])
+        resizer_mask = tv.transforms.Resize((256,256), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
+        resizer = lambda img, mask_croped, mask_resized: resize_image(img, mask_croped, mask_resized, (256, 256))
+
+        # bbox_before = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:]
+
+        depth_mask_garment = cropper(depth_mask_garment)
+        scale = depth_mask_garment.shape[0] / 256
+        depth_mask_garment = resizer_mask(depth_mask_garment[None,None,:,:])[0,0,:,:]
+
+        original = depth_image
+        original_mask = depth_mask
+
+        depth_mask_original = depth_mask
+        depth_mask_croped = cropper(depth_mask)
+        depth_mask = resizer_mask(depth_mask_croped[None,None,:,:])[0,0,:,:]
+        scaled = resizer(cropper(depth_image)[None,None,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,0,:,:]
+        normal_image = resizer(cropper(normal_image)[None,:,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,:,:,:]
+
+        # bbox_after = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:]
+        # scale = (bbox_before[2] - bbox_before[0]) / (bbox_after[2] - bbox_after[0])
+
+        root_dist = torch.tensor(np.sqrt(np.sum((root_joint[0] - info['camLoc']) ** 2)), device=self.device).to(torch.float32)
+        # root_dist = torch.tensor(np.abs(root_joint[0][0] - info['camLoc'][0]), device=self.device).to(torch.float32)
+        # scaled = root_dist - scaled
+        scaled = scaled - root_dist
 
         # root_joint_pixel = camera.transform_points_screen(torch.tensor(root_joint.astype(np.float32), device=combined.get_device()))[0,:2].round()
         # joints_pixel = camera.transform_points_screen(torch.tensor(joints.astype(np.float32), device=combined.get_device()))[0,:,:2]
@@ -115,7 +201,7 @@ class Dataset(Dataset):
         # bbox[:2] = bbox[:2] - offset
         # dims = dims + torch.abs(offset)
 
-        # cropped_2 = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), dims[1].to(torch.int64), dims[0].to(torch.int64))
+        # # cropped_2 = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), dims[1].to(torch.int64), dims[0].to(torch.int64))
 
 
         # left = torch.tensor([0], device=bbox.get_device())
@@ -142,24 +228,28 @@ class Dataset(Dataset):
         # upscaled = torch.nn.functional.pad(scaled, (left, right, top, bottom), value=-1)
 
 
-        # Testing
+        # Verifying
+
         # pc = kornia.geometry.depth.depth_to_3d(depth_image[None,None,:,:], K, normalize_points=True)
         # pc = pc[0,:,depth_mask].cpu().numpy()
         
-        all_faces = combined_faces
-        visible_faces_idx = pix_to_face * depth_mask
-        visible_faces_idx = visible_faces_idx[torch.nonzero(visible_faces_idx, as_tuple=True)]
-        visible_vertices = torch.unique(torch.flatten(all_faces[visible_faces_idx, :])).to(torch.long)
-        visible_vertices_mask = (v2i[visible_vertices, 1] < depth_image.shape[0]) * (v2i[visible_vertices, 0] < depth_image.shape[1])
-        visible_vertices = visible_vertices[visible_vertices_mask]
-        visible_pixels_mask = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]] != -1
-        visible_vertices = visible_vertices[visible_pixels_mask]
+        # all_faces = combined_faces
+        # visible_faces_idx = pix_to_face * depth_mask
+        # visible_faces_idx = visible_faces_idx[torch.nonzero(visible_faces_idx, as_tuple=True)]
+        # visible_vertices = torch.unique(torch.flatten(all_faces[visible_faces_idx, :])).to(torch.long)
+        # visible_vertices_mask = (v2i[visible_vertices, 1] < depth_image.shape[0]) * (v2i[visible_vertices, 0] < depth_image.shape[1])
+        # visible_vertices = visible_vertices[visible_vertices_mask]
+        # visible_pixels_mask = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]] != -1
+        # visible_vertices = visible_vertices[visible_pixels_mask]
 
-        x = combined[visible_vertices, :]
-        y = np.array(info['camLoc'], dtype=np.float32)
-        dist = torch.sqrt(torch.sum((x - torch.tensor(y[None,:], device=x.get_device())) ** 2, axis=-1))
-        dist_img = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]]
-        diff = torch.mean(torch.abs(dist - dist_img))
+        # x = combined[visible_vertices, :]
+        # y = np.array(info['camLoc'], dtype=np.float32)
+        # dist = torch.sqrt(torch.sum((x - torch.tensor(y[None,:], device=x.get_device())) ** 2, axis=-1))
+        # dist_img = depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]]
+        # diff = torch.mean(torch.abs(dist - dist_img))
+
+
+        # Plotting
 
         # combined = combined.cpu().numpy()
         # depth_image = depth_image.cpu().numpy()
@@ -190,24 +280,36 @@ class Dataset(Dataset):
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
         # ex.imshow(v2i_coverage).show()
 
-        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
+        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32).to(self.device)
+        rgb_image = resizer(cropper(rgb_image)[None,:,:,:], depth_mask_croped[None,:,:], depth_mask[None,:,:])[0,:,:,:]
 
-        view = torch.tensor(np.concatenate((np.array([info['zrot']], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
+        # view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), trans.astype(np.float32)), axis=0), device=self.device)
+
+        info = filter_info(info, frame.frame_nr)
 
         # image = self.transforms(image)
         data = { 'mesh_posed': mesh_posed,
-                 'depth_image': depth_image, 
-                 'mesh_mask': mesh_mask, 
-                 'depth_mask': depth_mask, 
+                 'depth_image': scaled, 
+                #  'depth_image': upscaled, 
+                 'depth_mask': depth_mask_garment, #depth_mask,  depth_mask_garment
+                 'depth_mask_full': depth_mask, #depth_mask, 
                  'normal_image': normal_image,
-                 'depth_mask_garment': depth_mask_garment,
-                 'depth_mask_body': depth_mask_body,
+                #  'depth_mask_garment': depth_mask_garment,
+                #  'depth_mask_body': depth_mask_body,
                  'rgb_image': rgb_image,
                  'view': view,
+                 'info': info,
+                 'root_dist': root_dist,
+                 'scale': scale,
+                 'root_joint_pixel': root_joint_pixel,
+                 'joints': joints,
+                 'original': original,
+                 'original_mask': original_mask,
                 }
 
+        if self.cache is not None:
             self.cache[frame] = data
-        return self.cache[frame]
+        return data
     
 
 if __name__ == '__main__':
@@ -216,7 +318,7 @@ if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     # dataset = Dataset('/home/workspace/local-data/train')
-    dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
+    dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True, device=device)
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
     for batch in train_dataloader:

data/pose_dataset.py

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

data/pose_renderer.py

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

data/preprocess.py

@@ -9,6 +9,7 @@ from tqdm import tqdm
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from utils.data_utils import mask_faces, Frame
 from superSMPL.IO import readPC2, writePC2
+from cloth3d.util import loadInfo
 
 
 def generate_faces(folder, frames):
@@ -30,12 +31,11 @@ def generate_frames_file(folder, file_name):
         garment_dir = os.path.join(folder, id)
         if not os.path.isdir(garment_dir):
             continue
-        garments = [f.replace('_registered.pc16','') for f in os.listdir(garment_dir) if f.endswith('_registered.pc16') and not f.startswith('Combined')]
-        assert len(garments) == 1
-        garment = garments[0]
-        verts = readPC2(os.path.join(garment_dir, garment + '_registered.pc16'), True)['V']
-        for i in range(verts.shape[0]):
-            frames.append(Frame(id, i, garment))
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        info = loadInfo(os.path.join(garment_dir, 'info'))
+        garments_str = ' '.join(garments)
+        for i in range(info['poses'].shape[1]):
+            frames.append(Frame(id, i, garments_str))
 
     with open(os.path.join(folder, file_name), 'w') as f:
         for i, frame in enumerate(frames):
@@ -72,7 +72,7 @@ def genereated_registered_and_mask(folder, is_skirt = False):
 if __name__ == '__main__':
     # folder = '/home/workspace/local-data/train'
     # folder = '/home/workspace/local-data/valid'
-    folder = 'dataset/CLOTH3D/train'
+    folder = '/home/workspace/data/dataset/cloth3d/train'
 
     # generate_faces(folder)
     # genereated_registered_and_mask(folder)

data/render.py

@@ -3,14 +3,14 @@ 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
-
+from utils.data_utils import intrinsic, extrinsic, get_camera
 
 class RenderCore:
     def __init__(self, max_depth=0, depth_scale=1):
@@ -70,45 +70,6 @@ class RenderCore:
         return depth
 
 
-""" CAMERA """
-def intrinsic():
-	RES_X = 640
-	RES_Y = 480
-	f_mm             = 50 # blender default
-	sensor_w_mm      = 36 # blender default
-	sensor_h_mm = sensor_w_mm * RES_Y / RES_X
-
-	fx_px = f_mm * RES_X / sensor_w_mm
-	fy_px = f_mm * RES_Y / sensor_h_mm
-
-	u = RES_X / 2
-	v = RES_Y / 2
-
-	return np.array([[fx_px, 0,     u],
-					 [0,     fy_px, v],
-					 [0,     0,     1]], np.float32)
-				
-
-def extrinsic(camLoc):
-	R_w2bc = np.array([[0, 1, 0],
-					   [0, 0, 1],
-					   [1, 0, 0]], np.float32)
-	
-	T_w2bc = -1 * R_w2bc.dot(camLoc)
-	
-	R_bc2cv = np.array([[1,  0,  0],
-						[0, -1,  0],
-						[0,  0, -1]], np.float32)
-	
-	R_w2cv = R_bc2cv.dot(R_w2bc)
-	T_w2cv = R_bc2cv.dot(T_w2bc)
-
-	return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1)
-	
-def proj(camLoc):
-	return intrinsic().dot(extrinsic(camLoc))
-
-
 class Render:
     def __init__(self, max_depth=0, depth_scale=1, size=(640, 480)):
         self.render_core = RenderCore(max_depth=max_depth)
@@ -145,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
@@ -161,16 +123,17 @@ class DepthRenderer(nn.Module):
         # R, T = look_at_view_transform(2.7, 0, 0) 
         # cameras = FoVPerspectiveCameras(device='cuda', R=R, T=T)
 
-        ext = torch.tensor(extrinsic(cam).astype(np.float32), device='cuda')
+        # ext = torch.tensor(extrinsic(cam).astype(np.float32), device='cuda')
         # ext = proj(cam)
         # fp = np.array([240 * (480 / 36), 320 * (640 / 36)], dtype=np.float32)
         img_size = np.array([480, 640], dtype=np.float32)
         # K = np.eye(4, dtype=np.float32)
         # K[:3,:3] = intrinsic()
-        K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+        # K = torch.tensor(intrinsic()[None, :, :], device='cuda')
         # cameras = PerspectiveCameras(device='cuda', R=ext[None, :, :3], T=ext[None, :, 3], focal_length=1,)# principal_point=fp[None, :], image_size=img_size[None, :])
 
-        cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(img_size[None, :], device='cuda'))
+        # cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(img_size[None, :], device='cuda'))
+        cameras, K = get_camera(cam, [480, 640], verts.get_device())
 
         raster_settings = RasterizationSettings(
                             image_size=(480, 640), 
@@ -184,13 +147,13 @@ class DepthRenderer(nn.Module):
                     raster_settings=raster_settings
                 )
         output = rasterizer(Meshes(verts[None,:,:], faces[None,:,:]))
-        v2i = cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
+        v2i = None #cameras.transform_points_screen(verts).round().to(torch.long)[:,:2]
         depth = output.zbuf[0,:,:,0]
         pix_to_face = output.pix_to_face[0,:,:,0]
 
         mask = depth != -1
         # dmin = depth.min()
-        dmean = depth[mask].mean()
+        # dmean = depth[mask].mean()
         # dmax = depth[mask].max()
         # depth[mask] = depth[mask] - dmean
         # depth[~mask] = 0.0
@@ -200,3 +163,33 @@ class DepthRenderer(nn.Module):
         # ex.imshow(depth.cpu().numpy()).show()
 
         return depth, mask, v2i, K, cameras, pix_to_face
+
+
+class PoseRenderer(nn.Module):
+    def __init__(self, device='cuda') -> None:
+        super().__init__()
+        colors = np.load('custom_data/pose_colors.npy')
+        self.textures = TexturesVertex(verts_features=torch.from_numpy(colors).to(device)[None, :, :])
+        self.cam = torch.tensor([0.0, 0.0, 1.0], device=device)
+
+
+    def forward(self, verts: torch.Tensor, faces: torch.Tensor):
+        cameras, K = get_camera(self.cam, [480, 640], verts.get_device())
+
+        raster_settings = RasterizationSettings(
+                            image_size=(480, 640), 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+        
+        renderer = MeshRenderer(
+            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
+            shader=HardPhongShader(device=verts.get_device(), cameras=cameras)
+        )
+        # create a pytorch3d vertex texture from the colors
+
+        output = renderer(Meshes(verts[None,:,:], faces[None,:,:], textures=self.textures))
+
+
+        return output

data/view_dataset.py

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

docker/Dockerfile

@@ -0,0 +1,50 @@
+# Use an existing image as the base
+# FROM nvcr.io/nvidia/pytorch:23.04-py3
+FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
+# FROM pytorch/pytorch:1.12.0-cuda11.3-cudnn8-devel
+
+# Add the current user to the image
+ARG USER_NAME=lhunor
+ARG USER_ID=11017
+ARG GROUP_NAME=guest
+ARG GROUP_ID=11000
+
+# Update the system and install required packages
+RUN apt-get update && apt-get upgrade -y && apt-get install -y sudo git apt-utils htop
+
+# Create a new group with the same GID as the host
+RUN groupadd --gid $GROUP_ID $GROUP_NAME
+
+# Create a new user with the same UID and GID as the host
+RUN useradd --uid $USER_ID --gid $GROUP_ID --shell /bin/bash --create-home $USER_NAME
+
+# don't require password with sudo, for convenience
+# not the safest thing to do, but hopefully okay inside the container
+RUN echo '%sudo ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
+
+# Add the new user to the sudo group
+RUN usermod -aG sudo $USER_NAME
+
+# install conda stuff still as root
+RUN conda update -n base conda
+RUN conda install -c fvcore -c iopath -c conda-forge fvcore iopath -y
+RUN conda install -c bottler nvidiacub -y
+RUN conda install pytorch3d -c pytorch3d -y
+
+# Set the default user to the new user
+USER $USER_NAME
+
+ENV PATH="$PATH:/home/$USER_NAME/.local/bin"
+RUN pip install --upgrade pip
+COPY requirements.txt /tmp/
+RUN pip install --requirement /tmp/requirements.txt
+
+# defining the volumes that will be mounted
+VOLUME [ "/data-net" ]
+VOLUME [ "/data-local" ]
+VOLUME [ "/home-net" ]
+
+WORKDIR /workspace
+
+# Start the container with a bash shell
+CMD ["/bin/bash"]

docker/docker-compose.yml

@@ -0,0 +1,61 @@
+services:
+  lhunor_torch:
+    image: lhunor_reconstruction:v1.13
+    container_name: lhunor_reconstruction_v1.13
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      #- '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/127-data2/lhunor/data:/data-local/127_2'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
+
+  lhunor_torch_2:
+    image: lhunor_reconstruction:v1.12
+    container_name: lhunor_reconstruction_v1.12
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      #- '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]
+
+  lhunor_torch_3:
+    image: lhunor_reconstruction:v2.0.1
+    container_name: lhunor_reconstruction_v2.0.1
+    ipc: host
+    tty: true
+    ulimits:
+      memlock: -1
+      stack: 67108864
+    volumes:
+      - '/home-net/lhunor:/home-net'
+      - '/data-net/guest/lhunor/data:/data-net'
+      #- '/data-fast/126-data1/lhunor/data:/data-local/126'
+      - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/128-data1/lhunor/data:/data-local/128'
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - capabilities: [gpu]

docker/requirements.txt

@@ -0,0 +1,9 @@
+plotly
+Cython
+scipy
+#deodr
+pandas
+pytorch_msssim
+kornia
+wandb
+trimesh

models/convnext.py

@@ -94,6 +94,7 @@ class ConvNeXtBase(nn.Module):
         layer_scale: float = 1e-6,
         block: Optional[Callable[..., nn.Module]] = None,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
+        input_features: int = 3,
         device: str = 'cuda',
     ) -> None:
         super().__init__()
@@ -115,7 +116,7 @@ class ConvNeXtBase(nn.Module):
         # Stem
         firstconv_output_channels = block_setting[0].input_channels
         self.stem = Conv2dNormActivation(
-                3,
+                input_features,
                 firstconv_output_channels,
                 kernel_size=4,
                 stride=4,
@@ -125,6 +126,7 @@ class ConvNeXtBase(nn.Module):
                 bias=True,
             ).to(device)
         
+        layers_list: List[nn.Module] = []
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
@@ -136,16 +138,17 @@ class ConvNeXtBase(nn.Module):
                 sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                 stage.append(block(cnf.input_channels, layer_scale, sd_prob))
                 stage_block_id += 1
-            self.stages.append(nn.Sequential(*stage).to(device))
+            layers_list.append(nn.Sequential(*stage))
             if cnf.out_channels is not None:
                 # Downsampling
-                self.downsamples.append(
+                layers_list.append(
                     nn.Sequential(
                         norm_layer(cnf.input_channels),
                         nn.Conv2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
                     ).to(device)
                 )
 
+        self.layers = nn.Sequential(*layers_list).to(device)
         self.act = nn.Tanh()
 
         for m in self.modules():
@@ -156,22 +159,17 @@ class ConvNeXtBase(nn.Module):
 
     def _forward_impl(self, x: Tensor) -> Tensor:
         x = self.stem(x)
-        features: List[Tensor] = []
-        for i in range(4 - 1):
-            x = self.stages[i](x)
-            features.append(x)
-            x = self.downsamples[i](x)
-        x = self.stages[-1](x)
+        x = self.layers(x)
         x = self.act(x)
 
-        return x, features
+        return x, None
 
     def forward(self, x: Tensor) -> Tensor:
         return self._forward_impl(x)
 
 
 
-class ConvNeXtReverseBase2(nn.Module):
+class ConvNeXtReverseDepthNormals(nn.Module):
     def __init__(
         self,
         block_setting: List[CNBlockConfig],
@@ -180,9 +178,12 @@ class ConvNeXtReverseBase2(nn.Module):
         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])):
@@ -214,8 +215,8 @@ class ConvNeXtReverseBase2(nn.Module):
 
         total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
         stage_block_id = 0
-        self.stages: List[nn.Module] = []
-        self.upsamples: List[nn.Module] = []
+        layers_list: List[nn.Module] = []
+
         for cnf in block_setting[1:]:
             if cnf.input_channels is not None:
                 # Bottlenecks
@@ -225,30 +226,45 @@ class ConvNeXtReverseBase2(nn.Module):
                     sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                     stage.append(block(cnf.input_channels, layer_scale, sd_prob))
                     stage_block_id += 1
-                self.stages.append(nn.Sequential(*stage).to(device))
+                layers_list.append(nn.Sequential(*stage).to(device))
             else:
-                self.stages.append(nn.Identity().to(device))
+                layers_list.append(nn.Identity().to(device))
 
             if cnf.out_channels is not None:
-                # Upnsampling
-                self.upsamples.append(
+                # 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).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_normals = nn.Sequential(*stage).to(device)
+
         self.last = nn.Conv2d(
                     block_setting[-1].input_channels,
                     1,
-                    kernel_size=7,
-                    padding=3,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+        self.last_normals = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3,
+                    kernel_size=3,
+                    padding='same',
                     ).to(device)
         # self.last = nn.Sequential(
         #                 Permute([0, 2, 3, 1]),
         #                 nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
         #                 Permute([0, 3, 1, 2]),
-        #             )
+        #             ).to(device)
 
 
         for m in self.modules():
@@ -259,17 +275,28 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        x, features = x
-        x = self.first(x)
+        if self.use_skip:
+            input, features = x
+        else:
+            input = x
+        x = self.first(input)
         for i in range(5 - 2):
-            x = self.stages[i](x + features[2-i])
+            if self.use_skip:
+                x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
+            else:
+                x = self.stages[i](x)
             x = self.upsamples[i](x)
         
         x = self.stages[-2](x)
-        x = self.upsamples[-1](x)
-        x = self.stages[-1](x)
-        x = self.last(x)
-        return x
+        f = self.upsamples[-1](x)
+
+        x = self.stages[-1](f)
+        depth = self.last(x)
+
+        x = self.stage_normals(f)
+        normals = self.last_normals(x)
+        
+        return depth, normals
 
     def forward(self, x) -> Tensor:
         return self._forward_impl(x)
@@ -282,10 +309,14 @@ class ConvNeXtReverseBase(nn.Module):
         stochastic_depth_prob: float = 0.0,
         layer_scale: float = 1e-6,
         block: Optional[Callable[..., nn.Module]] = None,
-        norm_layer: Optional[Callable[..., nn.Module]] = None
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip=False,
     ) -> None:
         super().__init__()
 
+        self.use_skip = use_skip
+
         if not block_setting:
             raise ValueError("The block_setting should not be empty")
         elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
@@ -297,53 +328,17 @@ class ConvNeXtReverseBase(nn.Module):
         if norm_layer is None:
             norm_layer = partial(LayerNorm2d, eps=1e-6)
 
-        layers: List[nn.Module] = []
-
-
-        # add bottleneck which will be executed at the end, to process after the final larger upsampling
-
-
-
-        # Stem
-        firstconv_output_channels = block_setting[0].input_channels
-
-        layers.append(
-            nn.Sequential(
-                nn.Conv2d(16, 1, kernel_size=7, padding=3, bias=True),
-
-            )
-        )
-
-        stage: List[nn.Module] = []
-        for _ in range(3):
-            stage.append(block(16, layer_scale, 0.0))
-        layers.append(nn.Sequential(*stage))
-        
-        layers.append(
-            # Convtr NormActivation(
-            #     firstconv_output_channels,
-            #     3,
-            #     kernel_size=4,
-            #     stride=4,
-            #     padding=0,
-            #     norm_layer=None,
-            #     activation_layer=None,
-            #     bias=True,
-            # )
-            nn.Sequential(
-                nn.ConvTranspose2d(
-                    firstconv_output_channels,
-                    16,
-                    kernel_size=4,
-                    stride=4,
-                    padding=0,
-                    ),
-            )
-        )
+        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
-        for cnf in block_setting:
+        layers_list: List[nn.Module] = []
+
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
                 # Bottlenecks
                 stage: List[nn.Module] = []
                 for _ in range(cnf.num_layers):
@@ -351,17 +346,26 @@ class ConvNeXtReverseBase(nn.Module):
                     sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
                     stage.append(block(cnf.input_channels, layer_scale, sd_prob))
                     stage_block_id += 1
-            layers.append(nn.Sequential(*stage))
+                layers_list.append(nn.Sequential(*stage).to(device))
+            else:
+                layers_list.append(nn.Identity().to(device))
+
             if cnf.out_channels is not None:
-                # Upsampling
-                layers.append(
+                # Upnsampling
+                layers_list.append(
                     nn.Sequential(
-                        norm_layer(cnf.out_channels),
-                        nn.ConvTranspose2d(cnf.out_channels, cnf.input_channels, kernel_size=2, stride=2),
-                    )
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    ).to(device)
                 )
 
-        self.features = nn.Sequential(*reversed(layers))
+        self.layers = nn.Sequential(*layers_list).to(device)
+        self.last = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    1,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
 
         for m in self.modules():
             if isinstance(m, (nn.Conv2d, nn.Linear)):
@@ -369,12 +373,20 @@ class ConvNeXtReverseBase(nn.Module):
                 if m.bias is not None:
                     nn.init.zeros_(m.bias)
 
-    def _forward_impl(self, x: Tensor) -> Tensor:
-        x = self.features(x)
-        # x = torch.tanh(x)
-        return x
 
-    def forward(self, x: Tensor) -> Tensor:
+    def _forward_impl(self, x) -> Tensor:
+        if self.use_skip:
+            raise NotImplementedError
+            input, features = x
+        else:
+            input = x
+        x = self.first(input)
+        x = self.layers(x)
+        out = self.last(x)
+        
+        return out
+
+    def forward(self, x) -> Tensor:
         return self._forward_impl(x)
 
 
@@ -383,13 +395,14 @@ def _convnext(
     stochastic_depth_prob: float,
     weights: Optional[WeightsEnum],
     reversed: bool,
+    input_features: int,
     **kwargs: Any,
 ) -> ConvNeXtBase:
 
     if reversed:
-        model = ConvNeXtReverseBase2(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob,**kwargs)
     else:
-        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, **kwargs)
+        model = ConvNeXtBase(block_setting, stochastic_depth_prob=stochastic_depth_prob, input_features=input_features, **kwargs)
 
     if weights is not None:
         model.load_state_dict(weights.get_state_dict())
@@ -398,10 +411,10 @@ def _convnext(
 
 
 class ConvNeXt(nn.Module):
-    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, **kwargs) -> None:
+    def __init__(self, block_setting, stochastic_depth_prob=0.0, reversed=False, input_features=3, **kwargs) -> None:
         super().__init__()
 
-        self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, **kwargs)
+        self.model = _convnext(block_setting, stochastic_depth_prob, None, reversed, input_features, **kwargs)
 
     def forward(self, x: Tensor) -> Tensor:
         x = self.model(x)

models/depth_estim.py

@@ -13,23 +13,69 @@ class ViewEncoder(nn.Module):
         super().__init__()
         self.features: int = features
 
+        self.flatten = nn.Flatten()
         self.lin1 = nn.LazyLinear(features).to(device)
         self.act1 = nn.GELU().to(device)
-        self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
+        self.lin2 = nn.Linear(features, features).to(device)
+        self.act2 = nn.Tanh().to(device)
     
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        x: torch.Tensor = self.lin1(input)
+        x: torch.Tensor = self.flatten(input)
+        x= self.lin1(x)
         x = self.act1(x)
         x = self.lin2(x)
-        x = x.view(x.shape[0],self.features,15,20)
+        x = self.act2(x)
+        x = torch.ones((x.shape[0], self.features, 8, 8), device=x.get_device()) * x[:,:,None,None]
 
         return x
 
-class DepthEstimation(nn.Module):
+
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+
+class DumbReconstructor(nn.Module):
     def __init__(self, view_features=128, device='cuda'):
         super().__init__()
 
+        # self.model = torch.nn.Sequential(
+        #     nn.ConvTranspose2d(768 + view_features, 128, 7, 4, 3),
+        #     nn.BatchNorm2d(128),
+        #     nn.ConvTranspose2d(128, 128, 7, 4, 3),
+        #     nn.Conv2d(128, 64, 7, 1, 'same'),
+        #     nn.BatchNorm2d(64),
+        #     nn.ConvTranspose2d(64, 32, 7, 4, 3),
+        #     nn.Conv2d(32, 1, 5, 1, 'same')
+        # ).to(device)
+
+        self.last = nn.Sequential(
+                        nn.Linear(in_features=1024, out_features=480*640, bias=True),
+                    ).to(device)
+
+    def forward(self, input) -> torch.Tensor:
+        x = input
+        x = self.last(x).view((1,1,480,640))
+        return x
+
+
+class DumbEncoder(nn.Module):
+    def __init__(self, view_features=128, device='cuda'):
+        super().__init__()
+
+        self.last = nn.Sequential(
+                        nn.Flatten(),
+                        nn.Linear(in_features=3*480*640, out_features=1024, bias=True),
+                    ).to(device)
+
+    def forward(self, input) -> torch.Tensor:
+        x = input
+        x = self.last(x)
+        return x, None
+
+
+class DepthEstimation(nn.Module):
+    def __init__(self, view_features=128, device='cuda', use_skip=False):
+        super().__init__()
+
         
         block_setting = [
             CNBlockConfig(96, 192, 3),
@@ -40,16 +86,18 @@ class DepthEstimation(nn.Module):
         
         block_setting_reverse = [
             CNBlockConfig(768 + view_features, 384, 3),
-            CNBlockConfig(384, 192, 9),
-            CNBlockConfig(192, 96, 3),
-            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(384 * ( 2 if use_skip else 1), 192, 9),
+            CNBlockConfig(192 * ( 2 if use_skip else 1), 96, 3),
+            CNBlockConfig(96 * ( 2 if use_skip else 1), 27, 3),
             CNBlockConfig(27, 9, 3),
             CNBlockConfig(9, None, 3),
         ]
 
         self.convnext = ConvNeXt(block_setting, device=device)
+        # self.convnext = DumbEncoder()
         self.flatten = nn.Flatten()
-        self.decoder = ConvNeXt(block_setting_reverse, reversed=True, device=device)
+        self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device)
+        # self.decoder = DumbReconstructor()
 
         self.view_encoder = ViewEncoder(view_features)
 
@@ -60,5 +108,7 @@ class DepthEstimation(nn.Module):
         view_features = self.view_encoder(view)
         # x = self.flatten(x)
         x = self.decoder([torch.concat((x, view_features), axis=1), features])
+        x = torch.tanh(x[0]), x[1]
+        # x = self.decoder(x)
         
         return x

models/pose_ae.py

@@ -0,0 +1,183 @@
+import os
+import sys
+
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+torch.set_float32_matmul_precision('high')
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from models.convnext import ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
+
+
+class ConvNeXtReversePose(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip=False,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        stage: List[nn.Module] = []
+        for _ in range(3):
+            stage.append(block(block_setting[0].input_channels, layer_scale, 0))
+        self.stage_smooth = nn.Sequential(*stage).to(device)
+
+        self.first = nn.Sequential(
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
+                        norm_layer(block_setting[0].out_channels),
+                    ).to(device)
+
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        layers_list: List[nn.Module] = []
+
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # Bottlenecks
+                stage: List[nn.Module] = []
+                for _ in range(cnf.num_layers):
+                    # adjust stochastic depth probability based on the depth of the stage block
+                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                    stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                    stage_block_id += 1
+                layers_list.append(nn.Sequential(*stage).to(device))
+            else:
+                layers_list.append(nn.Identity().to(device))
+
+            if cnf.out_channels is not None:
+                # Upsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    ).to(device)
+                )
+
+        self.layers = nn.Sequential(*layers_list[:-1]).to(device)
+        self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
+
+        stage: List[nn.Module] = []
+        for _ in range(cnf.num_layers):
+            stage.append(block(block_setting[-1].input_channels, layer_scale, 0))
+        self.stage_mask = nn.Sequential(*stage).to(device)
+
+        self.last = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+
+        self.last_mask = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    1,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        if self.use_skip:
+            raise NotImplementedError
+            input, features = x
+        else:
+            input = x
+        x = self.stage_smooth(input)
+        x = self.first(x)
+        f = self.layers(x)
+        x = self.stage_pose(f)
+        pose = torch.sigmoid(self.last(x))
+
+        x = self.stage_mask(f)
+        mask = torch.sigmoid(self.last_mask(x))
+        pose = pose * (mask > 0.5).to(torch.float32)
+
+        return pose, mask
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+class PoseAE(nn.Module):
+    def __init__(self, features=256, device='cuda'):
+        super().__init__()
+
+                
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, features, 9),
+            CNBlockConfig(features, None, 3),
+        ]
+
+        block_setting_reverse = [
+            CNBlockConfig(features, 384, 3),
+            CNBlockConfig(384, 192, 9),
+            CNBlockConfig(192, 96, 3),
+            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(27, 9, 3),
+            CNBlockConfig(9, None, 3),
+        ]
+
+        self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
+        self.decoder = ConvNeXtReversePose(block_setting_reverse, device=device)
+
+        self.flatten = nn.Flatten()
+
+
+    def forward(self, input) -> torch.Tensor:
+        mask, pose_in = input
+        x, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
+
+        pose_out, mask_out = self.decoder(x)
+
+        return pose_out, mask_out
+
+
+if __name__ == '__main__':
+    # test the model
+    model = PoseAE()
+    # model = torch.compile(model)
+    print(model)
+    print(sum(p.numel() for p in model.parameters() if p.requires_grad))
+    print(sum(p.numel() for p in model.parameters()))
+
+    x = torch.rand((4, 1, 512, 512)).to('cuda')
+    y = torch.rand((4, 3, 512, 512)).to('cuda')
+    output = model((x, y))
+
+
+
+    pass

models/pose_encoder.py

@@ -0,0 +1,260 @@
+import os
+import sys
+
+# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+torch.set_float32_matmul_precision('high')
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from models.convnext import ConvNeXt, CNBlockConfig, CNBlock, LayerNorm2d
+
+
+class ConvNeXtReverseMaskPose(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip=False,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        stage: List[nn.Module] = []
+        for _ in range(3):
+            stage.append(block(block_setting[0].input_channels, layer_scale, 0))
+        self.stage_smooth = nn.Sequential(*stage).to(device)
+
+        self.first = nn.Sequential(
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
+                        norm_layer(block_setting[0].out_channels),
+                    ).to(device)
+
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        layers_list: List[nn.Module] = []
+
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # Bottlenecks
+                stage: List[nn.Module] = []
+                for _ in range(cnf.num_layers):
+                    # adjust stochastic depth probability based on the depth of the stage block
+                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                    stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                    stage_block_id += 1
+                layers_list.append(nn.Sequential(*stage).to(device))
+            else:
+                layers_list.append(nn.Identity().to(device))
+
+            if cnf.out_channels is not None:
+                # Upsampling
+                layers_list.append(
+                    nn.Sequential(
+                        norm_layer(cnf.input_channels),
+                        nn.ConvTranspose2d(cnf.input_channels, cnf.out_channels, kernel_size=2, stride=2),
+                    ).to(device)
+                )
+
+        self.layers = nn.Sequential(*layers_list[:-1]).to(device)
+        self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
+
+        # for cnf in block_setting[-1:]:
+        #     # Bottlenecks
+        #     stage: List[nn.Module] = []
+        #     for _ in range(cnf.num_layers):
+        #         stage.append(block(cnf.input_channels, layer_scale, 0))
+        #     self.stage_mask = nn.Sequential(*stage).to(device)
+
+        self.last = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
+        # self.last_mask = nn.Conv2d(
+        #             block_setting[-1].input_channels,
+        #             1,
+        #             kernel_size=3,
+        #             padding='same',
+        #             ).to(device)
+
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Linear)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        if self.use_skip:
+            raise NotImplementedError
+            input, features = x
+        else:
+            input = x
+        x = self.stage_smooth(input)
+        x = self.first(x)
+        f = self.layers(x)
+        x = self.stage_pose(f)
+        pose = torch.sigmoid(self.last(x))
+
+        # x = self.stage_mask(f)
+        # mask = torch.sigmoid(self.last_mask(x))
+        
+        # return pose, mask
+        return pose, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device)
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+class PoseEncoder(nn.Module):
+    def __init__(self, features=128, device='cuda'):
+        super().__init__()
+
+                
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, features, 9),
+            CNBlockConfig(features, None, 3),
+        ]
+
+        block_setting_reverse = [
+            CNBlockConfig(features, 384, 3),
+            CNBlockConfig(384, 192, 9),
+            CNBlockConfig(192, 96, 3),
+            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(27, 9, 3),
+            CNBlockConfig(9, None, 3),
+        ]
+
+        rotation_block_setting = [
+            CNBlockConfig(64, 128, 3),
+            CNBlockConfig(128, 128, 3),
+            CNBlockConfig(128, 256, 9),
+            CNBlockConfig(256, None, 3),
+        ]
+
+        rotation_block_setting_reverse = [
+            CNBlockConfig(256, 128, 3),
+            CNBlockConfig(128, 64, 9),
+            CNBlockConfig(64, 32, 3),
+            CNBlockConfig(32, 16, 3),
+            CNBlockConfig(16, None, 3),
+        ]
+
+        self.encoder = ConvNeXt(block_setting, input_features=7, device=device)
+        self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
+        self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
+        self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)
+
+        self.flatten = nn.Flatten()
+
+
+    def forward(self, input) -> torch.Tensor:
+        mask, pose_in, pose_in_back, pose_out, mask_out = input
+        features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], dim=1))
+        x = features.view(features.size(0), features.size(1), -1)
+        y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
+        rotation = self.rotation_decoder(y)
+        rotation = torch.squeeze(rotation, dim=1)
+        rotation = torch.softmax(rotation, dim=1)
+
+        # x = torch.einsum('bij,bkj->bik', x, rotation)
+        x = torch.einsum('bij,bjk->bik', x, rotation)
+        x = x.view(x.size(0), x.size(1), 16, 16)
+        pose, mask_out_pred = self.decoder(x)
+        pose = pose * mask_out
+        # mask_out = torch.round(mask_out)
+        # mask_out = (mask_out > 0.5).to(torch.int64)
+        # pose = pose * (mask_out > 0.5).to(torch.float32)
+
+        return pose, mask_out_pred, rotation
+        # x = pose
+        # return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
+
+class DumbPoseEncoder(nn.Module):
+    def __init__(self, in_channels=7, out_channels=3, device='cuda'):
+        super(DumbPoseEncoder, self).__init__()
+
+        # Encoder
+        self.encoder_conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
+        self.encoder_conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.encoder_conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+        self.encoder_conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
+        self.encoder_pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Decoder
+        self.decoder_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
+        self.decoder_upsample_2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv2 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.decoder_upsample_3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.decoder_conv4 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        # Encoder
+        mask, pose_in, pose_out = x
+        x = torch.cat([mask, pose_in, pose_out], dim=1)
+        x1 = nn.functional.relu(self.encoder_conv1(x))
+        x = self.encoder_pool(x1)
+        x = nn.functional.relu(self.encoder_conv2(x))
+        x = nn.functional.relu(self.encoder_conv3(x))
+        x = nn.functional.relu(self.encoder_conv4(x))
+
+        # Decoder
+        x = self.decoder_upsample(x)
+        x = nn.functional.relu(self.decoder_conv1(x))
+        x = self.decoder_upsample_2(x)
+        x = nn.functional.relu(self.decoder_conv2(x))
+        x = self.decoder_upsample_3(x)
+        x = nn.functional.relu(self.decoder_conv3(x))
+        x = nn.functional.sigmoid(self.decoder_conv4(x))
+
+        return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
+
+
+if __name__ == '__main__':
+    # test the model
+    model = PoseEncoder()
+    # model = torch.compile(model)
+    print(model)
+    print(sum(p.numel() for p in model.parameters() if p.requires_grad))
+    print(sum(p.numel() for p in model.parameters()))
+
+    x = torch.rand((4, 1, 512, 512)).to('cuda')
+    y = torch.rand((4, 3, 512, 512)).to('cuda')
+    z = torch.rand((4, 3, 512, 512)).to('cuda')
+    output = model((x, y, z))
+
+    from torchviz import make_dot
+    graph = make_dot(output, params=dict(model.named_parameters()))
+    graph.render("model_graph", format="pdf")
+
+    pass

models/pose_manipulator.py

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

models/view_encoder.py

@@ -0,0 +1,108 @@
+import os
+import sys
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
+import torch
+import numpy as np
+from torch import nn, Tensor
+from torch.nn import functional as F
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+torch.set_float32_matmul_precision('high')
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+
+
+class ViewEncoder(nn.Module):
+    def __init__(self, bottleneck_size):
+        super(ViewEncoder, self).__init__()
+        self.bottleneck_size = bottleneck_size
+        
+        self.encoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(3*3, 64),
+                nn.BatchNorm1d(64),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(64, 128),
+                nn.BatchNorm1d(128),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(128, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(256, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(512, bottleneck_size),
+                nn.BatchNorm1d(bottleneck_size),)
+        ])
+
+        # Decoder
+        self.decoder = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(bottleneck_size, 512),
+                nn.BatchNorm1d(512),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(512*2, 256),
+                nn.BatchNorm1d(256),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(256*2, 128),
+                nn.BatchNorm1d(128),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(128*2, 64),
+                nn.BatchNorm1d(64),
+                nn.ReLU(),),
+            nn.Sequential(
+                nn.Linear(64*2, 3*3),
+                nn.Tanh(),),
+        ])
+    
+    def forward(self, x):
+        # Encode
+        f1 = self.encoder[0](x.view(-1, 3*3))
+        f2 = self.encoder[1](f1)
+        f3 = self.encoder[2](f2)
+        f4 = self.encoder[3](f3)
+        matrix = self.encoder[4](f4)
+
+        d1 = self.decoder[0](matrix)
+        d2 = self.decoder[1](torch.cat((d1, f4), dim=1))
+        d3 = self.decoder[2](torch.cat((d2, f3), dim=1))
+        d4 = self.decoder[3](torch.cat((d3, f2), dim=1))
+        x = self.decoder[4](torch.cat((d4, f1), dim=1))
+        
+        x = x * np.pi
+        return x.view(-1, 3, 3), matrix
+        
+        # # Decode
+        # x = self.decoder(matrix)
+        # x = x * np.pi
+        # return x.view(-1, 3, 3), matrix
+    
+
+if __name__ == '__main__':
+    # test the model
+    model = ViewEncoder(bottleneck_size=128).to('cuda')
+    # model = torch.compile(model)
+    print(model)
+    print(sum(p.numel() for p in model.parameters() if p.requires_grad))
+    print(sum(p.numel() for p in model.parameters()))
+
+    x = torch.rand((4, 3, 3)).to('cuda')
+    output = model(x)
+
+    # from torchviz import make_dot
+    # graph = make_dot(output, params=dict(model.named_parameters()))
+    # graph.render("model_graph", format="pdf")
+
+    pass

requirements.txt

@@ -1,4 +0,0 @@
-plotly
-scipy
-deodr
-pandas

scripts/pose_ae_train.py

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

scripts/pose_encoder_train.py

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

scripts/pose_manipulator_train.py

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

scripts/view_encoder_train.py

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

speed.py

@@ -0,0 +1,12 @@
+import numpy as np
+import time
+
+if __name__ == '__main__':
+
+    l = []
+    while True:
+        start = time.perf_counter()
+        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}')

speed_2.py

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

superSMPL/smpl.py

@@ -1,350 +0,0 @@
-import os
-import sys
-
-import pickle
-import scipy.io as sio
-from math import *
-
-import numpy as np
-import tensorflow as tf
-
-# Left-Right
-eval_joints = {
-        'All': [1, 2, 3, 4, 5, 7, 8, 9, 12, 15, 16, 17, 18, 19, 20, 21],
-        'Head': [12, 15], # Neck, Head
-        'Torso': [3, 6, 9], # Back joints
-        'Shoulder': [13, 14, 16, 17], # Inner, Outer 
-        'Elbow': [18, 19],
-        'Wrist': [20, 21],
-        'Hip': [1, 2],
-        'Knee': [4, 5],
-        'Foot': [7, 8, 10, 11]} # Ankle, Foot
-        
-class SMPL:
-
-    def __init__(self, path):
-        self.J_regressor = []
-        self.v_template = []
-        self.shapedirs = []
-        self.posedirs = []
-        self.weights = []
-        # Load SMPL models
-        # path = os.path.abspath(os.path.dirname(__file__))
-        with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
-            if sys.version_info[0]<3:
-                female = pickle.load(f)
-            else:
-                female = pickle.load(f, encoding="latin1")
-            self.J_regressor.append(tf.constant(np.array(female['J_regressor'].todense(), dtype=np.float32)))
-            self.v_template.append(tf.constant(female['v_template'], dtype=np.float32))
-            self.posedirs.append(tf.constant(female['posedirs'], dtype=np.float32))
-            self.shapedirs.append(tf.constant(female['shapedirs'], dtype=np.float32))
-            self.weights.append(tf.constant(female['weights'], dtype=np.float32))
-        with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
-            if sys.version_info[0]<3:
-                male = pickle.load(f)
-            else:
-                male = pickle.load(f, encoding="latin1")
-            self.J_regressor.append(tf.constant(np.array(male['J_regressor'].todense(), dtype=np.float32)))
-            self.v_template.append(tf.constant(male['v_template'], dtype=np.float32))
-            self.posedirs.append(tf.constant(male['posedirs'], dtype=np.float32))
-            self.shapedirs.append(tf.constant(male['shapedirs'], dtype=np.float32))
-            self.weights.append(tf.constant(male['weights'], dtype=np.float32))
-            kintree_table = male['kintree_table']
-            id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
-            self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
-        self.J_regressor = tf.stack(self.J_regressor, axis=0)
-        self.v_template = tf.stack(self.v_template, axis=0)
-        self.shapedirs = tf.stack(self.shapedirs, axis=0)
-        self.posedirs = tf.stack(self.posedirs, axis=0)
-        self.weights = tf.stack(self.weights, axis=0)
-        self.weights_one_hot = tf.reduce_mean(self.weights, axis=0)
-        self.weights_one_hot = tf.math.argmax(self.weights_one_hot, axis=-1)
-        self.weights_one_hot = tf.one_hot(self.weights_one_hot, 24)
-
-        # Topology no head/hands/feet
-        F = sio.loadmat(path + '/../Data/ssmpl_faces.mat')
-        self.f = np.array(F['faces_default'])
-        self.f_skirt = np.array(F['faces_skirt'])
-        self.nn = np.array(sio.loadmat(path + '/../Data/NN.mat'))
-        self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
-
-    def Q2R(self, w, x, y, z):
-        def diag(a, b):  # computes the diagonal entries,  1 - 2*a**2 - 2*b**2
-            return 1 - 2 * tf.pow(a, 2) - 2 * tf.pow(b, 2)
-
-        def tr_add(a, b, c, d):  # computes triangle entries with addition
-            return 2 * a * b + 2 * c * d
-
-        def tr_sub(a, b, c, d):  # computes triangle entries with subtraction
-            return 2 * a * b - 2 * c * d
-
-        m = [[diag(y, z), tr_sub(x, y, z, w), tr_add(x, z, y, w)],
-                [tr_add(x, y, z, w), diag(x, z), tr_sub(y, z, x, w)],
-                [tr_sub(x, z, y, w), tr_add(y, z, x, w), diag(x, y)]]
-            
-        R = tf.stack([tf.stack(m[i], axis=-1) for i in range(3)], axis=-2)
-        return R
-            
-    def rodrigues(self, r):
-        theta = tf.norm(r, axis=(1, 2), keepdims=True)
-        theta = tf.maximum(theta, np.finfo(np.float32).eps)
-        r_hat = r / theta
-        z_stick = tf.zeros(theta.get_shape().as_list()[0], dtype=tf.float32)
-        m = tf.stack(
-            (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
-            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), axis=1)
-        m = tf.reshape(m, (-1, 3, 3))
-        i_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros(
-            (theta.get_shape().as_list()[0], 3, 3), dtype=tf.float32)
-        # R = cos * i_cube + (1 - cos) * dot + tf.sin(theta) * m
-        R = i_cube + m * tf.sin(theta) + tf.matmul(m, m) * (1 - tf.cos(theta))
-        return R
-
-    def with_zeros(self, x):
-        return tf.concat((x, tf.constant([[0.0, 0.0, 0.0, 1.0]], dtype=tf.float32)), axis=0)
-
-    def pack(self, x):
-        return tf.concat((tf.zeros((x.get_shape().as_list()[0], 4, 3), dtype=tf.float32), x), axis=2)
-        
-    def run_rodrigues(self, pose):
-        batch_size = pose.shape[0]
-        R_out = self.rodrigues(tf.reshape(pose, (-1, 1, 3)))
-        R_out = tf.reshape(R_out, (batch_size, -1, 3, 3))
-        #R_out = []
-        #for i in range(batch_size):
-        #    R_out.append(self.rodrigues(tf.reshape(pose[i,:], (-1, 1, 3))))
-        #R_out = tf.stack(R_out, axis=0)
-        return R_out
-
-    def run(self, R, betas, remove=True):
-        batch_size = betas.shape[0]
-
-        #J_regressor = tf.gather(self.J_regressor, gender, axis=0)
-        #shapedirs = tf.gather(self.shapedirs, gender, axis=0)
-        #v_template = tf.gather(self.v_template, gender, axis=0)
-        
-        v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
-        v_shaped = tf.stack(v_shaped, axis=-1)
-
-        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=self.weights, parent=self.parent):
-            R = x[0]
-            v_shaped = x[1]
-            g = x[2]
-
-            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
-            R_cube_big = R
-            R_cube = R_cube_big[1:]
-            I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
-            lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
-            v_posed = v_shaped[:, :, g] + tf.tensordot(posedirs[g], lrotmin, axes=[[2], [0]])
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(weights[g], results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_posed, tf.ones((v_posed.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J), v_posed
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        J_out = tf.subtract(J_out, J0)
-        T_out = tf.subtract(T_out, J0)
-        v_posed = tf.subtract(v_posed, J0)
-        
-        if remove:
-            T_out = tf.gather(T_out, self.to_keep, axis=1)
-            v_posed = tf.gather(v_posed, self.to_keep, axis=1)
-        
-        return T_out, J_out, v_posed
-
-    def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None):
-        
-        v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
-        v_shaped = tf.stack(v_shaped, axis=-1)
-
-        if weights is None:
-            w = self.weights_one_hot[np.newaxis, :, :]
-            if filter is not None:
-                w = tf.gather(tf.gather(w, self.to_keep, axis=1), filter, axis=1)
-            if v_cloth.shape[1] == 14475:
-                w = tf.gather(w, self.to_keep, axis=1)
-        else:
-            w = weights
-
-        def single_run(x, J_regressor=self.J_regressor, w=w, parent=self.parent):
-            R = x[0]
-            v_shaped = x[1]
-            g = x[2]
-            v_cloth = x[3]
-
-            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
-            R_cube_big = R
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(w[0], results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J)
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out = tf.map_fn(single_run, (R, v_shaped, gender, v_cloth), (tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        T_out = tf.subtract(T_out, J0)
-
-        return T_out
-
-    def get_joints_template(self, betas):
-        v_shaped = []
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[0], axes=[[1], [2]]) + tf.reshape(self.v_template[0], shape=(1,27578,3))) # female
-        v_shaped.append(tf.tensordot(betas[:,:10], self.shapedirs[1], axes=[[1], [2]]) + tf.reshape(self.v_template[1], shape=(1,27578,3))) # male
-        v_shaped = tf.stack(v_shaped, axis=-1)
-
-        def single_run(x, J_regressor=self.J_regressor):
-            v_shaped = x[0]
-            g = x[1]
-            J = tf.matmul(J_regressor[g], v_shaped[:, :, g])
-            return J
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        J = tf.map_fn(single_run, (v_shaped, gender), (tf.float32))
-        J0 = tf.reshape(J[:,0,:], (-1, 1, 3))
-        
-        return tf.subtract(J, J0)
-
-
-    def run_garment(self, R, betas, v_cloth, J=None, fl_idx=None):
-        if J is None:
-            J = self.get_joints_template(betas)
-        
-        weights = tf.gather(self.weights, self.to_keep, axis=1)
-        
-        if fl_idx is not None:
-            weights = tf.gather(weights, fl_idx, axis=1)
-
-        def single_run(x, J_regressor=self.J_regressor, posedirs=self.posedirs, weights=weights, parent=self.parent):
-            R = x[0]
-            J = x[1]
-            g = x[2]
-            v_cloth = x[3]
-
-            R_cube_big = R
-            # R_cube = R_cube_big[1:]
-            # I_cube = tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((R_cube.get_shape()[0], 3, 3), dtype=tf.float32)
-            # lrotmin = tf.squeeze(tf.reshape((R_cube - I_cube), (-1, 1)))
-            
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(weights[g, :, :], results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.shape[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J)
-
-        gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out = tf.map_fn(single_run, (R, J, gender, v_cloth), (tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        J_out = tf.subtract(J_out, J0)
-        T_out = tf.subtract(T_out, J0)
-        
-        return T_out
-
-    def run_garment_with_weights(self, R, betas, v_cloth, weights, J=None):
-        
-        if J is None:
-            J = self.get_joints_template(betas)
-        
-        def single_run(x, parent=self.parent):
-            R = x[0]
-            J = x[1]
-            v_cloth = x[2]
-            weights = x[3]
-
-            R_cube_big = R
-       
-            results = []
-            results.append(self.with_zeros(tf.concat((R_cube_big[0], tf.reshape(J[0, :], (3, 1))), axis=1)))
-            for j in range(1, 24):
-                results.append(tf.matmul(results[parent[j]], self.with_zeros(tf.concat((R_cube_big[j], tf.reshape(J[j, :] - J[parent[j], :], (3, 1))), axis=1))))
-            stacked = tf.stack(results, axis=0)
-            results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
-            
-            J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
-            J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
-            
-            T = tf.tensordot(weights, results, axes=((1), (0)))
-            rest_shape_h = tf.concat((v_cloth, tf.ones((v_cloth.get_shape().as_list()[0], 1), dtype=tf.float32)), axis=1)
-            v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
-            v = tf.reshape(v, (-1, 4))[:, :3]
-
-            return v, tf.squeeze(J)
-
-        # gender = tf.cast(tf.greater_equal(betas[:, 10] , 0), tf.int32)
-        T_out, J_out = tf.map_fn(single_run, (R, J, v_cloth, weights), (tf.float32, tf.float32))
-        
-        J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3))
-        J_out = tf.subtract(J_out, J0)
-        T_out = tf.subtract(T_out, J0)
-        
-        return T_out
-
-    def run_backward(self, R, J, v_posed, weights):
-        batch_size = R.get_shape()[0]
-        
-        J_relative = J[:, 1:, :] - tf.gather(J, tf.constant(list(self.parent.values()), tf.int32), axis=1) # batch_size x 23 x 3
-
-        J_out = []
-        J_out.append(tf.expand_dims(J[:, 0, :], axis=-1)) # batch_size x 3 x 1
-        results = []
-        results.append(tf.expand_dims(tf.eye(3, dtype=tf.float32), axis=0) + tf.zeros((batch_size, 3, 3), dtype=tf.float32))
-        for j in range(1, 24):
-            results.append(tf.matmul(R[:, self.parent[j], :, :], results[self.parent[j]]))
-            J_out.append(tf.matmul(results[j], tf.expand_dims(J_relative[:, j-1, :], axis=-1)) + J_out[self.parent[j]]) # batch_size x 3 x 1
-        results = tf.transpose(tf.stack(results, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 3)
-        J_out = tf.transpose(tf.stack(J_out, axis=0), [1, 0, 2, 3]) # (batch_size x 24 x 3 x 1)
-
-        results = tf.matmul(R, results)
-        T = tf.transpose(tf.tensordot(results, weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 3)
-        v_posed = tf.expand_dims(v_posed, axis=-1) # (batch_size x 6890 x 3 x 1)
-        J = tf.expand_dims(J, axis=-1) # (batch_size x 24 x 3 x 1)
-        v = tf.matmul(T, v_posed) + tf.transpose(tf.tensordot(J_out - tf.matmul(results, J), weights, axes=((1), (1))), [0, 3, 1, 2]) # (batch_size x 6890 x 3 x 1)
-        v = tf.squeeze(v, -1) # (batch_size x 6890 x 3)
-        J_out = tf.squeeze(J_out, -1) # (batch_size x 24 x 3)
-
-        return v, J_out

superSMPL/smpl_torch.py

@@ -0,0 +1,229 @@
+# based on https://github.com/CalciferZh/SMPL/blob/master/smpl_torch_batch.py
+
+import os
+import sys
+import pickle
+import torch
+import numpy as np
+import scipy.io as sio
+from time import time
+from torch.nn import Module
+
+class SMPL(Module):
+  def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
+    
+    super(SMPL, self).__init__()
+    self.J_regressor = []
+    self.v_template = []
+    self.shapedirs = []
+    self.posedirs = []
+    self.weights = []
+    self.is_skirt = is_skirt
+    # Load SMPL models
+    # path = os.path.abspath(os.path.dirname(__file__))
+    with open(os.path.join(path, 'super_model_f.pkl'), 'rb') as f:
+        female = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(female['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(female['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(female['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(female['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(female['weights'], dtype=torch.float32))
+    with open(os.path.join(path, 'super_model_m.pkl'), 'rb') as f:
+        male = pickle.load(f, encoding="latin1")
+        self.J_regressor.append(torch.tensor(np.array(male['J_regressor'].todense(), dtype=np.float32)))
+        self.v_template.append(torch.tensor(male['v_template'], dtype=torch.float32))
+        self.posedirs.append(torch.tensor(male['posedirs'], dtype=torch.float32))
+        self.shapedirs.append(torch.tensor(male['shapedirs'], dtype=torch.float32))
+        self.weights.append(torch.tensor(male['weights'], dtype=torch.float32))
+        kintree_table = male['kintree_table']
+        id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
+        self.parent = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}
+    self.J_regressor = torch.stack(self.J_regressor, axis=0).to(device)
+    self.v_template = torch.stack(self.v_template, axis=0).to(device)
+    self.shapedirs = torch.stack(self.shapedirs, axis=0).to(device)
+    self.posedirs = torch.stack(self.posedirs, axis=0).to(device)
+    self.weights = torch.stack(self.weights, axis=0).to(device)
+
+    # Topology no head/hands/feet
+    F = sio.loadmat(os.path.join(path, 'ssmpl_faces.mat'))
+    self.f = np.array(F['faces_default'])
+    self.f_skirt = np.array(F['faces_skirt'])
+    self.nn = np.array(sio.loadmat(os.path.join(path, 'NN.mat')))
+    self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1)
+    self.f_all = np.asarray(female['f'])
+    # convert quad faces to triangles
+    self.f_all = np.concatenate([self.f_all[:,:3], self.f_all[:,[0,2,3]]], axis=0)
+    
+
+    self.weights_all = self.weights
+    self.weights = self.weights[:,self.to_keep,:]
+    
+    self.shape = (1,27578,3)
+    self.device = device
+
+    if orig_smpl:
+        self.to_keep = self.to_keep[self.to_keep < 6890]
+        self.J_regressor = self.J_regressor[:,:,:6890]
+        self.v_template = self.v_template[:,:6890,:]
+        self.shapedirs = self.shapedirs[:,:6890,:,:]
+        self.posedirs = self.posedirs[:,:6890,:,:]
+        self.weights_all = self.weights_all[:,:6890,:]
+        self.weights = self.weights_all[:,self.to_keep,:]
+        self.shape = (1, 6890, 3)
+
+  @staticmethod
+  def rodrigues(r):
+    """
+    Rodrigues' rotation formula that turns axis-angle tensor into rotation
+    matrix in a batch-ed manner.
+    Parameter:
+    ----------
+    r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].
+    Return:
+    -------
+    Rotation matrix of shape [batch_size * angle_num, 3, 3].
+    """
+    # eps = r.clone().normal_(std=1e-8)
+    # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
+    theta = torch.norm(r, dim=(1, 2), keepdim=True)
+    theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps)
+
+    theta_dim = theta.shape[0]
+    r_hat = r / theta
+    # cos = torch.cos(theta)
+    z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device)
+    m = torch.stack(
+      (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
+       -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
+    m = torch.reshape(m, (-1, 3, 3))
+    i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
+             + torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
+    A = r_hat.permute(0, 2, 1)
+    # dot = torch.matmul(A, r_hat)
+    # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
+    R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta))
+    return R
+
+  @staticmethod
+  def with_zeros(x):
+    """
+    Append a [0, 0, 0, 1] tensor to a [3, 4] tensor.
+    Parameter:
+    ---------
+    x: Tensor to be appended.
+    Return:
+    ------
+    Tensor after appending of shape [4,4]
+    """
+    ones = torch.tensor(
+      [[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
+    ).expand(x.shape[0],-1,-1).to(x.device)
+    ret = torch.cat((x, ones), dim=1)
+    return ret
+
+  @staticmethod
+  def pack(x):
+    """
+    Append zero tensors of shape [4, 3] to a batch of [4, 1] shape tensor.
+    Parameter:
+    ----------
+    x: A tensor of shape [batch_size, 4, 1]
+    Return:
+    ------
+    A tensor of shape [batch_size, 4, 4] after appending.
+    """
+    zeros43 = torch.zeros(
+      (x.shape[0], x.shape[1], 4, 3), dtype=torch.float32).to(x.device)
+    ret = torch.cat((zeros43, x), dim=3)
+    return ret
+
+
+  def forward(self, pose, betas, verts: torch.Tensor = None, remove: bool = True):
+    
+    """
+          Construct a compute graph that takes in parameters and outputs a tensor as
+          model vertices. Face indices are also returned as a numpy ndarray.
+          
+          20190128: Add batch support.
+          Parameters:
+          ---------
+          pose: Also known as 'theta', an [N, 24, 3] tensor indicating child joint rotation
+          relative to parent joint. For root joint it's global orientation.
+          Represented in a axis-angle format.
+          betas: Parameter for model shape. A tensor of shape [N, 10] as coefficients of
+          PCA components. Only 10 components were released by SMPL author.
+          trans: Global translation tensor of shape [N, 3].
+          Return:
+          ------
+          A 3-D tensor of [N * 6890 * 3] for vertices,
+          and the corresponding [N * 19 * 3] joint positions.
+    """
+    batch_num = betas.shape[0]
+
+    # v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template
+    v_shaped = []
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female
+    v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
+    v_shaped = torch.stack(v_shaped, axis=-1)
+    
+    # J = torch.matmul(self.J_regressor, v_shaped)
+    J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+    R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
+
+    if verts is not None:
+      v_posed = verts
+    else:
+      R_cube = R_cube_big[:, 1:, :, :]
+      I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \
+        torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device)
+      lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2)
+      v_posed = torch.stack([v_shaped[i, :, :, g.to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[g.to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
+
+    results = []
+    results.append(
+      self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
+    )
+    for i in range(1, 24):
+      results.append(
+        torch.matmul(
+          results[self.parent[i]],
+          self.with_zeros(
+            torch.cat(
+              (R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
+              dim=2
+            )
+          )
+        )
+      )
+    
+    stacked = torch.stack(results, dim=1)
+    results = stacked - \
+      self.pack(
+        torch.matmul(
+          stacked,
+          torch.reshape(
+            torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
+            (batch_num, 24, 4, 1)
+          )
+        )
+      )
+    
+    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+    
+    weights = self.weights if remove else self.weights_all
+    T = torch.stack([torch.tensordot(results[i], weights[g.to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    rest_shape_h = torch.cat(
+      (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
+    )
+    v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
+    v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
+    result = v
+
+    root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
+    result = result - root_joint
+    joints = joints - root_joint
+
+    if remove and verts is None:
+      result = result[:,self.to_keep,:]
+
+    return result, joints

test.py

@@ -1,5 +1,6 @@
 import os
 os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 
 import sys
 
@@ -17,16 +18,212 @@ from utils.utils import Accumulator
 from models.depth_estim import DepthEstimation
 from models.convnext import ConvNeXt
 
+from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
+ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
+
+
+def plot3d(depth, mesh, info, root_dist, scale, mask, root_joint_pixel):
+    import kornia
+    import plotly.graph_objects as go
+    from utils.data_utils import  get_camera
+    import torchvision as tv
+    cameras, K = get_camera(info['camLoc'][0], [480, 640], depth.get_device())
+    mask_orig = mask[:,None,:,:]
+    mask, mask_croped = restore(mask[:,None,:,:], scale, root_joint_pixel)
+    depth, _ = restore(depth, scale, root_joint_pixel, (mask_orig.to(torch.bool), mask_croped.to(torch.bool)), root_dist)# * mask
+
+    pc = depth_to_3d(depth, mask.to(torch.bool), root_joint_pixel, root_dist, K)
+    pc = pc.detach().cpu().numpy()
+
+
+    go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+                        ]).show()
+
+
+def restore(img, scale, center, masks=None, root_dist=0.0):
+    import torchvision as tv
+    import numpy as np
+    from utils.data_utils import resize_image
+    if masks is None:
+        resizer = tv.transforms.Resize((torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64)), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)
+    else:
+        resizer = lambda img, mask_croped, mask_resized: resize_image(img, mask_croped[:,0,:,:], mask_resized[:,0,:,:], (torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64)))
+    
+    cropper = lambda img, off: tv.transforms.functional.crop(img, off[1], off[0], 480, 640)
+
+    if masks is None:
+        out = resizer(img + root_dist)
+    else:
+        out = resizer(img + root_dist, masks[0], masks[1])
+
+    out_croped = out
+    offset = (np.array(np.round([out.shape[2]/2, out.shape[2]/2]))[None,:] - center.cpu().numpy()).astype(np.int64)
+    out = torch.stack([cropper(out[i], offset[i]) for i in range(img.shape[0])])
+    out = out
+
+    return out, out_croped
+    
+
+def depth_to_3d(depth, depth_mask, root_joint_pixel, root_dist, K):
+    import kornia
+    # pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
+    root_joint_image = torch.zeros_like(depth)
+    root_joint_image[:,:,root_joint_pixel[:,1], root_joint_pixel[:,0]] = root_dist
+    root_reprojected = kornia.geometry.depth.depth_to_3d(root_joint_image, K, normalize_points=True)[0,:,root_joint_pixel[:,1], root_joint_pixel[:,0]][:,0]
+    pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
+    pc = pc[0,:,depth_mask[0,0,:,:]]
+    # pc[2,:] -= root_dist
+    # root_reprojected[2] = root_dist
+    pc = pc - root_reprojected[:,None]
+    return pc
+
+
+def plot3d_with_mesh(depth, mesh, info, root_dist, scale, mask, root_joint_pixel, joints, original, original_mask, device='cuda'):
+    import plotly.graph_objects as go
+    import plotly.express as ex
+    from cloth3d.util import zRotMatrix
+    import pytorch3d
+    import numpy as np
+    from utils.data_utils import intrinsic, extrinsic, get_camera
+    import torchvision as tv
+    cameras, K = get_camera(info['camLoc'][0], [480, 640], depth.get_device())
+    # cameras, K = get_camera(info['camLoc'][0], [640, 480], depth.get_device())
+    # K = torch.tensor(intrinsic()[None, :, :], device=device)
+    # pc = cameras.unproject_points(depth)
+
+    # depth = original[None, :, :, :]
+    # mask = original_mask[None, :, :, :]
+
+    mask_orig = mask[:,None,:,:]
+    mask, mask_croped = restore(mask[:,None,:,:], scale, root_joint_pixel)
+    depth, _ = restore(depth, scale, root_joint_pixel, (mask_orig.to(torch.bool), mask_croped.to(torch.bool)), root_dist)# * mask
+
+    pc = depth_to_3d(depth, mask.to(torch.bool), root_joint_pixel, root_dist, K)
+    pc = pc.detach().cpu().numpy()
+
+    mesh = mesh[0] - joints[0,0,:]
+    # zRot = torch.tensor(zRotMatrix(-info['zrot']), device=device)
+    # mesh = torch.matmul(zRot, mesh.T).T
+    # mesh = mesh - info['trans'][0].to(torch.float32).to(device)
+    # zRot = torch.tensor(zRotMatrix(info['zrot']), device=device)
+    # mesh = torch.matmul(zRot, mesh.T).T
+
+    # view = pytorch3d.transforms.axis_angle_to_matrix(torch.tensor([0.0, 0.0, torch.pi/2], device=device))
+    # transform = pytorch3d.transforms.Rotate(view)
+    # mesh = transform.transform_points(mesh)
+
+    mesh = mesh.detach().cpu().numpy()
+
+    go.Figure(data = [go.Scatter3d(x=pc[2,:], y=pc[0,:], z=-pc[1,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+                      go.Scatter3d(x=-mesh[:,0], y=mesh[:,1], z=mesh[:,2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+                      ]).show()
+
+    pass
+
+# temp
+def plot_gt(depth, gt, mask):
+    import kornia
+    import plotly.graph_objects as go
+    from data.render import intrinsic
+    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+    pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
+    pc = pc[0,:,mask[0]].detach().cpu().numpy()
+    pc_gt = kornia.geometry.depth.depth_to_3d(gt[:,None,:,:], K, normalize_points=True)
+    pc_gt = pc_gt[0,:,mask[0]].detach().cpu().numpy()
+    go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+                        go.Scatter3d(x=pc_gt[0,:], y=pc_gt[1,:], z=pc_gt[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+                        ]).show()
+
+    pass
+
+
+def chamfer_loss(depth_image, vertices, depth_mask):
+    import kornia
+    import pytorch3d
+    from data.render import intrinsic
+    from pytorch3d.structures import Pointclouds
+    from utils.pytorch3d_chamfer import chamfer_distance
+    from pytorch3d.ops.points_normals import estimate_pointcloud_normals
+    from utils.data_utils import get_normals_from_depth
+    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+    pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
+    pc = [torch.permute(pc[i,:,depth_mask[i]], [1,0]) for i in range(depth_image.shape[0])]
+    pc_norm = get_normals_from_depth(depth_image, depth_mask, K)
+    pc_norm = [pc_norm[i,depth_mask[i],:] for i in range(pc_norm.shape[0])]
+    mesh_gt = [vertices[i, :, :] for i in range(vertices.shape[0])]
+    pc_pc = Pointclouds(pc, normals=pc_norm)
+    mesh_gt_pc = Pointclouds(mesh_gt)
+    mesh_norm = list(torch.unbind(estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)))
+    mesh_gt_pc = Pointclouds(mesh_gt, normals=mesh_norm)
+
+    chamfer = chamfer_distance(pc_pc, mesh_gt_pc)
+
+    return chamfer
+
+
+# def normal_loss(depth_image, depth_mask, normal_image_gt):
+#     import kornia
+#     import pytorch3d
+#     from data.render import intrinsic
+#     K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+#     pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
+#     pc = torch.stack([pc[i,:,depth_mask[i]] for i in range(depth_image.shape[0])])
+#     normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
+#     pass
+
+def depth_reconst_loss(y_pred, y_true, mask):
+    import kornia
+    from data.render import intrinsic
+    from pytorch3d.structures import Pointclouds
+    from pytorch3d.loss.chamfer import chamfer_distance as chamfer_pytorch
+    # from utils.pytorch3d_chamfer import chamfer_distance
+    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+    pred_pc = kornia.geometry.depth.depth_to_3d(y_pred, K, normalize_points=True)
+    pred_pc = [torch.permute(pred_pc[i, :, mask[i]], [1,0]) for i in range(y_pred.shape[0])]
+    pred_pc_pc = Pointclouds(pred_pc)
+    gt_pc = kornia.geometry.depth.depth_to_3d(y_true, K, normalize_points=True)
+    gt_pc = [torch.permute(gt_pc[i, :, mask[i]], [1,0]) for i in range(y_true.shape[0])]
+    gt_pc_pc = Pointclouds(gt_pc)
+    chamfer, _ = chamfer_pytorch(pred_pc_pc, gt_pc_pc)
+    return chamfer
+
+
+def normals_reconst_loss(y_pred, normals_gt, mask):
+    import kornia
+    from data.render import intrinsic
+    from utils.data_utils import  get_normals_from_depth
+    K = torch.tensor(intrinsic()[None, :, :], device='cuda')
+    normal_image = torch.permute(get_normals_from_depth(y_pred, mask, K), [0,3,1,2])
+
+    mse = loss_fn(normal_image * mask[:,None,:,:], normals_gt * mask[:,None,:,:])
+    return mse
+
 
 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(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
 
+    # input = torch.concat((batch['rgb_image'], batch['normal_image']), axis=1)
+    input = batch['rgb_image']
+    # import plotly.express as ex
+    # ex.imshow(batch['depth_image'][0,:,:].cpu().numpy(), zmin=-1, zmax=5).show()
+
     # 1. Forward pass
-    y_pred = model([batch['rgb_image'], batch['view']])
+    y_pred, y_pred_normals = model([input, batch['view']])
 
     # 2. Calculate loss
-    loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
+    mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
+    mse_normals = loss_fn(y_pred_normals * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:])
+    # chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
+    msssim = ms_ssim_module(y_pred * batch['depth_mask'][:,None,:,:] + 5, batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:] + 5)
+    chamfer_reconst = depth_reconst_loss(y_pred, batch['depth_image'][:,None,:,:], batch['depth_mask'])
+    normals_reconst = normals_reconst_loss(y_pred * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:], batch['depth_mask'])
+    # loss = mse #+ 1e-1 * (chamfer + normal)
+    loss = mse #+ mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst
+
+    # plot_gt(y_pred, batch)
 
     # 3. Zero grad optimizer
     optimizer.zero_grad()
@@ -45,10 +242,21 @@ def valid_step(batch):
     model.eval() # put the model in evaluation mode for testing (inference) 
     # 1. Forward pass
     with torch.inference_mode():
-        pred = model([batch['rgb_image'], batch['view']])
+        # input = torch.concat((batch['rgb_image'], batch['normal_image']), axis=1)
+        input = batch['rgb_image']
+
+        pred, pred_normals = model([input, batch['view']])
 
         # 2. Calculate the loss
-        loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
+        mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
+        mse_normals = loss_fn(pred_normals * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:])
+        # chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
+        msssim = ms_ssim_module(pred * batch['depth_mask'][:,None,:,:] + 5, batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:] + 5)
+        chamfer_reconst = depth_reconst_loss(pred, batch['depth_image'][:,None,:,:],  batch['depth_mask'])
+        normals_reconst = normals_reconst_loss(pred * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:], batch['depth_mask'])
+        # loss = mse #+ 1e-1 * (chamfer + normal)
+        loss = mse #+ mse_normals #+ chamfer_reconst + mse + normals_reconst
+
 
     return { 'valid/loss': loss.detach().cpu().numpy() }
 
@@ -63,17 +271,25 @@ def batch_to_device(batch, device):
 if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
-    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
-    train_dataloader = DataLoader(train_dataset, batch_size=2, drop_last=True, num_workers=0)
-    valid_dataloader = DataLoader(valid_dataset, batch_size=2, drop_last=True, num_workers=0)
+    random_view = False
+    use_skip = False
+    batch_size = 1
+    use_cache = True
+
+    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()
     valid_accumulator = Accumulator()
 
-    model = DepthEstimation(device=device)
+    model = DepthEstimation(device=device, use_skip=use_skip)
 
     loss_fn = nn.MSELoss()
+    loss_fn_3d = nn.MSELoss()
     optimizer = optim.Adam(model.parameters(), lr=0.001)
+    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9999, verbose=False)
+    # optimizer = optim.SGD(model.parameters(), lr=0.01)
 
     iteration = 0
     while True:
@@ -82,12 +298,15 @@ if __name__ == '__main__':
             train_loss = train_step(batch)
             train_accumulator.update(train_loss)
 
+            lr_scheduler.step()
+            iteration += 1
+
         if iteration % 1000 == 0:
             print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
             train_accumulator.reset()
 
-            iteration += 1
-
+        # TODO: change when batchsize != 1
+        if iteration % 100 == 0:
             for batch in valid_dataloader:
                 valid_loss = valid_step(batch)
                 valid_accumulator.update(valid_loss)
@@ -97,4 +316,5 @@ if __name__ == '__main__':
 
 
 
+
     pass 

test/perspective.py

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

utils/checkpoints.py

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

utils/data_utils.py

@@ -5,7 +5,8 @@ import torch
 import numpy as np
 import scipy.io as sio
 from math import cos, sin
-
+from pytorch3d.utils import cameras_from_opencv_projection
+import torchvision as tv
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 # from superSMPL.IO import readPC2Frame, readOBJ
@@ -46,15 +47,88 @@ def read_frames_file(file):
 
 
 def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
-    depth_dilated = kornia.morphology.dilation(depth_image[:,None,:,:], torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
+    depth_dilated = kornia.morphology.dilation(depth_image, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
     depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
     depth_dilated = torch.permute(depth_dilated, [0,2,3,1])
-    depth_dilated[depth_mask,:] = depth_image[:,:,:,None][depth_mask,:]
+    depth_image = torch.permute(depth_image, [0,2,3,1])
+    depth_dilated[depth_mask,:] = depth_image[depth_mask,:]
     depth_dilated = torch.permute(depth_dilated, [0,3,1,2])
     normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
     normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]
     return normals
 
 
+def resize_image(img, mask_croped, mask_resized, size):
+    resizer = tv.transforms.Resize(size, interpolation=tv.transforms.functional.InterpolationMode.BILINEAR)
+    dilated = kornia.morphology.dilation(img, torch.ones((3,3), device=img.get_device()), engine='convolution')
+    dilated = kornia.morphology.dilation(dilated, torch.ones((3,3), device=img.get_device()), engine='convolution')
+    dilated = torch.permute(dilated, [0,2,3,1])
+    img = torch.permute(img, [0,2,3,1])
+    dilated[mask_croped, :] = img[mask_croped, :]
+    dilated = torch.permute(dilated, [0,3,1,2])
+    out = resizer(dilated)
+    out = out * mask_resized[:,None,:,:]
+    return out
+
+
 def center_crop_resize(img, new_center, size):
     pass
+
+
+""" CAMERA """
+def intrinsic(RES_X=640, RES_Y=480):
+	f_mm             = 50 # blender default
+	sensor_w_mm      = 36 # blender default
+	sensor_h_mm = sensor_w_mm * RES_Y / RES_X
+
+	fx_px = f_mm * RES_X / sensor_w_mm
+	fy_px = f_mm * RES_Y / sensor_h_mm
+
+	u = RES_X / 2
+	v = RES_Y / 2
+
+	return np.array([[fx_px, 0,     u],
+					 [0,     fy_px, v],
+					 [0,     0,     1]], np.float32)
+				
+
+def extrinsic(camLoc):
+	R_w2bc = np.array([[0, 1, 0],
+					   [0, 0, 1],
+					   [1, 0, 0]], np.float32)
+	
+	T_w2bc = -1 * R_w2bc.dot(camLoc)
+	
+	R_bc2cv = np.array([[1,  0,  0],
+						[0, -1,  0],
+						[0,  0, -1]], np.float32)
+	
+	R_w2cv = R_bc2cv.dot(R_w2bc)
+	T_w2cv = R_bc2cv.dot(T_w2bc)
+
+	return np.concatenate((R_w2cv, T_w2cv[:,None]), axis=1)
+	
+
+def proj(camLoc):
+	return intrinsic().dot(extrinsic(camLoc))
+
+
+def get_camera(cam, img_size, device, res_x=640, res_y=480):
+    size = np.array(img_size, dtype=np.float32)
+    ext = torch.tensor(extrinsic(cam).astype(np.float32), device=device)
+    K = torch.tensor(intrinsic(res_x, res_y)[None, :, :], device=device)
+    cameras = cameras_from_opencv_projection(R=ext[None, :, :3], tvec=ext[None, :, 3], camera_matrix=K, image_size=torch.tensor(size[None, :], device=device))
+    return cameras, K
+
+
+def filter_info(info, frame_nr):
+    frame_specific = set(['poses', 'trans'])
+    general = set(['shape', 'camLoc', 'zrot', 'gender'])
+    out = dict()
+    for key in info.keys():
+        if key in frame_specific:
+            out[key] = info[key][:, frame_nr]
+        elif key in general:
+            out[key] = info[key]
+    
+    return out

utils/pose_coloring.py

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

utils/pytorch3d_chamfer.py

@@ -0,0 +1,141 @@
+# taken from pytorch3d.loss.chamfer, but only one directional
+
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+from pytorch3d.ops.knn import knn_gather, knn_points
+from pytorch3d.loss.chamfer import _handle_pointcloud_input, _validate_chamfer_reduction_inputs
+
+
+def chamfer_distance(
+    x,
+    y,
+    x_lengths=None,
+    y_lengths=None,
+    x_normals=None,
+    y_normals=None,
+    weights=None,
+    batch_reduction: Union[str, None] = "mean",
+    point_reduction: str = "mean",
+    norm: int = 2,
+):
+    """
+    Chamfer distance between two pointclouds x and y.
+
+    Args:
+        x: FloatTensor of shape (N, P1, D) or a Pointclouds object representing
+            a batch of point clouds with at most P1 points in each batch element,
+            batch size N and feature dimension D.
+        y: FloatTensor of shape (N, P2, D) or a Pointclouds object representing
+            a batch of point clouds with at most P2 points in each batch element,
+            batch size N and feature dimension D.
+        x_lengths: Optional LongTensor of shape (N,) giving the number of points in each
+            cloud in x.
+        y_lengths: Optional LongTensor of shape (N,) giving the number of points in each
+            cloud in y.
+        x_normals: Optional FloatTensor of shape (N, P1, D).
+        y_normals: Optional FloatTensor of shape (N, P2, D).
+        weights: Optional FloatTensor of shape (N,) giving weights for
+            batch elements for reduction operation.
+        batch_reduction: Reduction operation to apply for the loss across the
+            batch, can be one of ["mean", "sum"] or None.
+        point_reduction: Reduction operation to apply for the loss across the
+            points, can be one of ["mean", "sum"].
+        norm: int indicates the norm used for the distance. Supports 1 for L1 and 2 for L2.
+
+    Returns:
+        2-element tuple containing
+
+        - **loss**: Tensor giving the reduced distance between the pointclouds
+          in x and the pointclouds in y.
+        - **loss_normals**: Tensor giving the reduced cosine distance of normals
+          between pointclouds in x and pointclouds in y. Returns None if
+          x_normals and y_normals are None.
+    """
+    _validate_chamfer_reduction_inputs(batch_reduction, point_reduction)
+
+    if not ((norm == 1) or (norm == 2)):
+        raise ValueError("Support for 1 or 2 norm.")
+
+    x, x_lengths, x_normals = _handle_pointcloud_input(x, x_lengths, x_normals)
+    y, y_lengths, y_normals = _handle_pointcloud_input(y, y_lengths, y_normals)
+
+    return_normals = x_normals is not None and y_normals is not None
+
+    N, P1, D = x.shape
+    P2 = y.shape[1]
+
+    # Check if inputs are heterogeneous and create a lengths mask.
+    is_x_heterogeneous = (x_lengths != P1).any()
+    x_mask = (
+        torch.arange(P1, device=x.device)[None] >= x_lengths[:, None]
+    )  # shape [N, P1]
+
+    if y.shape[0] != N or y.shape[2] != D:
+        raise ValueError("y does not have the correct shape.")
+    if weights is not None:
+        if weights.size(0) != N:
+            raise ValueError("weights must be of shape (N,).")
+        if not (weights >= 0).all():
+            raise ValueError("weights cannot be negative.")
+        if weights.sum() == 0.0:
+            weights = weights.view(N, 1)
+            if batch_reduction in ["mean", "sum"]:
+                return (
+                    (x.sum((1, 2)) * weights).sum() * 0.0,
+                    (x.sum((1, 2)) * weights).sum() * 0.0,
+                )
+            return ((x.sum((1, 2)) * weights) * 0.0, (x.sum((1, 2)) * weights) * 0.0)
+
+    cham_norm_x = x.new_zeros(())
+
+    x_nn = knn_points(x, y, lengths1=x_lengths, lengths2=y_lengths, norm=norm, K=1)
+
+    cham_x = x_nn.dists[..., 0]  # (N, P1)
+
+    if is_x_heterogeneous:
+        cham_x[x_mask] = 0.0
+
+    if weights is not None:
+        cham_x *= weights.view(N, 1)
+
+    if return_normals:
+        # Gather the normals using the indices and keep only value for k=0
+        x_normals_near = knn_gather(y_normals, x_nn.idx, y_lengths)[..., 0, :]
+
+        cham_norm_x = 1 - torch.abs(
+            F.cosine_similarity(x_normals, x_normals_near, dim=2, eps=1e-6)
+        )
+
+        if is_x_heterogeneous:
+            cham_norm_x[x_mask] = 0.0
+
+        if weights is not None:
+            cham_norm_x *= weights.view(N, 1)
+
+    # Apply point reduction
+    cham_x = cham_x.sum(1)  # (N,)
+    if return_normals:
+        cham_norm_x = cham_norm_x.sum(1)  # (N,)
+    if point_reduction == "mean":
+        x_lengths_clamped = x_lengths.clamp(min=1)
+        cham_x /= x_lengths_clamped
+        if return_normals:
+            cham_norm_x /= x_lengths_clamped
+
+    if batch_reduction is not None:
+        # batch_reduction == "sum"
+        cham_x = cham_x.sum()
+        if return_normals:
+            cham_norm_x = cham_norm_x.sum()
+        if batch_reduction == "mean":
+            div = weights.sum() if weights is not None else max(N, 1)
+            cham_x /= div
+            if return_normals:
+                cham_norm_x /= div
+
+    cham_dist = cham_x 
+    cham_normals = cham_norm_x if return_normals else None
+
+    return cham_dist, cham_normals

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()