vto conversion, some progress, not working properl

configs/direct_1.py

@@ -0,0 +1,77 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'direct_1'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = False
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+config['input_features'] = 6
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_body_posed',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        }

configs/direct_2.py

@@ -0,0 +1,78 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '0'
+config['name'] = 'direct_2'
+config['notes'] = 'with curvature input' 
+config['tags'] = ['torch', 'direct'] 
+config['script'] = 'scripts/direct_train.py'
+
+config['pretrained_model'] = ''
+config['use_curvature'] = True
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/train'
+config['valid_dir'] = '/data-local/127/valid'
+config['test_dir'] = '/data-net/dataset/test_ready'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = False
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+config['input_features'] = 10
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['is_label'] = False
+config['kernel_size'] = 7
+config['bottleneck_size'] = 64 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 128, 3),
+    CNBlockConfig(128, 256, 9),
+    CNBlockConfig(256, config['bottleneck_size'], 9),
+    CNBlockConfig(config['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['bottleneck_size'], 256, 3),
+    CNBlockConfig(256, 128, 9),
+    CNBlockConfig(128, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+    # CNBlockConfig(3, 3, 3),
+]
+
+config['encoder_block_settings'] = block_setting
+config['decoder_block_settings'] = block_setting_reverse
+config['activation'] = 'tanh'
+
+
+config['lr'] = 0.0001
+config['reg_loss_coeff'] = 0.1
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'pose_shape',
+                        'uv_curvature',
+                        'mesh_mask',
+                        'mesh_body_posed',
+                        'mesh_posed',
+                        'uv_unposed',
+                        'mesh_unposed',
+                        'uv_static',
+                        'uv_mask',
+                        'uv_body_posed',
+                        }

custom_utils/custom_utils.py

@@ -43,6 +43,7 @@ class UV2Mesh(nn.Module):
         if uv.shape[1] == 3 or uv.shape[1] == 7:
             uv = torch.permute(uv, [0,2,3,1])
         if uv.shape[1] != 512:
+            raise NotImplementedError
             uv = Resize(uv[:,:,None], size=(512, 256), interpolation=InterpolationMode.BILINEAR)
             
         # batch_size = uv.shape[0]

data/convert_vto.py

@@ -0,0 +1,259 @@
+import numpy as np
+import os
+import sys
+import pickle
+import plotly.express as ex
+import plotly.graph_objects as go
+import torch
+import open3d as o3d
+import trimesh
+# from vedo import trimesh as trimesh_plot
+from scipy.spatial import KDTree
+from scipy.sparse import lil_matrix
+from scipy.optimize import minimize
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from smpl.smpl_torch import SMPL
+
+# import mask faces
+from custom_utils.custom_utils import _maskFaces
+
+smpl = SMPL(path='smpl')
+
+def energy_function(params, source_points, target_points, indices):
+    num_points = source_points.shape[0]
+    num_dims = source_points.shape[1]
+    
+    # Reshape the parameters
+    weights = params[:num_points * num_dims].reshape((num_points, num_dims))
+    translation = params[num_points * num_dims:num_points * num_dims + num_dims]
+    control_points = params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
+    
+    # Compute the transformed source points
+    transformed_points = source_points + np.dot(weights, control_points.T)
+    
+    # Compute the TPS energy
+    tps_energy = 0.5 * np.sum(np.square(target_points[indices] - transformed_points))
+    
+    return tps_energy
+
+def non_rigid(source_points, target_points):
+    tree = KDTree(target_points)
+    distances, indices = tree.query(source_points)
+
+    # Initialize the optimization parameters
+    num_points = source_points.shape[0]
+    num_dims = source_points.shape[1]
+    initial_params = np.zeros(num_points * num_dims + num_dims + num_points)
+
+    # Define the optimization bounds
+    bounds = [(None, None)] * (num_points * num_dims)
+    bounds += [(None, None)] * num_dims
+    bounds += [(None, None)] * (num_points * num_dims)
+
+    # Perform the optimization
+    result = minimize(energy_function, initial_params, method='L-BFGS-B',
+                    args=(source_points, target_points, indices), bounds=bounds)
+
+    # Get the optimized parameters
+    optimized_params = result.x
+
+    # Reshape the parameters
+    optimized_weights = optimized_params[:num_points * num_dims].reshape((num_points, num_dims))
+    optimized_translation = optimized_params[num_points * num_dims:num_points * num_dims + num_dims]
+    optimized_control_points = optimized_params[num_points * num_dims + num_dims:].reshape((-1, num_dims))
+
+    transformed_points = source_points + np.dot(optimized_weights, optimized_control_points.T)
+    return transformed_points
+
+
+def register_smpl():
+    # generate body with zero pose and shape
+    pose = torch.zeros((1, 72))
+    shape = torch.zeros((1, 11))
+    smpl_out = smpl(pose, shape)[0][0]
+
+    with open('/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl', "rb") as f:
+        data = pickle.load(f)
+        garment = data['vertices'][100]
+        faces = data['faces']
+
+    # rotate garment x degrees along y axis
+    # convert degrees to radians
+    degrees = 107
+    theta = degrees * np.pi / 180
+    rot = np.array([[np.cos(theta), 0, np.sin(theta)],
+                    [0, 1, 0],
+                    [-np.sin(theta), 0, np.cos(theta)]])
+    garment = np.dot(rot, garment.T).T
+
+    # rotate garment x degrees along x axis
+    # convert degrees to radians
+    degrees = 5
+    theta = degrees * np.pi / 180
+    rot = np.array([[1, 0, 0],
+                    [0, np.cos(theta), -np.sin(theta)],
+                    [0, np.sin(theta), np.cos(theta)]])
+    garment = np.dot(rot, garment.T).T
+
+    # rotate garment x degrees along z axis
+    # convert degrees to radians
+    degrees = -5
+    theta = degrees * np.pi / 180
+    rot = np.array([[np.cos(theta), -np.sin(theta), 0],
+                    [np.sin(theta), np.cos(theta), 0],
+                    [0, 0, 1]])
+    garment = np.dot(rot, garment.T).T
+
+
+    garment[:,0] -= 0.06
+    garment[:,1] -= 0.03
+    garment[:,2] -= 0.015
+    
+    # visualize the body and garment
+    # fig = go.Figure(data=[go.Mesh3d(x=smpl_out[:,0], y=smpl_out[:,1], z=smpl_out[:,2], i=smpl.f[:,0], j=smpl.f[:,1], k=smpl.f[:,2])])
+    # fig.add_trace(go.Mesh3d(x=garment[:,0], y=garment[:,1], z=garment[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2]))
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    # fig.show()
+
+    # non_rigid(garment, smpl_out.numpy())
+
+    # Load the source and target meshes
+    source_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(garment), o3d.utility.Vector3iVector(faces))
+    target_mesh = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(smpl_out), o3d.utility.Vector3iVector(smpl.f))
+    # flip source and target mesh
+    source_mesh, target_mesh = target_mesh, source_mesh
+
+    source_mesh.compute_vertex_normals()
+    source_mesh = source_mesh.filter_smooth_laplacian(number_of_iterations=10)
+    # plot the smoothed mesh with plotly
+    # fig = go.Figure(data=[go.Mesh3d(x=np.asarray(source_mesh.vertices)[:,0], y=np.asarray(source_mesh.vertices)[:,1], z=np.asarray(source_mesh.vertices)[:,2], i=np.asarray(source_mesh.triangles)[:,0], j=np.asarray(source_mesh.triangles)[:,1], k=np.asarray(source_mesh.triangles)[:,2])])
+    # fig.show()
+
+    # Perform initial alignment (translation and/or scaling)
+    # You can modify this step based on your specific needs
+    # For example, you can manually set an initial transformation
+    init_transformation = [[1.0, 0.0, 0.0, 0.0],
+                        [0.0, 1.0, 0.0, 0.0],
+                        [0.0, 0.0, 1.0, 0.0],
+                        [0.0, 0.0, 0.0, 1.0]]
+
+    # Create the ICP object and set the parameters
+    icp = o3d.pipelines.registration.TransformationEstimationPointToPoint()
+    criteria = o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=2000)
+        
+    # Perform ICP registration
+    reg_result = o3d.pipelines.registration.registration_icp(o3d.geometry.PointCloud(source_mesh.vertices),
+                                                             o3d.geometry.PointCloud(target_mesh.vertices),
+                                                             0.075,
+                                                             init_transformation, icp, criteria)
+
+    # Get the aligned source mesh after registration
+    aligned_mesh = source_mesh.transform(reg_result.transformation)
+
+    # Visualize the aligned mesh and target mesh with plotly
+    fig = go.Figure(data=[go.Mesh3d(x=np.asarray(aligned_mesh.vertices)[:,0], y=np.asarray(aligned_mesh.vertices)[:,1], z=np.asarray(aligned_mesh.vertices)[:,2], i=np.asarray(aligned_mesh.triangles)[:,0], j=np.asarray(aligned_mesh.triangles)[:,1], k=np.asarray(aligned_mesh.triangles)[:,2])])
+    fig.add_trace(go.Mesh3d(x=np.asarray(target_mesh.vertices)[:,0], y=np.asarray(target_mesh.vertices)[:,1], z=np.asarray(target_mesh.vertices)[:,2], i=np.asarray(target_mesh.triangles)[:,0], j=np.asarray(target_mesh.triangles)[:,1], k=np.asarray(target_mesh.triangles)[:,2]))
+    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    fig.show()
+
+    # Get the vertices of the aligned source mesh and target mesh
+    aligned_source_vertices = np.asarray(aligned_mesh.vertices)
+    target_vertices = np.asarray(target_mesh.vertices)
+
+    # Perform nearest neighbor search to find corresponding vertex indices
+    kd_tree = o3d.geometry.KDTreeFlann(target_mesh)
+    correspondences = []
+    matched = set()
+    for i in range(len(aligned_source_vertices)):
+        k = 1
+        _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
+        # while idx[-1] in matched:
+        #     k += 1
+        #     _, idx, _ = kd_tree.search_knn_vector_3d(aligned_source_vertices[i], k)
+        # matched.add(idx[-1])
+        correspondences.append((i, idx[-1]))
+
+    # correspondences is a list of tuples where each tuple contains the corresponding vertex indices
+    # For example, correspondences = [(0, 10), (1, 5), (2, 15), ...]
+
+    # Visualize the correspondences with plotly
+    # fig = go.Figure(data=[go.Scatter3d(x=aligned_source_vertices[:,0], y=aligned_source_vertices[:,1], z=aligned_source_vertices[:,2], mode='markers', marker=dict(color='red', size=1))])
+    # fig.add_trace(go.Scatter3d(x=target_vertices[:,0], y=target_vertices[:,1], z=target_vertices[:,2], mode='markers', marker=dict(color='blue', size=1)))
+    # for corr in correspondences:
+    #     fig.add_trace(go.Scatter3d(x=[aligned_source_vertices[corr[0],0], target_vertices[corr[1],0]], y=[aligned_source_vertices[corr[0],1], target_vertices[corr[1],1]], z=[aligned_source_vertices[corr[0],2], target_vertices[corr[1],2]], mode='lines', line=dict(color='green', width=1)))
+    # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    # fig.show()
+
+    registered = np.zeros((14475, 3))
+    mask = np.zeros((14475), dtype=bool)
+    for i in range(len(correspondences)):
+        registered[correspondences[i][1]] = aligned_source_vertices[correspondences[i][0]]
+        mask[correspondences[i][1]] = True
+
+    # F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
+
+    # registered_masked, masked_faces = _maskFaces(registered, np.asarray(target_mesh.triangles), mask)
+    registered_masked, masked_faces = registered, np.asarray(target_mesh.triangles)
+
+    # plot aligned source mesh with the target topology
+    fig = go.Figure(data=[go.Mesh3d(x=registered_masked[:,0], y=registered_masked[:,1], z=registered_masked[:,2], i=masked_faces[:,0], j=masked_faces[:,1], k=masked_faces[:,2])])
+    fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+    fig.show()
+
+    return smpl
+
+
+def convert_vto(folder, out_folder):
+    os.makedirs(out_folder, exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_posed_body_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'pose_shape'), exist_ok=True)
+    F = np.load(f'smpl/smpl_faces.npy').astype(np.int32)
+
+    files = os.listdir(folder)
+    # files = ['/data-net/code/vto-dataset/tshirt/simulations/tshirt_shape00_55_27.pkl']
+    for file in files:
+        # read pickle files and get vetices and faces from them
+        with open(os.path.join(folder, file), "rb") as f:
+            data = pickle.load(f)
+            vertices = data['vertices']
+            faces = data['faces']
+            pose = data['pose']
+            shape = data['shape']
+            trans = data['translation']
+            id = file.split('.')[0]
+            # pose[:, :3] = 0
+            id = id.replace('_', '-')
+            subject = id.split('-')[2]
+            sequence = '_'.join(id.split('-')[2:])
+            # cmu_data = np.load(f'/data-net/code/vto-garment-collisions/assets/CMU/{subject}/{sequence}_poses.npz')
+            # shape = np.tile(cmu_data['betas'], (shape.shape[0], 1))
+            # pose = cmu_data['poses']
+
+            body = smpl(torch.tensor(pose), torch.tensor(np.concatenate((shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)))[0].numpy()
+            pose_shape = np.concatenate((pose, shape, np.zeros((shape.shape[0],1),dtype=np.float32)),axis=1)
+            
+            # for i in range(vertices.shape[0]):
+            #     np.savez_compressed(os.path.join(out_folder, 'mesh_posed_body_no_orient', f'{id}_{i}.npz'), mesh_posed_body_no_orient=body[i])
+            #     np.save(os.path.join(out_folder, 'pose_shape', f'{id}_{i}.npy'), pose_shape[i])
+
+
+            # plot mesh with plotly
+            frame = 0
+            fig = go.Figure(data=[go.Mesh3d(x=vertices[frame,:,0], y=vertices[frame,:,1], z=vertices[frame,:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])])
+            fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,1], z=body[frame,:,2], i=F[:,0], j=F[:,1], k=F[:,2]))
+            fig.update_layout(scene=dict(xaxis=dict(range=[-1,1],), yaxis=dict(range=[-1,1],), zaxis=dict(range=[-1,1],),),)
+            fig.show()
+
+            
+            pass
+    pass
+
+
+if __name__ == '__main__':
+    # register_smpl()
+    convert_vto('/data-net/code/vto-dataset/dress/simulations', '/data-net/code/vto-dataset/dress_converted')
+    pass

data/dataset_vto.py

@@ -0,0 +1,40 @@
+import os
+import sys
+
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+if args['debug']:
+    os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+import math
+import time
+import torch
+import random
+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.utils import Frame
+
+class Dataset(Dataset):
+    def __init__(self, folder):
+        pass
+
+    def __len__(self):
+        return len(self.frames)
+
+
+    def __getitem__(self, idx):
+        frame : Frame = self.frames[idx]
+
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader

models/base.py

@@ -35,10 +35,10 @@ class VariationalEncoder(nn.Module):
 
 class AEEncoder(nn.Module):
     
-    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda'):
+    def __init__(self, block_setting: CNBlockConfig, kernel_size: int = 7, activation='tanh', is_label=False, nr_of_labels=None, device='cuda', input_features=3):
         super(AEEncoder, self).__init__()
 
-        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device)
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, is_label=is_label, nr_of_labels=nr_of_labels, device=device, input_features=input_features)
         if activation == 'tanh':
             self.act = nn.Tanh()
         elif activation == 'sigmoid':

models/cond_encoder.py

@@ -8,7 +8,7 @@ class AutoEncoder(nn.Module):
     def __init__(self, config):
         super(AutoEncoder, self).__init__()
 
-        self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7)
+        self.encoder : AEEncoder = AEEncoder(config['encoder_block_settings'], config['kernel_size'], config['activation'], is_label=config['is_label'], device=config['device'], nr_of_labels=7, input_features=config['input_features'])
         self.decoder : Decoder = Decoder(config['decoder_block_settings'], kernel_size=config['kernel_size'], is_label=config['is_label'], device=config['device'], nr_of_classes=6)
 
 

models/direct.py

@@ -0,0 +1,29 @@
+import torch
+from torch import nn
+from models.convnext_orig import ConvNeXt, CNBlock, CNBlockConfig
+from typing import Any, Callable, List, Optional, Sequence
+
+
+class DirectModel(nn.Module):
+    
+    def __init__(self, block_setting: List[CNBlockConfig], kernel_size: int = 7, input_features: int = 16, device: str = 'cuda'):
+        super(DirectModel, self).__init__()
+
+        self.encoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, input_features=input_features, device=device)
+        self.decoder = ConvNeXt(block_setting=block_setting, kernel_size=kernel_size, device=device, reversed=True)
+
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.mean_branch = nn.Sequential(*layers).to(device)
+        layers = [CNBlock(block_setting[-1].input_channels, kernel_size=kernel_size) for _ in range(5)]
+        self.var_branch = nn.Sequential(*layers).to(device)
+
+        self.flatten_mean = nn.Flatten()
+        self.flatten_var = nn.Flatten()
+
+
+    def forward(self, inputs):
+        body, template = inputs
+
+        x = self.encoder(x)
+
+        return mean

run.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.pyr_64_cape_m import config
+from configs.direct_1 import config
 import argparse
 
 # parser = argparse.ArgumentParser(description='')

scripts/direct_train.py

@@ -0,0 +1,281 @@
+import os
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+import importlib
+import argparse
+parser = argparse.ArgumentParser(description='')
+parser.add_argument('--name', type=str)
+parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
+args = vars(parser.parse_args())
+
+config_module = importlib.import_module(args['config'])
+config = config_module.config
+if args['debug']:
+    from configs.direct_1 import config
+    # from configs.ae_static_orig_2 import config
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+import time
+import argparse
+import wandb
+import random
+import numpy as np
+import torch
+# from torchmetrics import MultiScaleStructuralSimilarityIndexMeasure as MSSSIM
+from custom_utils.custom_utils import UV2Mesh, get_summary_html, maskFacesBatch, plot_3d_collision
+import plotly.graph_objects as go
+import plotly.express as ex
+import torch.optim as optim
+
+if config['is_skirt']:
+    from losses.losses import edge_loss_skirt as edge_loss
+else:
+    from losses.losses import edge_loss
+
+from data.dataset import Dataset
+# from torch.utils.data import DataLoader
+from losses.losses import collision_loss, normal_loss
+from utils.postprocess import fix_collision
+from custom_utils.custom_utils import plot_3d
+from utils.checkpoints import CheckpointSaver, load_model
+from custom_utils.custom_utils import Accumulator
+from utils.utils import InfiniteDataLoader as DataLoader
+from models.cond_encoder import AutoEncoder as Model
+from smpl.smpl_torch import SMPL
+
+from pytorch_msssim import MS_SSIM
+
+
+def init_wandb():
+    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('scripts/direct_train.py')
+    artifact.add_file('losses/losses.py')
+    artifact.add_file('models/base.py')
+    artifact.add_file('models/cond_encoder.py')
+    artifact.add_file('models/convnext.py')
+    artifact.add_file('data/dataset.py')
+    config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
+    for config_file_name in config_file_names:
+        file_name = '/'.join(config_file_name.split('.')) + '.py'
+        if os.path.exists(file_name):
+            artifact.add_file(file_name)
+    run.log_artifact(artifact)
+
+uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+
+
+def train_step(batch):
+    model.train()
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
+    if config['use_curvature']:
+        input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
+    else:
+        input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
+    pred = model([input, batch['uv_mask']])
+
+    mesh_pred = uv2mesh(pred)
+    mesh_orig = batch['mesh_unposed']
+    shape = batch['pose_shape'][:, 72:]
+    gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+    shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+    mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+    mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+ 
+    e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+    c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+    c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+    # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+    l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+    n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+    loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
+
+    model.eval()
+    with torch.inference_mode():
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/euclid_error': metric, 'train/collision_ratio': c_ratio,
+             'train/ms_ssim': ms_ssim, 'train/collision_loss': c_loss, 'train/l2_mesh_loss': l2_mesh_loss,
+             'train/normal_loss': n_loss }
+        
+
+def valid_step(batch):
+    model.eval()
+
+    with torch.inference_mode():
+        if config['use_curvature']:
+            input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
+        else:
+            input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
+        pred = model([input, batch['uv_mask']])
+        mesh_pred = uv2mesh(pred)
+        mesh_orig = batch['mesh_unposed']
+        shape = batch['pose_shape'][:, 72:]
+        gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+        shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+        mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+        mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+    
+        e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask']))
+        c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask'])
+        c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
+        # l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((batch['uv_unposed'] - pred) ** 2), axis=-1), axis=[1, 2]))
+        l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, axis=-1), axis=-1))
+        n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
+
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_unposed'] + 1))
+        loss = l2_mesh_loss + (e_loss + c_loss + n_loss) * config['reg_loss_coeff']
+
+        metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_mask'])
+    
+    return { 'valid/loss': loss, 'valid/edge_loss': e_loss, 'valid/euclid_error': metric, 'valid/collision_ratio': c_ratio,
+             'valid/ms_ssim': ms_ssim, 'valid/collision_loss': c_loss, 'valid/l2_mesh_loss': l2_mesh_loss,
+             'valid/normal_loss': n_loss }
+       
+
+def plot_everything(batch: dict, type: str):
+    model.eval()
+    with torch.inference_mode():
+        if config['use_curvature']:
+            input = torch.cat([batch['uv_curvature'], batch['uv_static']], axis=1)
+        else:
+            input = torch.cat([batch['uv_body_posed'], batch['uv_static']], axis=1)
+        pred = model([input, batch['uv_mask']]).detach()
+
+    mesh_pred = uv2mesh(pred).cpu().numpy()
+    mesh_orig = batch['mesh_unposed'].cpu().numpy()
+    
+    orig_img = torch.permute(batch['uv_unposed'], [0,2,3,1])
+    uv_static_img = torch.permute(batch['uv_static'], [0,2,3,1])
+    uv_body_posed_img = torch.permute(batch['uv_body_posed'], [0,2,3,1])
+    pred_img = torch.permute(pred, [0,2,3,1])
+    img = torch.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
+    imgs = [wandb.Image((img[i].cpu().numpy() + 1) / 2.0) for i in range(4)]
+    
+    F = np.load("smpl/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
+    mask = batch['mesh_mask'].cpu().numpy().astype(bool)
+
+    shape = batch['pose_shape'][:, 72:]
+    gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+    shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]).cpu().numpy()
+
+    garment_unposed, faces = maskFacesBatch(mesh_pred, F, mask)
+    garment_unposed = [garment_unposed[i] + shape_offsets[i, mask[i], :] for i in range(mask.shape[0])]
+    gt = [mesh_orig[i][mask[i], :] + shape_offsets[i, mask[i], :] for i in range(pred.shape[0])]
+
+    for i in range(mesh_orig.shape[0]):
+        mesh_pred[i, mask[i], :] += shape_offsets[i, mask[i], :]
+    mesh_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], torch.tensor(mesh_pred, device=config['device']))[0].cpu().numpy()
+
+    gt_posed = [batch['mesh_posed'].cpu().numpy()[i, mask[i], :] for i in range(mask.shape[0])]
+
+    mesh_posed_fixed = fix_collision(torch.tensor(mesh_posed, device=config['device']), batch['mesh_body_posed'], torch.tensor(mask, device=config['device'])).cpu().numpy()
+    body = batch['mesh_body_posed'].cpu().numpy()
+
+    garment_posed = [mesh_posed[i, mask[i], :] for i in range(body.shape[0])]
+    collision = [np.concatenate((body[i], garment_posed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_faces = [np.concatenate((F, faces[i] + config['smpl_res']), axis=0) for i in range(len(garment_posed))]
+    collision_color = [np.concatenate((np.ones_like(body[i], dtype=np.int32)*100, np.ones_like(garment_posed[i], dtype=np.int32)*230), axis=0) for i in range(len(garment_posed))]
+
+    garment_posed_fixed = [mesh_posed_fixed[i, mask[i], :] for i in range(body.shape[0])]
+    collision_fixed = [np.concatenate((body[i], garment_posed_fixed[i]), axis=0) for i in range(len(garment_posed))]
+    collision_gt = [np.concatenate((body[i], gt_posed[i]), axis=0) for i in range(len(garment_posed))]
+
+    figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
+    figs_coll = { type + "/garment_" + str(i) + '_collision': plot_3d_collision(collision_gt[i], collision[i], collision_fixed[i], collision_faces[i], collision_color[i]) for i in range(4) }
+
+    return { **figs, **figs_coll, **{ type + '/images': imgs, } }
+
+
+def train():
+    valid_iterator = iter(valid_dataloader)
+
+    step = 0
+    for batch in train_dataloader:
+        train_loss = train_step(batch)
+        train_accumulator.update(train_loss)
+
+        if step % 25 == 0:
+            data_valid = next(valid_iterator)
+            valid_loss = valid_step(data_valid)
+            valid_accumulator.update(valid_loss)
+
+        if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
+            to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
+            if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
+                train_figs = plot_everything(batch, 'train')
+                valid_figs = plot_everything(data_valid, 'valid')
+                to_log = { **to_log, **train_figs, **valid_figs }
+            if step == 0:
+                to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
+
+            if not args['debug']:
+                wandb.log(to_log, step=step)
+            train_accumulator.reset()
+            valid_accumulator.reset()
+            print(step)
+
+        if step % 10000 == 0:
+            ckpt_saver_enc(model.encoder, step, optimizer, scaler)
+            ckpt_saver_dec(model.decoder, step, optimizer, scaler)
+
+        step += 1
+
+
+if __name__ == '__main__':
+    if not args['debug']:
+        init_wandb()
+    
+    use_amp = True
+    skirt_modifier = '_skirt' if config['is_skirt'] else ''
+
+    # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
+    valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
+    train_accumulator = Accumulator()
+    valid_accumulator = Accumulator()
+
+    model : Model = Model(config)
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
+
+    # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
+
+    if config['pretrained_model'] != '':
+        load_model(model.encoder, config['pretrained_model'] + '_enc')
+        load_model(model.decoder, config['pretrained_model'] + '_dec')
+
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
+
+    run_name = config['name']
+    print(f'started training: {run_name}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
+    train()
+
+
+

smpl/smpl_torch.py

@@ -206,9 +206,9 @@ class SMPL(Module):
         )
       )
     
-    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
+    joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.device)], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
-    T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    T = torch.stack([torch.tensordot(results[i], self.weights_all[torch.greater(g, 0.5).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
     )
@@ -230,4 +230,4 @@ class SMPL(Module):
     if remove and verts is None:
       result = result[:,self.to_keep,:]
 
-    return result, joints, v_shaped, v_posed
+    return result, joints