344 lines
14 KiB
Python
344 lines
14 KiB
Python
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] #cloud.get_neighbors(k=neighbourhood)
|
|
#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
|
|
|
|
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]
|
|
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)
|
|
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, 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) |