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

This reverts commit 05341fd48f.

.gitignore

@@ -161,4 +161,10 @@ cython_debug/
 #.idea/
 
 .vscode/
-temp/
+temp/
+dataset/
+wandb_login.py
+custom_data/dijsktra*
+checkpoints/
+
+*.out

cloth3d/read.py

@@ -0,0 +1,169 @@
+import os
+import sys
+import numpy as np
+from PIL import Image
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from cloth3d.smpl_np import SMPLModel
+from cloth3d.util import loadInfo, zRotMatrix, proj, mesh2UV, uv_to_pixel
+from superSMPL.IO import readOBJ, readPC2Frame
+
+class DataReader:
+
+    def __init__(self, folder, smpl_path):
+        # Data paths
+        self.SRC = folder
+        # SMPL model
+        smpl_path = smpl_path
+        self.smpl = {
+            'f': SMPLModel(os.path.join(smpl_path, 'super_model_f.pkl')),
+            'm': SMPLModel(os.path.join(smpl_path, 'super_model_m.pkl'))
+        }
+        
+    """ 
+    Read sample info 
+    Input:
+    - sample: name of the sample e.g.:'01_01_s0'
+    """
+    def read_info(self, sample):
+        info_path = os.path.join(self.SRC, sample, 'info')
+        return loadInfo(info_path)
+        
+    """ Human data """
+    """
+    Read SMPL parameters for the specified sample and frame
+    Inputs:
+    - sample: name of the sample
+    - frame: frame number
+    """
+    def read_smpl_params(self, sample, frame):
+        # Read sample data
+        info = self.read_info(sample)
+        # SMPL parameters
+        gender = 'm' if info['gender'] else 'f'
+        pose = info['poses'][:, frame].reshape(self.smpl[gender].pose_shape)
+        shape = info['shape']
+        trans = info['trans'][:, frame].reshape(self.smpl[gender].trans_shape)
+        return gender, pose, shape, trans
+    
+    """
+    Computes human mesh for the specified sample and frame
+    Inputs:
+    - sample: name of the sample
+    - frame: frame number
+    - absolute: True for absolute vertex locations, False for locations relative to SMPL root joint	
+    Outputs:
+    - V: human mesh vertices
+    - F: mesh faces
+    """
+    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 if use_trans else None)
+        V -= J[0:1]
+        # Apply rotation on z-axis
+        zRot = zRotMatrix(zrot)
+        return zRot.dot(V.T).T, self.smpl[gender].faces, zRot.dot(trans.T).T[None,:], zRot.dot((J -  J[0:1] + trans[None,:]).T).T[None,:], zRot.dot((J -  J[0:1]).T).T[None,:]
+    
+    """ Garment data """
+    """
+    Reads garment vertices location for the specified sample, garment and frame
+    Inputs:
+    - sample: name of the sample
+    - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+    - frame: frame number
+    - TODO: absolute: True for absolute vertex locations, False for locations relative to SMPL root joint	
+    Outputs:
+    - V: 3D vertex locations for the specified sample, garment and frame
+    """
+    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 + '.pc16')
+        V = readPC2Frame(pc16_path, frame, True)
+        # Read sample data
+        if trans is not None:
+            # Transform to absolute
+            V += trans[None]
+        zRot = zRotMatrix(zrot)
+        return zRot.dot(V.T).T
+
+    """
+    Reads garment faces for the specified sample and garment
+    Inputs:
+    - sample: name of the sample
+    - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+    Outputs:
+    - F: mesh faces	
+    """
+    def read_garment_topology(self, sample, garment):
+        # Read OBJ file
+        obj_path = os.path.join(self.SRC, sample, garment + '.obj')
+        return readOBJ(obj_path)[1]
+
+    """	
+    Reads garment UV map for the specified sample and garment
+    Inputs:
+    - sample: name of the sample
+    - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+    Outputs:
+    - Vt: UV map vertices
+    - Ft: UV map faces		
+    """
+    def read_garment_UVMap(self, sample, garment):
+        # Read OBJ file
+        obj_path = os.path.join(self.SRC, sample, garment + '.obj')
+        return readOBJ(obj_path)[2:]		
+
+    """
+    Reads vertex colors of the specified sample and garment
+    Inputs:
+    - sample: name of the sample
+    - garment: type of garment (e.g.: 'Tshirt', 'Jumpsuit', ...)
+    - F: mesh faces
+    - Vt: UV map vertices
+    - Ft: UV map faces
+    Output
+    - C: RGB colors
+    """
+    def read_garment_vertex_colors(self, sample, garment, F, Vt, Ft):
+        # Read sample texture info
+        texture = self.read_info(sample)['outfit'][garment]['texture']
+        # Plain color
+        if texture['type'] == 'color': return (255*texture['data']).astype(np.int32)
+        # Image texture
+        img_path = os.path.join(self.SRC, sample, garment + '.png')
+        img = Image.open(img_path)
+        # Get color of each UV vertex
+        color = np.array([img.getpixel(uv_to_pixel(vt)) for vt in Vt])
+        # Compute correspondece between V (mesh) and Vt (UV map)
+        m2uv = mesh2UV(F, Ft)
+        return np.array([color[list(m2uv[idx])].mean(0) for idx in sorted(list(m2uv.keys()))], np.uint8)
+        
+    """ Scene data """
+    """
+    Read camera location and compute projection matrix
+    Input:
+    - sample: name of the sample
+    Output:
+    - P: camera projection matrix (3 x 4)
+    """
+    def read_camera(self, sample):
+        # Read sample data
+        info = self.read_info(sample)
+        # Camera location
+        camLoc = info['camLoc']
+        # Camera projection matrix
+        return proj(camLoc)
+        
+# TESTING
+if __name__ == '__main__':
+    sample = '135_02_s8'
+    frame = 0
+    garment = 'Tshirt'
+    
+    reader = DataReader()
+    F = reader.read_garment_topology(sample, garment)
+    Vt, Ft = reader.read_garment_UVMap(sample, garment)
+    C = reader.read_garment_vertex_colors(sample, garment, F, Vt, Ft)

cloth3d/smpl_np.py

@@ -0,0 +1,189 @@
+import sys
+import numpy as np
+import pickle
+
+class SMPLModel():
+  def __init__(self, model_path):
+    """
+    SMPL model.
+
+    Parameter:
+    ---------
+    model_path: Path to the SMPL model parameters, pre-processed by
+    `preprocess.py`.
+
+    """
+    with open(model_path, 'rb') as f:
+      if sys.version_info[0] == 2: 
+          params = pickle.load(f) # Python 2.x
+      elif sys.version_info[0] == 3: 
+          params = pickle.load(f, encoding='latin1') # Python 3.x
+      self.J_regressor = params['J_regressor']
+      self.weights = params['weights']
+      self.posedirs = params['posedirs']
+      self.v_template = params['v_template']
+      self.shapedirs = params['shapedirs']
+      self.faces = params['f']
+      self.kintree_table = params['kintree_table']
+
+    id_to_col = {
+      self.kintree_table[1, i]: i for i in range(self.kintree_table.shape[1])
+    }
+    self.parent = {
+      i: id_to_col[self.kintree_table[0, i]]
+      for i in range(1, self.kintree_table.shape[1])
+    }
+
+    self.pose_shape = [24, 3]
+    self.beta_shape = [10]
+    self.trans_shape = [3]
+
+    self.pose = np.zeros(self.pose_shape)
+    self.beta = np.zeros(self.beta_shape)
+    self.trans = np.zeros(self.trans_shape)
+
+    self.verts = None
+    self.J = None
+    self.R = None
+
+    self.update()
+
+  def set_params(self, pose=None, beta=None, trans=None):
+    """
+    Set pose, shape, and/or translation parameters of SMPL model. Verices of the
+    model will be updated and returned.
+
+    Parameters:
+    ---------
+    pose: Also known as 'theta', a [24,3] matrix indicating child joint rotation
+    relative to parent joint. For root joint it's global orientation.
+    Represented in a axis-angle format.
+
+    beta: Parameter for model shape. A vector of shape [10]. Coefficients for
+    PCA component. Only 10 components were released by MPI.
+
+    trans: Global translation of shape [3].
+
+    Return:
+    ------
+    Updated vertices.
+
+    """
+    if pose is not None:
+      self.pose = pose
+    if beta is not None:
+      self.beta = beta
+    if trans is not None:
+      self.trans = trans
+    else:
+      self.trans = np.zeros(self.trans_shape)
+    self.update()
+    return self.verts, self.J
+
+  def update(self):
+    """
+    Called automatically when parameters are updated.
+
+    """
+    # how beta affect body shape
+    v_shaped = self.shapedirs.dot(self.beta) + self.v_template
+    # joints location
+    self.J = self.J_regressor.dot(v_shaped)
+    pose_cube = self.pose.reshape((-1, 1, 3))
+    # rotation matrix for each joint
+    self.R = self.rodrigues(pose_cube)
+    I_cube = np.broadcast_to(
+      np.expand_dims(np.eye(3), axis=0),
+      (self.R.shape[0]-1, 3, 3)
+    )
+    lrotmin = (self.R[1:] - I_cube).ravel()
+    # how pose affect body shape in zero pose
+    v_posed = v_shaped + self.posedirs.dot(lrotmin)
+    # world transformation of each joint
+    G = np.empty((self.kintree_table.shape[1], 4, 4))
+    G[0] = self.with_zeros(np.hstack((self.R[0], self.J[0, :].reshape([3, 1]))))
+    for i in range(1, self.kintree_table.shape[1]):
+      G[i] = G[self.parent[i]].dot(
+        self.with_zeros(
+          np.hstack(
+            [self.R[i],((self.J[i, :]-self.J[self.parent[i],:]).reshape([3,1]))]
+          )
+        )
+      )
+    G = G - self.pack(
+      np.matmul(
+        G,
+        np.hstack([self.J, np.zeros([24, 1])]).reshape([24, 4, 1])
+        )
+      )
+    # transformation of each vertex
+    T = np.tensordot(self.weights, G, axes=[[1], [0]])
+    rest_shape_h = np.hstack((v_posed, np.ones([v_posed.shape[0], 1])))
+    v = np.matmul(T, rest_shape_h.reshape([-1, 4, 1])).reshape([-1, 4])[:, :3]
+    self.verts = v + self.trans.reshape([1, 3])
+
+  def rodrigues(self, r):
+    """
+    Rodrigues' rotation formula that turns axis-angle vector into rotation
+    matrix in a batch-ed manner.
+
+    Parameter:
+    ----------
+    r: Axis-angle rotation vector of shape [batch_size, 1, 3].
+
+    Return:
+    -------
+    Rotation matrix of shape [batch_size, 3, 3].
+
+    """
+    theta = np.linalg.norm(r, axis=(1, 2), keepdims=True)
+    # avoid zero divide
+    theta = np.maximum(theta, np.finfo(np.float64).tiny)
+    r_hat = r / theta
+    cos = np.cos(theta)
+    z_stick = np.zeros(theta.shape[0])
+    m = np.dstack([
+      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]
+    ).reshape([-1, 3, 3])
+    i_cube = np.broadcast_to(
+      np.expand_dims(np.eye(3), axis=0),
+      [theta.shape[0], 3, 3]
+    )
+    A = np.transpose(r_hat, axes=[0, 2, 1])
+    B = r_hat
+    dot = np.matmul(A, B)
+    R = cos * i_cube + (1 - cos) * dot + np.sin(theta) * m
+    return R
+
+  def with_zeros(self, x):
+    """
+    Append a [0, 0, 0, 1] vector to a [3, 4] matrix.
+
+    Parameter:
+    ---------
+    x: Matrix to be appended.
+
+    Return:
+    ------
+    Matrix after appending of shape [4,4]
+
+    """
+    return np.vstack((x, np.array([[0.0, 0.0, 0.0, 1.0]])))
+
+  def pack(self, x):
+    """
+    Append zero matrices of shape [4, 3] to vectors of [4, 1] shape in a batched
+    manner.
+
+    Parameter:
+    ----------
+    x: Matrices to be appended of shape [batch_size, 4, 1]
+
+    Return:
+    ------
+    Matrix of shape [batch_size, 4, 4] after appending.
+
+    """
+    return np.dstack((np.zeros((x.shape[0], 4, 3)), x))

cloth3d/util.py

@@ -0,0 +1,122 @@
+import numpy as np
+import scipy.io as sio
+from math import cos, sin
+
+def loadInfo(filename):
+	'''
+	this function should be called instead of direct sio.loadmat
+	as it cures the problem of not properly recovering python dictionaries
+	from mat files. It calls the function check keys to cure all entries
+	which are still mat-objects
+	'''
+	data = sio.loadmat(filename, struct_as_record=False, squeeze_me=True)
+	del data['__globals__']
+	del data['__header__']
+	del data['__version__']
+	return _check_keys(data)
+
+def _check_keys(dict):
+	'''
+	checks if entries in dictionary are mat-objects. If yes
+	todict is called to change them to nested dictionaries
+	'''
+	for key in dict:
+		if isinstance(dict[key], sio.matlab.mio5_params.mat_struct):
+			dict[key] = _todict(dict[key])
+	return dict        
+
+def _todict(matobj):
+	'''
+	A recursive function which constructs from matobjects nested dictionaries
+	'''
+	dict = {}
+	for strg in matobj._fieldnames:
+		elem = matobj.__dict__[strg]
+		if isinstance(elem, sio.matlab.mio5_params.mat_struct):
+			dict[strg] = _todict(elem)
+		elif isinstance(elem, np.ndarray) and np.any([isinstance(item, sio.matlab.mio5_params.mat_struct) for item in elem]):
+			dict[strg] = [None] * len(elem)
+			for i,item in enumerate(elem):
+				if isinstance(item, sio.matlab.mio5_params.mat_struct):
+					dict[strg][i] = _todict(item)
+				else:
+					dict[strg][i] = item
+		else:
+			dict[strg] = elem
+	return dict
+
+# Computes matrix of rotation around z-axis for 'zrot' radians
+def zRotMatrix(zrot):
+	c, s = cos(zrot), sin(zrot)
+	return np.array([[c, -s, 0],
+					 [s,  c, 0],
+					 [0,  0, 1]], np.float32)
+
+
+""" 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))
+	
+""" 
+Mesh to UV map
+Computes correspondences between 3D mesh and UV map
+NOTE: 3D mesh vertices can have multiple correspondences with UV vertices
+"""
+def mesh2UV(F, Ft):
+	m2uv = {v: set() for f in F for v in f}
+	for f, ft in zip(F, Ft):
+		for v, vt in zip(f, ft):
+			m2uv[v].add(vt)
+	# m2uv = {k:list(v) for k,v in m2uv.items()}
+	return m2uv
+	
+# Maps UV coordinates to texture space (pixel)
+IMG_SIZE = 2048 # all image textures have this squared size
+def uv_to_pixel(vt):
+	px = vt * IMG_SIZE # scale to image plane
+	px %= IMG_SIZE # wrap to [0, IMG_SIZE]
+	# Note that Blender graphic engines invert vertical axis
+	return int(px[0]), int(IMG_SIZE - px[1]) # texel X, texel Y
+
+
+def quads2tris(F):
+    out = []
+    for f in F:
+        if len(f) == 3: out += [f]
+        elif len(f) == 4: out += [[f[0],f[1],f[2]],
+                                [f[0],f[2],f[3]]]
+        else: print("This should not happen...")
+    return np.array(out, np.int32)

data/dataset.py

@@ -1,11 +1,13 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '1'
+# 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
@@ -13,55 +15,314 @@ import plotly.express as ex
 
 from torch.utils.data import Dataset
 
-from render import Render
+
+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
 
 
 class Dataset(Dataset):
-    def __init__(self, folder):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, random_view=True):
         self.folder = folder
-        self.transforms = tv.transforms.ToTensor()
+        self.folder_rgb = folder_rgb
-        self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))
+        self.device = device
+        self.random_view = random_view
 
-        self.render = Render(max_depth=10)
+        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.body_faces = np.load('custom_data/smpl_faces.npy')
+        self.reader = DataReader(folder, 'superSMPL')
         
 
+        self.cache = dict() if use_cache else None
+
     def __len__(self):
         return len(self.frames)
 
 
+    # TODO: make it work with two part outfits 
     def __getitem__(self, idx):
-        path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
+        frame : Frame = self.frames[idx]
-        mesh_posed = np.load(path)['mesh_posed']
+        if self.cache is not None:
+            if frame in self.cache.keys():
+                return self.cache[frame]
+
+        path = os.path.join(self.folder, frame.id)
+        info = self.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 = 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))
+        # 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, 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)
+
+        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)
 
-        id = self.frames[idx].split('_')[0]
-        faces = np.load(os.path.join(self.folder, 'faces', id + '.npy'))
-        mask = np.load(os.path.join(self.folder, 'mesh_mask', id + '.npy'))
-        body = np.load(os.path.join(self.folder, 'mesh_body_posed', self.frames[idx][:-1] + 'y'))
 
-        info_folder = '/home/workspace/data/dataset/' + ('train' if self.folder.split('/')[-1] == 'train' else 'validation') + '/' + id + '/info.mat'
-        info = sio.loadmat(info_folder)
-        combined = np.concatenate((body, mesh_posed[mask, :]))
-        combined_faces = np.concatenate((self.body_faces, faces + 14475))
 
         start = time.perf_counter()
-        depth_image = self.render.render(combined, combined_faces, info['camLoc'][0])[:,:,0]
+        # depth_image = self.render.render(mesh_posed[mask, :], faces, info['camLoc'][0])[:,:,0]
-        t = time.perf_counter() - start
+        # 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, 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
+        # 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,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]
+        # joints_no_trans_pixel = camera_no_trans.transform_points_screen(torch.tensor(joints_no_trans.astype(np.float32), device=combined.get_device()))[0,:,:2]
+        # ratio = ((joints_pixel[1,:] - joints_pixel[2,:]) ** 2).sum().sqrt() / ((joints_no_trans_pixel[1,:] - joints_no_trans_pixel[2,:]) ** 2).sum().sqrt()
+
+        # center_offset = torch.tensor([240, 320], device=root_joint_pixel.get_device()) - root_joint_pixel
+        # center_offset_ratio = center_offset / torch.tensor([480, 640], device=root_joint_pixel.get_device())
+        # bbox = tv.ops.masks_to_boxes(depth_mask[None,:,:])[0,:]
+        # width = bbox[2] - bbox[0]
+        # height = bbox[3] - bbox[1]
+        # width = width + (width % 2)
+        # height = height + (height % 2)
+        # dims = torch.tensor([width, height], device=width.get_device())
+        # # crop_size = torch.max(dims)
+        # # top = (root_joint_pixel[1] - height / 2).round().to(torch.int64)
+        # # left = (root_joint_pixel[0] - width / 2).round().to(torch.int64)
+        
+        # cropped = tv.transforms.functional.crop(depth_image, bbox[1].to(torch.int64), bbox[0].to(torch.int64), height.to(torch.int64), width.to(torch.int64))
+        # new_root_joint_pixel = root_joint_pixel - bbox[:2]
+        # offset = dims / 2 - new_root_joint_pixel
+        # 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))
+
+
+        # left = torch.tensor([0], device=bbox.get_device())
+        # right = torch.tensor([0], device=bbox.get_device())
+        # top = torch.tensor([0], device=bbox.get_device())
+        # bottom = torch.tensor([0], device=bbox.get_device())
+        # if offset[0] > 0:
+        #     left = 2 * offset[0]
+        # else:
+        #     right = -1 * 2 * offset[0]
+        # if offset[1] > 0:
+        #     top = 2 * offset[1]
+        # else:
+        #     bottom = -1 * 2 * offset[1]
+        # # centered = tv.transforms.functional.crop(cropped, new_top.to(torch.int64), new_left.to(torch.int64), new_height.to(torch.int64), new_width.to(torch.int64))
+        # centered = torch.nn.functional.pad(cropped, (left.to(torch.int64), right.to(torch.int64), top.to(torch.int64), bottom.to(torch.int64)), value=-1)
+        # orig = depth_image[root_joint_pixel.to(torch.int64)[1], root_joint_pixel.to(torch.int64)[0]].cpu().numpy()
+        # new = centered[(centered.shape[0] // 2), (centered.shape[1] // 2)].cpu().numpy()
+        # assert new == orig, 'Center pad did not work'
+
+        # scaled = tv.transforms.Resize((centered.shape[1] / ratio).round().int().item(), interpolation=tv.transforms.functional.InterpolationMode.NEAREST)(centered[None,None,:,:])[0,0,:,:]
+        # left, right = math.floor((depth_image.shape[1] - scaled.shape[1]) / 2), round((depth_image.shape[1] - scaled.shape[1]) / 2)
+        # top, bottom = math.floor((depth_image.shape[0] - scaled.shape[0]) / 2), round((depth_image.shape[0] - scaled.shape[0]) / 2)
+        # upscaled = torch.nn.functional.pad(scaled, (left, right, top, bottom), value=-1)
+
+
+        # 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]
+
+        # 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()
+        # dist_img = depth_image[v2i[idx, 1], v2i[idx, 0]]
+        
+        # depth_image[v2i[visible_vertices, 1], v2i[visible_vertices, 0]] = 10
+
+        # diff = np.abs(dist - dist_img)
+
+        # import plotly.express as ex
+        # import plotly.graph_objects as go
+        # ex.imshow(deodr_depth[:,:,0]).show()
+        # ex.imshow(depth_image.cpu().numpy()).show()
+        # ex.imshow(normal_image[0].cpu().numpy()).show()
+
+        # trans = np.array(info['trans'], dtype=np.float32)[:, self.frames[idx].frame_nr]
+        # root = info['camLoc'].copy()
+        # root[1:] = 0
+        # go.Figure(data = [go.Scatter3d(x=combined[:, 0], y=combined[:, 1], z=combined[:, 2], marker=dict(size=3, line=dict(width=0)), mode='markers'),
+        #                   go.Scatter3d(x=[combined[idx, 0]], y=[combined[idx, 1]], z=[combined[idx, 2]], marker=dict(size=10, line=dict(width=0)), mode='markers'),
+        #                   go.Scatter3d(x=[root[0]], y=[root[1]], z=[root[2]], marker=dict(size=10, line=dict(width=0)), mode='markers'),
+        #                  ]).show()
+
+        # 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()
+
+        # v2i_coverage = np.zeros((480, 640), dtype=np.float32)
+        # 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(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([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)
-        return { 'mesh_posed': mesh_posed, 'depth_image': depth_image, 'mask': mask }
+        data = { 'mesh_posed': mesh_posed,
+                 '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,
+                 '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 data
+    
 
 if __name__ == '__main__':
     from torch.utils.data import DataLoader
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    dataset = Dataset('/home/workspace/local-data/train')
+    # dataset = Dataset('/home/workspace/local-data/train')
-    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True)
+    dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005.txt', random_view=True, device=device)
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
 
     for batch in train_dataloader:
         batch
         pass
 
-    pass 
+    pass 

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

@@ -2,12 +2,17 @@ import os
 import sys
 
 import numpy as np
+import scipy as sp
+import scipy.io as sio
 from tqdm import tqdm
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from utils.utils import mask_faces
+from utils.data_utils import mask_faces, Frame
+from superSMPL.IO import readPC2, writePC2
+from cloth3d.util import loadInfo
 
-def generate_faces(folder):
+
+def generate_faces(folder, frames):
     frames = os.listdir(os.path.join(folder, 'mesh_mask'))
     F = np.load('custom_data/smpl_faces.npy')
     os.makedirs(os.path.join(folder, 'faces'), exist_ok=True)
@@ -18,8 +23,57 @@ def generate_faces(folder):
         np.save(os.path.join(folder, 'faces', frame), faces)
 
 
+def generate_frames_file(folder, file_name):
+    ids = os.listdir(folder)#[:1]
+    frames = []
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+        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):
+            f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')
+
+
+def genereated_registered_and_mask(folder, is_skirt = False):
+    ids = os.listdir(folder)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+        garments = [f.replace('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')]
+        
+        assert len(garments) == 1
+
+        if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')):
+            continue
+        
+        garment = garments[0]
+
+        verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+        map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1))
+        verts = verts[:, map, :] # ---------------> This is the garment after registration
+        raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1))
+        mask = raw_mask > 0
+        # boundaries = raw_mask == 2
+
+        writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True)
+        np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask)
+
+
 if __name__ == '__main__':
     # folder = '/home/workspace/local-data/train'
-    folder = '/home/workspace/local-data/valid'
+    # folder = '/home/workspace/local-data/valid'
+    folder = '/home/workspace/data/dataset/cloth3d/train'
 
-    generate_faces(folder)
+    # generate_faces(folder)
+    # genereated_registered_and_mask(folder)
+    generate_frames_file(folder, 'frames_all.txt')

data/render.py

@@ -1,15 +1,22 @@
+import os
+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
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from deodr.triangulated_mesh import TriMesh
+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):
         self.mesh = None
         self.depthScale = depth_scale
-        self.scene = Scene3D(0)
+        self.scene = Scene3D()
         self.set_max_depth(max_depth)
 		
     def set_mesh(self, vertices, faces):
@@ -63,44 +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)
@@ -121,4 +90,106 @@ class Render:
         depth[~mask] = 0.0
         mask = mask.astype(np.float32)
 
-        return depth
+        return depth
+
+
+from torch import nn
+from pytorch3d.renderer import (
+    look_at_view_transform,
+    FoVPerspectiveCameras, 
+    PerspectiveCameras,
+    PointLights, 
+    DirectionalLights, 
+    Materials, 
+    RasterizationSettings, 
+    MeshRenderer, 
+    MeshRasterizer,  
+    SoftPhongShader,
+    TexturesUV,
+    TexturesVertex,
+    HardPhongShader,
+)
+from pytorch3d.structures import Meshes
+from pytorch3d.utils import cameras_from_opencv_projection
+
+
+class DepthRenderer(nn.Module):
+    def __init__(self) -> None:
+         super().__init__()
+
+
+
+    def forward(self, verts, faces, cam):
+        # 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 = 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')
+        # 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, K = get_camera(cam, [480, 640], verts.get_device())
+
+        raster_settings = RasterizationSettings(
+                            image_size=(480, 640), 
+                            blur_radius=0.0, 
+                            faces_per_pixel=1, 
+                            bin_size=0
+                        )
+
+        rasterizer = MeshRasterizer(
+                    cameras=cameras, 
+                    raster_settings=raster_settings
+                )
+        output = rasterizer(Meshes(verts[None,:,:], faces[None,:,:]))
+        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()
+        # dmax = depth[mask].max()
+        # depth[mask] = depth[mask] - dmean
+        # depth[~mask] = 0.0
+        # mask = mask.astype(np.float32)
+
+        # import plotly.express as ex
+        # 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

@@ -0,0 +1,429 @@
+# taken from torchvision.models.convnext
+
+from functools import partial
+from typing import Any, Callable, List, Optional, Sequence
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+from torchvision.ops.misc import Conv2dNormActivation, Permute
+from torchvision.ops.stochastic_depth import StochasticDepth
+from torchvision.models._api import Weights, WeightsEnum
+
+
+
+__all__ = [
+    "ConvNeXt",
+    "ConvNeXt_Tiny_Weights",
+    "ConvNeXt_Small_Weights",
+    "ConvNeXt_Base_Weights",
+    "ConvNeXt_Large_Weights",
+    "convnext_tiny",
+    "convnext_small",
+    "convnext_base",
+    "convnext_large",
+]
+
+
+class LayerNorm2d(nn.LayerNorm):
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.permute(0, 2, 3, 1)
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        x = x.permute(0, 3, 1, 2)
+        return x
+
+
+class CNBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        layer_scale: float,
+        stochastic_depth_prob: float,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+    ) -> None:
+        super().__init__()
+        if norm_layer is None:
+            norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+        self.block = nn.Sequential(
+            nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim, bias=True),
+            Permute([0, 2, 3, 1]),
+            norm_layer(dim),
+            nn.Linear(in_features=dim, out_features=4 * dim, bias=True),
+            nn.GELU(),
+            nn.Linear(in_features=4 * dim, out_features=dim, bias=True),
+            Permute([0, 3, 1, 2]),
+        )
+        self.layer_scale = nn.Parameter(torch.ones(dim, 1, 1) * layer_scale)
+        self.stochastic_depth = StochasticDepth(stochastic_depth_prob, "row")
+
+    def forward(self, input: Tensor) -> Tensor:
+        result = self.layer_scale * self.block(input)
+        result = self.stochastic_depth(result)
+        result += input
+        return result
+
+
+class CNBlockConfig:
+    # Stores information listed at Section 3 of the ConvNeXt paper
+    def __init__(
+        self,
+        input_channels: int,
+        out_channels: Optional[int],
+        num_layers: int,
+    ) -> None:
+        self.input_channels = input_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+
+    def __repr__(self) -> str:
+        s = self.__class__.__name__ + "("
+        s += "input_channels={input_channels}"
+        s += ", out_channels={out_channels}"
+        s += ", num_layers={num_layers}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+
+class ConvNeXtBase(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,
+        input_features: int = 3,
+        device: str = 'cuda',
+    ) -> None:
+        super().__init__()
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        self.stages: List[nn.Module] = []
+        self.downsamples: List[nn.Module] = []
+
+        # Stem
+        firstconv_output_channels = block_setting[0].input_channels
+        self.stem = Conv2dNormActivation(
+                input_features,
+                firstconv_output_channels,
+                kernel_size=4,
+                stride=4,
+                padding=0,
+                norm_layer=norm_layer,
+                activation_layer=None,
+                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
+        for cnf in block_setting:
+            # 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))
+            if cnf.out_channels is not None:
+                # Downsampling
+                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():
+            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) -> Tensor:
+        x = self.stem(x)
+        x = self.layers(x)
+        x = self.act(x)
+
+        return x, None
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x)
+
+
+
+class ConvNeXtReverseDepthNormals(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)
+
+        # Stem
+        # firstconv_output_channels = block_setting[0].input_channels
+        # self.first = nn.Sequential(
+        #             nn.ConvTranspose2d(
+        #                 firstconv_output_channels,
+        #                 firstconv_output_channels,
+        #                 kernel_size=2,
+        #                 stride=2,
+        #                 padding=0,
+        #                 ),
+        #             norm_layer(firstconv_output_channels),
+        #         ).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).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=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():
+            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:
+            input, features = x
+        else:
+            input = x
+        x = self.first(input)
+        for i in range(5 - 2):
+            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)
+        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)
+
+
+class ConvNeXtReverseBase(nn.Module):
+    def __init__(
+        self,
+        block_setting: List[CNBlockConfig],
+        stochastic_depth_prob: float = 0.0,
+        layer_scale: float = 1e-6,
+        block: Optional[Callable[..., nn.Module]] = None,
+        norm_layer: Optional[Callable[..., nn.Module]] = None,
+        device: str = 'cuda',
+        use_skip=False,
+    ) -> None:
+        super().__init__()
+
+        self.use_skip = use_skip
+
+        if not block_setting:
+            raise ValueError("The block_setting should not be empty")
+        elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
+            raise TypeError("The block_setting should be List[CNBlockConfig]")
+
+        if block is None:
+            block = CNBlock
+
+        if norm_layer is None:
+            norm_layer = partial(LayerNorm2d, eps=1e-6)
+
+        self.first = nn.Sequential(
+                        nn.ConvTranspose2d(block_setting[0].input_channels, block_setting[0].out_channels, kernel_size=2, stride=2),
+                        norm_layer(block_setting[0].out_channels),
+                    ).to(device)
+
+        total_stage_blocks = sum(cnf.num_layers for cnf in block_setting)
+        stage_block_id = 0
+        layers_list: List[nn.Module] = []
+
+        for cnf in block_setting[1:]:
+            if cnf.input_channels is not None:
+                # Bottlenecks
+                stage: List[nn.Module] = []
+                for _ in range(cnf.num_layers):
+                    # adjust stochastic depth probability based on the depth of the stage block
+                    sd_prob = stochastic_depth_prob * stage_block_id / (total_stage_blocks - 1.0)
+                    stage.append(block(cnf.input_channels, layer_scale, sd_prob))
+                    stage_block_id += 1
+                layers_list.append(nn.Sequential(*stage).to(device))
+            else:
+                layers_list.append(nn.Identity().to(device))
+
+            if cnf.out_channels is not None:
+                # Upnsampling
+                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)
+        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)):
+                nn.init.trunc_normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+
+    def _forward_impl(self, x) -> Tensor:
+        if self.use_skip:
+            raise NotImplementedError
+            input, features = x
+        else:
+            input = x
+        x = self.first(input)
+        x = self.layers(x)
+        out = self.last(x)
+        
+        return out
+
+    def forward(self, x) -> Tensor:
+        return self._forward_impl(x)
+
+
+def _convnext(
+    block_setting: List[CNBlockConfig],
+    stochastic_depth_prob: float,
+    weights: Optional[WeightsEnum],
+    reversed: bool,
+    input_features: int,
+    **kwargs: Any,
+) -> ConvNeXtBase:
+
+    if reversed:
+        model = ConvNeXtReverseBase(block_setting, stochastic_depth_prob=stochastic_depth_prob,**kwargs)
+    else:
+        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())
+
+    return model
+
+
+class ConvNeXt(nn.Module):
+    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, input_features, **kwargs)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.model(x)
+        return x
+        
+
+
+if __name__ == '__main__':
+    model = ConvNeXt().to('cuda')
+    test = model(torch.zeros((4, 3, 512, 256)).to('cuda'))
+
+    pass

models/depth_estim.py

@@ -3,15 +3,112 @@ import sys
 
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from models.convnext import ConvNeXt, CNBlockConfig
+
+class ViewEncoder(nn.Module):
+    def __init__(self, features=128, device='cuda'):
+        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).to(device)
+        self.act2 = nn.Tanh().to(device)
+    
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        x: torch.Tensor = self.flatten(input)
+        x= self.lin1(x)
+        x = self.act1(x)
+        x = self.lin2(x)
+        x = self.act2(x)
+        x = torch.ones((x.shape[0], self.features, 8, 8), device=x.get_device()) * x[:,:,None,None]
+
+        return x
+
+
+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):
+    def __init__(self, view_features=128, device='cuda', use_skip=False):
         super().__init__()
 
-        self.layer_1 = nn.Linear(in_features=3, out_features=3)  
+        
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, 768, 9),
+            CNBlockConfig(768, None, 3),
+        ]
+        
+        block_setting_reverse = [
+            CNBlockConfig(768 + view_features, 384, 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),
+        ]
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.convnext = ConvNeXt(block_setting, device=device)
-        x = self.layer_1(x)
+        # self.convnext = DumbEncoder()
+        self.flatten = nn.Flatten()
+        self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device)
+        # self.decoder = DumbReconstructor()
+
+        self.view_encoder = ViewEncoder(view_features)
+
+
+    def forward(self, input) -> torch.Tensor:
+        image, view = input
+        x, features = self.convnext(image)
+        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

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

@@ -0,0 +1,138 @@
+import numpy as np
+from struct import pack, unpack
+
+
+"""
+Reads PC2 files, and proposed format PC16 files
+Inputs:
+- file: path to .pc2/.pc16 file
+- float16: False for PC2 files, True for PC16
+Output:
+- data: dictionary with .pc2/.pc16 file data
+NOTE: 16-bit floats lose precision with high values (positive or negative),
+	  we do not recommend using this format for data outside range [-2, 2]
+"""
+def readPC2(file, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	data = {}
+	bytes = 2 if float16 else 4
+	dtype = np.float16 if float16 else np.float32
+	with open(file, 'rb') as f:
+		# Header
+		data['sign'] = f.read(12)
+		# data['version'] = int.from_bytes(f.read(4), 'little')
+		data['version'] = unpack('<i', f.read(4))[0]
+		# Num points
+		# data['nPoints'] = int.from_bytes(f.read(4), 'little')
+		data['nPoints'] = unpack('<i', f.read(4))[0]
+		# Start frame
+		data['startFrame'] = unpack('f', f.read(4))
+		# Sample rate
+		data['sampleRate'] = unpack('f', f.read(4))
+		# Number of samples
+		# data['nSamples'] = int.from_bytes(f.read(4), 'little')
+		data['nSamples'] = unpack('<i', f.read(4))[0]
+		# Animation data
+		size = data['nPoints']*data['nSamples']*3*bytes
+		data['V'] = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
+		data['V'] = data['V'].reshape(data['nSamples'], data['nPoints'], 3)
+		
+	return data
+	
+"""
+Reads an specific frame of PC2/PC16 files
+Inputs:
+- file: path to .pc2/.pc16 file
+- frame: number of the frame to read
+- float16: False for PC2 files, True for PC16
+Output:
+- T: mesh vertex data at specified frame
+"""
+def readPC2Frame(file, frame, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	assert frame >= 0 and isinstance(frame,int), 'Frame must be a positive integer'
+	bytes = 2 if float16 else 4
+	dtype = np.float16 if float16 else np.float32
+	with open(file,'rb') as f:
+		# Num points
+		f.seek(16)
+		# nPoints = int.from_bytes(f.read(4), 'little')
+		nPoints = unpack('<i', f.read(4))[0]
+		# Number of samples
+		f.seek(28)
+		# nSamples = int.from_bytes(f.read(4), 'little')
+		nSamples = unpack('<i', f.read(4))[0]
+		if frame > nSamples:
+			print("Frame index outside size")
+			print("\tN. frame: " + str(frame))
+			print("\tN. samples: " + str(nSamples))
+			return
+		# Read frame
+		size = nPoints * 3 * bytes
+		f.seek(size * frame, 1) # offset from current '1'
+		T = np.frombuffer(f.read(size), dtype=dtype).astype(np.float32)
+	return T.reshape(nPoints, 3)
+
+"""
+Writes PC2 and PC16 files
+Inputs:
+- file: path to file (overwrites if exists)
+- V: 3D animation data as a three dimensional array (N. Frames x N. Vertices x 3)
+- float16: False for writing as PC2 file, True for PC16
+This function assumes 'startFrame' to be 0 and 'sampleRate' to be 1
+NOTE: 16-bit floats lose precision with high values (positive or negative),
+	  we do not recommend using this format for data outside range [-2, 2]
+"""
+def writePC2(file, V, float16=False):
+	assert file.endswith('.pc2') and not float16 or file.endswith('.pc16') and float16, 'File format not consistent with specified input format'
+	if float16: V = V.astype(np.float16)
+	else: V = V.astype(np.float32)
+	with open(file, 'wb') as f:
+		# Create the header
+		headerFormat='<12siiffi'
+		headerStr = pack(headerFormat, b'POINTCACHE2\0',
+						1, V.shape[1], 0, 1, V.shape[0])
+		f.write(headerStr)
+		# Write vertices
+		f.write(V.tobytes())
+
+"""
+Reads OBJ files
+Only handles vertices, faces and UV maps
+Input:
+- file: path to .obj file
+Outputs:
+- V: 3D vertices
+- F: 3D faces
+- Vt: UV vertices
+- Ft: UV faces
+Correspondence between mesh and UV map is implicit in F to Ft correspondences
+If no UV map data in .obj file, it shall return Vt=None and Ft=None
+"""
+def readOBJ(file):
+	V, Vt, F, Ft = [], [], [], []
+	with open(file, 'r') as f:
+		T = f.readlines()
+	for t in T:
+		# 3D vertex
+		if t.startswith('v '):
+			v = [float(n) for n in t.replace('v ','').split(' ')]
+			V += [v]
+		# UV vertex
+		elif t.startswith('vt '):
+			v = [float(n) for n in t.replace('vt ','').split(' ')]
+			Vt += [v]
+		# Face
+		elif t.startswith('f '):
+			idx = [n.split('/') for n in t.replace('f ','').split(' ')]
+			f = [int(n[0]) - 1 for n in idx]
+			F += [f]
+			# UV face
+			if '/' in t:
+				f = [int(n[1]) - 1 for n in idx]
+				Ft += [f]
+	V = np.array(V, np.float32)
+	Vt = np.array(Vt, np.float32)
+	if Ft: assert len(F) == len(Ft), 'Inconsistent .obj file, mesh and UV map do not have the same number of faces' 
+	else: Vt, Ft = None, None
+	return V, F, Vt, Ft

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,4 +1,7 @@
 import os
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+
 import sys
 
 import torch
@@ -7,20 +10,220 @@ import torch.optim as optim
 import torch.nn as nn
 
 from torch.utils.data import DataLoader
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
 from data.dataset import Dataset
 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['mesh_posed'])
+    y_pred, y_pred_normals = model([input, batch['view']])
 
     # 2. Calculate loss
-    loss = loss_fn(y_pred, batch['mesh_posed'])
+    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()
@@ -36,13 +239,24 @@ def train_step(batch):
 
 def valid_step(batch):
     ### Testing
-    model.eval() # put the model in evaluation mode for testing (inference)
+    model.eval() # put the model in evaluation mode for testing (inference) 
     # 1. Forward pass
     with torch.inference_mode():
-        pred = model(batch['mesh_posed'])
+        # 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['mesh_posed'])
+        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() }
 
@@ -57,18 +271,25 @@ def batch_to_device(batch, device):
 if __name__ == '__main__':
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
+    random_view = False
+    use_skip = False
+    batch_size = 1
+    use_cache = True
 
-    train_dataset = Dataset('/home/workspace/local-data/train')
+    train_dataset = Dataset('/data-net/dataset/cloth3d/train', '/data-net/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache)
-    valid_dataset = Dataset('/home/workspace/local-data/valid')
+    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=4, drop_last=True, num_workers=2)
+    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
-    valid_dataloader = DataLoader(valid_dataset, batch_size=4, drop_last=True, num_workers=2)
+    valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
     train_accumulator = Accumulator()
     valid_accumulator = Accumulator()
 
-    model = DepthEstimation()
+    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:
@@ -77,19 +298,23 @@ if __name__ == '__main__':
             train_loss = train_step(batch)
             train_accumulator.update(train_loss)
 
-            if iteration % 1000 == 0:
+            lr_scheduler.step()
-                print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
-                train_accumulator.reset()
-
             iteration += 1
 
-        for batch in valid_dataloader:
+        if iteration % 1000 == 0:
-            valid_loss = valid_step(batch)
+            print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
-            valid_accumulator.update(valid_loss)
+            train_accumulator.reset()
+
+        # TODO: change when batchsize != 1
+        if iteration % 100 == 0:
+            for batch in valid_dataloader:
+                valid_loss = valid_step(batch)
+                valid_accumulator.update(valid_loss)
 
             print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
             valid_accumulator.reset()
 
 
 
+
     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

@@ -0,0 +1,134 @@
+import os
+import sys
+import kornia
+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
+
+
+class Frame:
+    def __init__(self, id, frame_nr, garment):
+        self.id = id
+        self.frame_nr = frame_nr
+        self.garment = garment
+
+    def __hash__(self):
+        return hash((self.id, self.frame_nr))
+
+    def __eq__(self, other):
+        return (self.id, self.frame_nr) == (other.id, other.frame_nr)
+
+
+def mask_faces(F, mask):
+    F_mask = np.take(mask, F.reshape(-1), axis=0)
+    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+
+    V_ind = np.zeros(mask.shape, dtype=np.int32)
+    V_ind[mask] = np.arange(mask.sum())
+    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+
+    return F_new
+
+
+def read_frames_file(file):
+    frames = []
+    with open(file, 'r') as f:
+        lines = f.readlines()
+        for line in lines:
+            id, frame_nr, garment = line.split()
+            frames.append(Frame(id, int(frame_nr), garment))
+    return frames
+
+
+def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
+    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_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()
@@ -25,13 +25,3 @@ class Accumulator():
         self.dict = dict()
 
 
-
-def mask_faces(F, mask):
-    F_mask = np.take(mask, F.reshape(-1), axis=0)
-    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
-
-    V_ind = np.zeros(mask.shape, dtype=np.int32)
-    V_ind[mask] = np.arange(mask.sum())
-    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
-
-    return F_new