try

auto_encoder_curv.py

@@ -95,7 +95,7 @@ def dice_dense(prediction, ground_truth, weight_map=None):
     return tf.reduce_mean(dice_score)
 
 
-@tf.function
+# @tf.function
 def train_step(batch):
     with tf.GradientTape() as tape:
         pred = model([batch['uv_curvature'], batch['uv_mask']], training=True)

configs/config_abl_overfit_jumpsuit_2.py

@@ -19,7 +19,7 @@ config['only_curvatures'] =           True
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '5'
+config['gpu'] = '3'
 config['name'] = 'abl_overfit_jumpsuit_2'
 config['notes'] = 'remove other conditions from input as well'
 

curvature_analysis.py

@@ -0,0 +1,348 @@
+import os
+import sys
+os.environ["CUDA_VISIBLE_DEVICES"] = '4'
+
+
+import tqdm
+import argparse
+import random
+import string
+import time
+import numpy as np
+import trimesh as tm
+import cmapy as cm
+import scipy as sp
+import scipy.io as sio
+import plotly.graph_objects as go
+from pyntcloud import PyntCloud
+from scipy.sparse import csr_matrix, coo_matrix
+from scipy.sparse.csgraph import floyd_warshall, dijkstra
+from plotly.subplots import make_subplots
+from custom_utils.custom_utils import _maskFaces
+from postprocessing import fix_collision
+
+from custom_utils.supersmpl_io import readPC2, readPC2Frame, writePC2, readOBJ
+from custom_utils.custom_utils import _maskFaces
+from auto_encoder_curv import dice_dense
+
+def postprocess(mesh, faces, mask=None, body=None, do_smoothing=True):
+    if body is not None:
+        mesh_fixed = fix_collision(mesh, body[np.newaxis,:,:], mask[np.newaxis,:])[0].numpy()[mask, :]
+    else:
+        mesh_fixed = mesh
+
+    if do_smoothing:
+        m = tm.Trimesh(mesh_fixed, faces, process=False, validate=False)
+        outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
+        outliers = np.nonzero(outliers)[0]
+        outliers_n_neighbours = get_closest_n_points(mesh_fixed, outliers, 2)
+        outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
+        faces_fixed = np.concatenate((faces, outliers_faces), axis=0)
+        m = tm.Trimesh(mesh_fixed, faces_fixed, process=False, validate=False)
+        tm.smoothing.filter_humphrey(m, iterations=1, alpha=0.6, beta=0.4)
+
+        mesh = m.vertices
+    else:
+        mesh = mesh_fixed
+
+    return mesh
+
+def export_to_ply(mesh, file):
+    ply = tm.exchange.ply.export_ply(mesh, encoding='binary', vertex_normal=True, include_attributes=True)
+    output_file = open(file, "wb+")
+    output_file.write(ply)
+    output_file.close()
+
+
+def get_closest_n_points(verts, idx, n):
+    dist = sp.spatial.distance.cdist(verts[idx, :], verts)
+    dist[dist == 0] = 1e9
+    idx_sort = np.argsort(dist, axis=-1)
+    closest = idx_sort[:, :n]
+    return closest
+
+
+def calculate_dist(folder, neighbourhood):
+    out = dict()
+    frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
+
+    for frame in tqdm.tqdm(frames):
+        id = frame.split('_')[0]
+        if id in out.keys():
+            continue
+
+        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
+        k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
+        out[id] = k_neighbors
+
+    return out
+
+
+def histedges_equalN(x, nbin):
+    npt = len(x)
+    return np.interp(np.linspace(0, npt, nbin + 1),
+                     np.arange(npt),
+                     np.sort(x))
+
+
+def get_bins(values, nr_of_bins):
+    flattened = []
+    for c in values:
+        flattened.extend(c.flatten())
+    flattened = np.array(flattened)
+    bins = histedges_equalN(flattened, nr_of_bins)
+    return bins
+
+
+def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6):
+    # id = frame.split('_')[0]
+    # frame_nr = int(frame.split('_')[1].split('.')[0])
+
+    # mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
+    # _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+    # v = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
+    
+    m = tm.Trimesh(v, faces, process=False, validate=False)
+
+    outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
+    outliers = np.nonzero(outliers)[0]
+    outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
+    outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
+    faces = np.concatenate((faces, outliers_faces), axis=0)
+
+    verts = v
+
+    m = tm.Trimesh(verts, faces, process=False, validate=False)
+    export_to_ply(m, temp_file)
+
+    cloud = PyntCloud.from_file(temp_file)
+
+    k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood]
+    # k_neighbors = cloud.get_neighbors(k=neighbourhood)
+    # tree = sp.spatial.cKDTree(verts)
+    # k_neighbors = tree.query(verts, 25)[1]
+    #k_neighbors = k_neighbors_cache[id] # np.argsort(dist[mask, :][:, mask], axis=-1)[:, :neighbourhood] #cloud.get_neighbors(k=neighbourhood)
+    ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
+    cloud.add_scalar_field("curvature", ev=ev)
+    curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
+    
+    bins = np.load(f'custom_data/train_bins.npy')
+    # bins = get_bins(curvature, nr_of_labels)
+    labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
+
+    frame_labels = np.zeros((curvature.shape[0]), dtype=np.float32)
+    for j in range(bins.shape[0] - 2):
+        frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
+    frame_labels[curvature >= bins[-2]] = labels[-1]
+       
+    return frame_labels, curvature
+
+import tensorflow as tf
+from configs.config_abl_overfit_jumpsuit_2 import config
+from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision, UV2MeshLabel
+
+from custom_utils.smpl import SMPL
+from models.cvae_model import CVAE
+
+uv2mesh = UV2Mesh(res=256, is_skirt=False, is_adaptive=False, is_smpl_orig=False)
+uv2meshlabel = UV2MeshLabel(res=256, is_skirt=False)
+
+def get_prediction(folder, frame, curv):
+    smpl = SMPL(use_one_hot=False, orig_smpl=False)
+    model = CVAE(config)
+    load_model(model, 'checkpoints/' + config['name'])
+
+    id = frame.split('_')[0]
+    frame_nr = frame.split('_')[1].split('.')[0]
+
+    mesh_mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
+    uv_mask = np.load(os.path.join(folder, 'uv_mask', f'{id}.npz'))['uv_mask'].astype(np.float32)  
+    pose_shape = np.load(os.path.join(folder, 'pose_shape', f'{id}_{frame_nr}.npy')).astype(np.float32)  
+    pose_shape[:3] = 0
+    uv_unposed = np.load(os.path.join(folder, 'uv_unposed_no_shape', f'{id}_{frame_nr}.npz'))['uv_unposed_no_shape']
+    uv_static = np.load(os.path.join(folder, 'uv_static', f'{id}.npz'))['uv_static']
+    uv_body_posed = np.load(os.path.join(folder, 'uv_posed_body_no_orient_2', f'{id}_{frame_nr}.npz'))['uv_posed_body_no_orient']
+
+    pred, _, _ = model([uv_unposed[None,:,:,:], uv_static[None,:,:,:], uv_body_posed[None,:,:,:], curv[None,:,:,:], uv_mask[None,:,:]], training=False)
+    mesh_pred = uv2mesh(pred)
+
+    shape = pose_shape[None, 72:]
+    gender = tf.cast(tf.greater_equal(shape[:, 10] , 0.5), tf.int32)
+    shape_offsets = tf.stack([tf.gather(tf.tensordot(shape[i,:10][tf.newaxis, :], smpl.shapedirs[gender[i]], axes=[[1], [2]])[0,:,:], smpl.to_keep, axis=0) for i in range(shape.shape[0])])
+
+    mesh_pred_shaped = mesh_pred + shape_offsets * mesh_mask[tf.newaxis,:,tf.newaxis]
+    mesh_pred_posed = smpl.run_garment(pose_shape[None, :72], pose_shape[None, 72:], mesh_pred_shaped)
+    
+    return mesh_pred_posed.numpy()[0,mesh_mask,:]
+
+
+def visualize_curvature_labels(folder, results_folder, frame):
+    id = frame.split('_')[0]
+    frame_nr = frame.split('_')[1].split('.')[0]
+
+    mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy')).astype(bool)
+    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)
+    verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}_{frame_nr}.npz'))['mesh_posed'][mask,:]
+    pred = np.load(os.path.join(results_folder, frame + '.npy'))
+    pred = postprocess(pred, faces)
+
+    uv_curvature = np.load(os.path.join(folder, 'uv_curvature', f'{id}_{frame_nr}.npz'))['curvature_uv']
+    mesh_curvature = np.load(os.path.join(folder, 'mesh_curvature', f'{id}_{frame_nr}.npz'))['mesh_curvature'][mask]
+    mesh_curvature_raw = np.load(os.path.join(folder, 'mesh_curvature_raw', f'{id}_{frame_nr}.npz'))['mesh_curvature_raw'][mask]
+    uv_curvature = tf.one_hot(uv_curvature, 7)
+    pred_new = get_prediction(folder, frame, uv_curvature)
+    pred_new = postprocess(pred_new, faces)
+
+    curvature_labels_gt, curv_gt = generate_curvatures(verts, faces, mask, 25)
+    curvature_labels_pred, _ = generate_curvatures(pred, faces, mask, 25)
+    curvatures_new = (uv2meshlabel(uv_curvature[None,:,:,:])[0].numpy() * mask[:,None])[mask, :]
+    curvatures_pred_new, curv_gt_new = generate_curvatures(pred_new, faces, mask, 25)
+    dice = dice_dense(tf.one_hot(curvature_labels_pred, 7)[None,:], tf.one_hot(curvature_labels_gt, 7)[None,:])
+    dice_new = dice_dense(tf.one_hot(curvatures_pred_new, 7)[None,:], curvatures_new[None,:])
+    curvatures_new = curvatures_new.argmax(axis=-1) - 1
+
+    colors_binned_all = []
+
+    for curvature_labels in [curvature_labels_gt, curvature_labels_pred, curvatures_new, curvatures_pred_new]:
+        colors = np.array([cm.color('jet', 1 - i) for i in np.linspace(0, 1, int(curvature_labels.max()+1))])
+        colors_binned = np.zeros((curvature_labels.shape[0], 3))
+        for i in range(int(curvature_labels.max())+1):
+            # print(np.count_nonzero(curvature_labels == i))
+            colors_binned[curvature_labels == i, :] = colors[i]
+        colors_binned_all.append(colors_binned)
+
+    fig = make_subplots(rows=1, cols=4, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(4)] ])
+
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=verts[:,0],
+                y=verts[:,1],
+                z=verts[:,2],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors_binned_all[0]
+            ),
+            row=1, col=1)
+
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=-pred[:,0],
+                y=pred[:,2],
+                z=pred[:,1],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors_binned_all[1],
+                text=f'dice: {dice}'
+            ),
+            row=1, col=2)
+    
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=verts[:,0],
+                y=verts[:,1],
+                z=verts[:,2],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors_binned_all[2]
+            ),
+            row=1, col=3)
+
+    fig.add_trace(go.Mesh3d(
+                # 8 vertices of a cube
+                x=-pred_new[:,0],
+                y=pred_new[:,2],
+                z=pred_new[:,1],
+                # i, j and k give the vertices of triangles
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2],
+                vertexcolor=colors_binned_all[3],
+                text=f'dice: {dice_new}'
+            ),
+            row=1, col=4)
+
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1),
+                    scene3_aspectmode='manual', scene3_aspectratio=dict(x=1, y=1, z=1),
+                    scene4_aspectmode='manual', scene4_aspectratio=dict(x=1, y=1, z=1),
+                    # scene5_aspectmode='manual', scene5_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                yaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),
+                zaxis = dict(range=[-1,1], showbackground=False, showticklabels=False, showaxeslabels=False, showspikes=False, visible=False),)
+        
+    fig.update_layout(scene = helper,
+                        scene2 = helper,
+                        scene3 = helper,
+                        scene4 = helper,
+                        # scene5 = helper,
+                        )
+
+
+    # meshes = [go.Mesh3d(
+    #             x=x[:,0],
+    #             y=x[:,1],
+    #             z=x[:,2],
+    #             i = faces[:,0],
+    #             j = faces[:,1],
+    #             k = faces[:,2],
+    #             vertexcolor=colors_binned
+    #             )]
+
+    # fig = go.Figure(data=meshes)
+    # fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1))
+    # fig.update_layout(
+    #     scene = dict(
+    #         xaxis = dict(nticks=4, range=[-1,1],),
+    #                     yaxis = dict(nticks=4, range=[-1,1],),
+    #                     zaxis = dict(nticks=4, range=[-1,1],),)
+    #                     )
+
+    return fig, dice
+
+import wandb
+
+def init_wandb():
+    run = wandb.init(project='plots', name='review_5', dir='/home/workspace/local-data/temp/')
+    artifact = wandb.Artifact('code', type='code')
+    artifact.add_file('curvature_analysis.py')
+    run.log_artifact(artifact)
+
+if __name__ == '__main__':
+    # init_wandb()
+
+    is_skirt = False
+    skirt_modifier = '_skirt' if is_skirt else ''
+    F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
+    # k_neighbours(k=25)
+
+    folder = '/home/workspace/local-data/train'
+    results_folder = '/home/workspace/data/results/post/abl_overfit_jumpsuit_2-'
+
+    frames_file = 'frames_jumpsuit_30k.txt'
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        lines = f.readlines()
+        lines = [line.strip() for line in lines]
+    frames = lines
+
+    temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
+    if not is_skirt:
+        dist = np.load('custom_data/dijkstra.npy')
+    else:
+        dist = np.load('custom_data/dijkstra_skirt.npy')
+
+    i = 0
+    while True:
+        # frame = np.random.choice(frames)
+        frame = '02806_129'
+        fig, dice = visualize_curvature_labels(folder, results_folder, frame)
+        wandb.log({f'{frame}-dice:{dice}': fig}, step=i)
+        i += 1
+        print(frame)

cvae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_all_post import config
+from configs.config_abl_overfit_jumpsuit_2 import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -41,7 +41,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAENew as CVAE
+from models.cvae_model import CVAE
 
 
 def init_wandb():

evaluate_ae.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from configs.config_abl_all_post import config
+from configs.config_abl_overfit_jumpsuit_2 import config
 
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 import tensorflow as tf
@@ -42,7 +42,7 @@ from losses import collision_loss as collision_penalty
 from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mesh, get_summary_html, maskFacesBatch, _maskFaces, plot_3d_collision
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
-from models.cvae_model import CVAENew as CVAE
+from models.cvae_model import CVAE
 from custom_utils.supersmpl_io import writeOBJ
 # from unif_metric import chamfer_and_score
 
@@ -104,7 +104,7 @@ if __name__ == '__main__':
 
     iteration = ''
 
-    load_model(model, 'checkpoints/post/' + checkpoint_name, iteration=iteration)
+    load_model(model, 'checkpoints/' + checkpoint_name, iteration=iteration)
     # load_model(model, 'checkpoints/' + 'abl_all_120k')
 
     output_folder = '/home/workspace/data/results/post/' + checkpoint_name + '-' + iteration