126 lines
3.6 KiB
Python
126 lines
3.6 KiB
Python
import os
|
|
|
|
|
|
import sys
|
|
import shutil
|
|
import tqdm
|
|
import random
|
|
import math
|
|
import torch
|
|
import numpy as np
|
|
import scipy as sp
|
|
import scipy.io as sio
|
|
import plotly.graph_objects as go
|
|
import plotly.express as ex
|
|
from pytorch3d.structures import Meshes
|
|
|
|
from custom_utils.custom_utils import _maskFaces
|
|
from losses.losses import _nearest_neighbour
|
|
|
|
|
|
# from hbertiche/PBNS
|
|
""" COLLISION FUNCTIONS START """
|
|
# F = np.load("smpl/smpl_faces.npy")
|
|
to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
|
|
F_fix_coll = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[to_keep,:][None,:,:], [4,1,1]), device='cuda')
|
|
|
|
|
|
def fix_collision(V: torch.Tensor, B: torch.Tensor, mask: torch.Tensor, thr=.004):
|
|
mesh = Meshes(B, F_fix_coll)
|
|
B_vn = torch.stack(mesh.verts_normals_list())
|
|
|
|
# compute correspondences
|
|
idx = _nearest_neighbour(V, B)
|
|
# compute loss
|
|
B_idx = torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
|
|
D = V - B_idx
|
|
b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
|
|
dot = torch.einsum('abc,abc->ab', D, b_vn)
|
|
loss = torch.minimum(dot - thr, torch.zeros_like(dot)) * mask
|
|
B_offset_closest = B_idx + (b_vn / 2e2)
|
|
output = torch.where(loss[:,:,None] < 0, B_offset_closest, V)
|
|
|
|
return output
|
|
|
|
|
|
""" COLLISION FUNCTIONS END """
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
F = np.load("custom_utils/smpl_faces.npy")
|
|
posed = np.load('/home/workspace/local-data/train/mesh_posed/00016_10.npz')['mesh_posed'][np.newaxis,:,:]
|
|
body = np.load('/home/workspace/local-data/train/mesh_body_posed/00016_10.npy')[np.newaxis,:,:]
|
|
mask = np.load('/home/workspace/local-data/train/mesh_mask/00016.npy')
|
|
posed_2 = np.load('/home/workspace/local-data/train/mesh_posed/00016_15.npz')['mesh_posed'][np.newaxis,:,:]
|
|
|
|
fixed = fix_collision(posed_2, body, mask).numpy()
|
|
|
|
x = posed_2[0, mask,:]
|
|
y = body[0]
|
|
_, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
|
|
fig = go.Figure(data=[
|
|
go.Mesh3d(
|
|
x=x[:,0],
|
|
y=x[:,1],
|
|
z=x[:,2],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2],
|
|
),
|
|
go.Mesh3d(
|
|
x=y[:,0],
|
|
y=y[:,1],
|
|
z=y[:,2],
|
|
i = F[:,0],
|
|
j = F[:,1],
|
|
k = F[:,2],
|
|
),
|
|
])
|
|
|
|
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],),)
|
|
)
|
|
|
|
fig.show()
|
|
|
|
x = fixed[0, mask,:]
|
|
y = body[0]
|
|
_, faces = _maskFaces(np.zeros((14475,3), np.float32), F, mask)
|
|
fig = go.Figure(data=[
|
|
go.Mesh3d(
|
|
x=x[:,0],
|
|
y=x[:,1],
|
|
z=x[:,2],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2],
|
|
),
|
|
go.Mesh3d(
|
|
x=y[:,0],
|
|
y=y[:,1],
|
|
z=y[:,2],
|
|
i = F[:,0],
|
|
j = F[:,1],
|
|
k = F[:,2],
|
|
),
|
|
])
|
|
|
|
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],),)
|
|
)
|
|
|
|
fig.show()
|
|
|
|
print() |