diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index 2756698..61a8feb 100644 --- a/auto_encoder_curv.py +++ b/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) diff --git a/curvature_analysis.py b/curvature_analysis.py index 7bfe55a..eaf5bea 100644 --- a/curvature_analysis.py +++ b/curvature_analysis.py @@ -123,7 +123,8 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6): 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') # get_bins(curvature, nr_of_labels) + 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) @@ -133,6 +134,44 @@ def generate_curvatures(v, faces, mask, neighbourhood=25, nr_of_labels=6): return frame_labels +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] @@ -144,12 +183,23 @@ def visualize_curvature_labels(folder, results_folder, frame): 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] + 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 = generate_curvatures(verts, faces, mask, 25) curvature_labels_pred = generate_curvatures(pred, faces, mask, 25) - dice = dice_dense(curvature_labels_pred[None,:], curvature_labels_gt[None,:]) + curvatures_new = (uv2meshlabel(uv_curvature[None,:,:,:])[0].numpy() * mask[:,None])[mask, :] + curvatures_pred_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]: + 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): @@ -157,20 +207,7 @@ def visualize_curvature_labels(folder, results_folder, frame): colors_binned[curvature_labels == i, :] = colors[i] colors_binned_all.append(colors_binned) - fig = make_subplots(rows=1, cols=3, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(3)] ]) - - 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=np.ones((verts.shape[0],3))*200 - ), - row=1, col=1) + 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 @@ -183,8 +220,8 @@ def visualize_curvature_labels(folder, results_folder, frame): k = faces[:,2], vertexcolor=colors_binned_all[0] ), - row=1, col=2) - + row=1, col=1) + fig.add_trace(go.Mesh3d( # 8 vertices of a cube x=-pred[:,0], @@ -197,12 +234,39 @@ def visualize_curvature_labels(folder, results_folder, frame): 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), + 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), @@ -212,7 +276,7 @@ def visualize_curvature_labels(folder, results_folder, frame): fig.update_layout(scene = helper, scene2 = helper, scene3 = helper, - # scene4 = helper, + scene4 = helper, # scene5 = helper, ) @@ -241,13 +305,13 @@ def visualize_curvature_labels(folder, results_folder, frame): import wandb def init_wandb(): - run = wandb.init(project='plots', name='review_4', dir='/home/workspace/local-data/temp/') + 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() + # init_wandb() is_skirt = False skirt_modifier = '_skirt' if is_skirt else '' @@ -272,7 +336,8 @@ if __name__ == '__main__': i = 0 while True: - frame = np.random.choice(frames) + # 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