luni/cvpr
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

compare predictions with diff curvs

Browse Source
This commit is contained in:
luni
2023-01-29 17:34:54 +00:00
parent cbfa129999
commit 1c24bc0d0e
2 changed files with 90 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_curv.py
+1 -1
View File
@@ -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)
curvature_analysis.py
+89 -24
View File
@@ -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