Spaces:
Sleeping
Sleeping
File size: 3,137 Bytes
99a05f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import numpy as np
import os
import json
import trimesh
import seaborn as sns
# Load the combined dca train, val and test npzs
dir = '/is/cluster/work/stripathi/pycharm_remote/dca_contact/data/dataset_extras'
trainval_npz = np.load(os.path.join(dir, 'hot_dca_trainval.npz'), allow_pickle=True)
test_npz = np.load(os.path.join(dir, 'hot_dca_test.npz'), allow_pickle=True)
# combine the two npz
combined_npz = {}
for key in trainval_npz.keys():
combined_npz[key] = np.concatenate([trainval_npz[key], test_npz[key]], axis=0)
segmentation_path = 'data/smpl_vert_segmentation.json'
with open(segmentation_path, 'rb') as f:
part_segmentation = json.load(f)
combine_keys = {'leftFoot': ['leftToeBase'],
'rightFoot': ['rightToeBase'],
'leftHand': ['leftHandIndex1'],
'rightHand': ['rightHandIndex1'],
'spine': ['spine1', 'spine2'],
'head': ['neck'],}
for key in combine_keys:
for subkey in combine_keys[key]:
part_segmentation[key] += part_segmentation[subkey]
del part_segmentation[subkey]
# reverse the part segmentation
part_segmentation_rev = {}
for part in part_segmentation:
for vert in part_segmentation[part]:
part_segmentation_rev[vert] = part
# count the number of contact instances per vertex
per_vert_contact_count = np.zeros(6890)
for cls in combined_npz['contact_label']:
per_vert_contact_count += cls
# calculate the maximum contact count per part
part_contact_max = {}
for part in part_segmentation:
part_contact_max[part] = np.max(per_vert_contact_count[part_segmentation[part]])
# calculate the contact probability globally
contact_prob = np.zeros(6890)
for vid in range(6890):
contact_prob[vid] = (per_vert_contact_count[vid] / max(per_vert_contact_count)) ** 0.3
# save the contact probability mesh
outdir = "/is/cluster/work/stripathi/pycharm_remote/dca_contact/hot_analysis"
# load template smpl mesh
mesh = trimesh.load_mesh('data/smpl/smpl_neutral_tpose.ply')
vertex_colors = trimesh.visual.interpolate(contact_prob, 'jet')
# set the vertex colors of the mesh
mesh.visual.vertex_colors = vertex_colors
# save the mesh
out_path = os.path.join(outdir, "contact_probability_mesh.obj")
mesh.export(out_path)
# # calculate the contact probability per part
# contact_prob = np.zeros(6890)
# for vid in range(6890):
# if 'Hand' in part_segmentation_rev[vid]:
# contact_prob[vid] = (per_vert_contact_count[vid] / part_contact_max[part_segmentation_rev[vid]]) ** 0.4 if 'Hand' not in part_segmentation_rev[vid] else (per_vert_contact_count[vid] / part_contact_max[part_segmentation_rev[vid]]) ** 0.8
#
# # save the contact probability mesh
# outdir = "/is/cluster/work/stripathi/pycharm_remote/dca_contact/hot_analysis"
#
# # load template smpl mesh
# mesh = trimesh.load_mesh('data/smpl/smpl_neutral_tpose.ply')
# vertex_colors = trimesh.visual.interpolate(contact_prob, 'jet')
# # set the vertex colors of the mesh
# mesh.visual.vertex_colors = vertex_colors
# # save the mesh
# out_path = os.path.join(outdir, "contact_probability_mesh_part.obj")
# mesh.export(out_path)
|