registration progress save
This commit is contained in:
+282
-28
@@ -7,6 +7,9 @@ import torch
|
||||
from torch.nn.functional import mse_loss
|
||||
from pytorch3d.structures import Meshes
|
||||
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
|
||||
from pymeshfix import MeshFix
|
||||
|
||||
# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
|
||||
|
||||
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
||||
from data.cloth3d_io import (
|
||||
@@ -26,7 +29,40 @@ from data.cloth3d_io import (
|
||||
from smpl.smpl_torch import SMPL
|
||||
# from custom_utils.custom_utils import _maskFaces
|
||||
|
||||
smpl = SMPL('cuda', 'smpl')
|
||||
device = 'cuda'
|
||||
smpl = SMPL(device, 'smpl')
|
||||
|
||||
|
||||
import torch.nn.functional as F
|
||||
from pytorch3d.structures import Meshes
|
||||
|
||||
def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Given two PyTorch3D meshes and a set of vertex indices for mesh1,
|
||||
find the closest vertices from mesh2 to these vertices.
|
||||
|
||||
Args:
|
||||
mesh1: A PyTorch3D Meshes object representing the first mesh.
|
||||
mesh2: A PyTorch3D Meshes object representing the second mesh.
|
||||
vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
|
||||
|
||||
Returns:
|
||||
A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
|
||||
"""
|
||||
# Extract the vertices of mesh1 and mesh2.
|
||||
vertices1 = mesh1.verts_packed()
|
||||
vertices2 = mesh2.verts_packed()
|
||||
|
||||
# Extract the vertices of mesh1 that correspond to the given indices.
|
||||
vertices1_selected = vertices1[vertex_indices]
|
||||
|
||||
# Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
|
||||
distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
|
||||
|
||||
# Find the indices of the closest vertices in mesh2.
|
||||
closest_indices = torch.argmin(distances, dim=1)
|
||||
|
||||
return closest_indices
|
||||
|
||||
|
||||
def _maskFaces(V, F, mask, return_unfixed=False):
|
||||
@@ -90,6 +126,8 @@ def build_lookup_table(repeated):
|
||||
from data.cloth3d_io import readPC2Frame
|
||||
import scipy.io as sio
|
||||
import scipy
|
||||
import scipy.spatial as spatial
|
||||
|
||||
def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
|
||||
pc16_path = os.path.join(folder, id, garment + '.pc16')
|
||||
verts = readPC2Frame(pc16_path, frame, True)
|
||||
@@ -97,8 +135,8 @@ def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorie
|
||||
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
|
||||
|
||||
if deorient:
|
||||
R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
|
||||
# R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
|
||||
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
|
||||
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
|
||||
verts = R.apply(verts)
|
||||
|
||||
return verts, map
|
||||
@@ -108,77 +146,293 @@ def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
|
||||
verts = readPC2Frame(pc16_path, frame, True)
|
||||
|
||||
if deorient:
|
||||
R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
|
||||
# R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
|
||||
R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
|
||||
# R = spatial.transform.Rotation.from_euler([0,0,-zrot])
|
||||
verts = R.apply(verts)
|
||||
|
||||
return verts
|
||||
|
||||
def edge_loss_custom(x, y, edges):
|
||||
x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
|
||||
x_e = torch.sum(x_e ** 2, -1)
|
||||
x_e = torch.sqrt(x_e + 1e-9)
|
||||
y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
|
||||
y_e = torch.sum(y_e ** 2, -1)
|
||||
y_e = torch.sqrt(y_e + 1e-9)
|
||||
d_e = y_e - x_e
|
||||
loss = torch.mean(d_e ** 2, dim=-1)[0]
|
||||
return loss
|
||||
|
||||
|
||||
def crotch(mesh, body, crotch_idx, nr=75):
|
||||
crotch_middle_on_body = crotch_idx
|
||||
# create kdtree for body and mesh
|
||||
body_tree = spatial.cKDTree(body)
|
||||
mesh_tree = spatial.cKDTree(mesh)
|
||||
# find closest point on mesh from crotch middle on body along the z axis
|
||||
|
||||
crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
|
||||
|
||||
crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
|
||||
crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
|
||||
|
||||
body_verts = body[crotch_body]
|
||||
mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
|
||||
mesh_verts = spatial.cKDTree(mesh_verts)
|
||||
closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
|
||||
crotch_mesh = crotch_mesh[closest_mesh_verts]
|
||||
|
||||
return crotch_body, crotch_mesh
|
||||
|
||||
|
||||
def find_closest_point_along_z(point, verts, eps=0.01):
|
||||
min_dist = np.inf
|
||||
closest_point = None
|
||||
for i, v in enumerate(verts):
|
||||
dist = np.linalg.norm(point - v)
|
||||
if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
|
||||
min_dist = dist
|
||||
closest_point = i
|
||||
return closest_point
|
||||
|
||||
|
||||
def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
|
||||
"""
|
||||
Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
|
||||
|
||||
Args:
|
||||
mesh: A PyTorch3D Meshes object representing the mesh.
|
||||
|
||||
Returns:
|
||||
A tensor of shape (E,) representing the lengths of all edges in the mesh,
|
||||
where E is the total number of edges in the mesh.
|
||||
"""
|
||||
# Extract the vertices and faces of the mesh.
|
||||
vertices = mesh.verts_packed()
|
||||
faces = mesh.faces_packed()
|
||||
|
||||
# Compute the pairwise distances between all pairs of vertices that form an edge.
|
||||
edge_indices = faces.view(-1, 3)
|
||||
edge_vertices = vertices[edge_indices]
|
||||
edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
|
||||
|
||||
return edge_lengths
|
||||
|
||||
|
||||
def register(folder, frame):
|
||||
id, frame_nr, garment = frame.split()
|
||||
frame_nr = int(frame_nr)
|
||||
info = loadInfo(os.path.join(folder, id, 'info.mat'))
|
||||
raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
|
||||
# convert it to indices
|
||||
mask = raw_mask > 0
|
||||
mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
|
||||
body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
|
||||
template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
|
||||
faces = quads2tris(faces)
|
||||
pose = torch.zeros(72)
|
||||
shape = torch.zeros(11)
|
||||
shape[-1] = 1
|
||||
|
||||
# boundaries = raw_mask == 2
|
||||
# # invert map and use it to invert boudaries
|
||||
# map_inv = np.zeros(map.max()+1, dtype=np.int32)
|
||||
# map_inv[map] = np.arange(len(map))
|
||||
# boundaries = boundaries[map_inv]
|
||||
# # get the boundary indices
|
||||
|
||||
# mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
|
||||
|
||||
|
||||
|
||||
# pose = torch.zeros(1, 72).to('cuda')
|
||||
shape = torch.zeros(1, 11).to(device)
|
||||
shape[:, :10] = torch.tensor(info['shape']).to(device)
|
||||
shape[:, -1] = torch.tensor(info['gender']).to(device)
|
||||
# for star pose
|
||||
pose[:,5] = 3.1415/4
|
||||
pose[:,8] = -3.1415/4
|
||||
star_body = smpl(pose.unsqueeze(0), shape.unsqueeze(0))[0][0].cpu().numpy()
|
||||
# pose[:, 5] = 3.1415/4
|
||||
# pose[:,8] = -3.1415/4
|
||||
STAR_POSE = np.zeros((24, 3), np.float32)
|
||||
STAR_POSE[1,2] = .15
|
||||
STAR_POSE[2,2] = -.15
|
||||
pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
|
||||
star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
|
||||
# rotate star body with 90 degrees around x axis
|
||||
R = spatial.transform.Rotation.from_euler('x', np.pi/2)
|
||||
star_body = R.apply(star_body).astype(np.float32)
|
||||
|
||||
is_skirt = (garment == 'Dress' or garment == 'Skirt')
|
||||
skirt_mod = '_skirt' if is_skirt else ''
|
||||
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
|
||||
|
||||
plot(template, faces)
|
||||
##################### mask body #####################
|
||||
new_crotch = 4729
|
||||
|
||||
# star_body = star_body[mask]
|
||||
# new_crotch = np.zeros((14475), bool)
|
||||
# new_crotch[4729] = True
|
||||
# new_crotch = new_crotch[mask]
|
||||
# new_crotch = np.where(new_crotch)[0][0]
|
||||
# _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
|
||||
|
||||
# plot(template, faces)
|
||||
# v = mesh_posed[mask,:]
|
||||
|
||||
##################### boundaries #####################
|
||||
# meshfix = MeshFix(template, faces)
|
||||
# holes = meshfix.extract_holes()
|
||||
# new_border = []
|
||||
# for border in np.array(holes.points):
|
||||
# # find the index of the vertex in the original mesh
|
||||
# ind = np.argwhere(np.all(template == border, axis=1))
|
||||
# new_border.extend(ind.ravel())
|
||||
# mesh1_boundaries = torch.tensor(new_border).to(device)
|
||||
# np.save('temp/body_boundaries.npy', new_border)
|
||||
mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
|
||||
|
||||
# meshfix = MeshFix(template, faces)
|
||||
# holes = meshfix.extract_holes()
|
||||
# new_border = []
|
||||
# for border in np.array(holes.points):
|
||||
# # find the index of the vertex in the original mesh
|
||||
# ind = np.argwhere(np.all(template == border, axis=1))
|
||||
# new_border.extend(ind.ravel())
|
||||
# mesh1_boundaries = torch.tensor(new_border).to(device)
|
||||
# np.save('temp/body_boundaries.npy', new_border)
|
||||
# mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
|
||||
|
||||
##################### crotch #####################
|
||||
crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
|
||||
|
||||
|
||||
|
||||
|
||||
# export star_body as obj
|
||||
# writeOBJ('temp/star_body.obj', star_body, F)
|
||||
|
||||
|
||||
template = torch.tensor(template).to(device)
|
||||
template_orig = template.clone()
|
||||
faces = torch.tensor(faces).to(device)
|
||||
star_body = torch.tensor(star_body).to(device)
|
||||
F = torch.tensor(F).to(device)
|
||||
|
||||
# verts = template.cpu().numpy()
|
||||
# star_body = star_body.cpu().numpy()
|
||||
# faces = faces.detach().cpu().numpy()
|
||||
# F = F.detach().cpu().numpy()
|
||||
# fig = go.Figure(data=[
|
||||
# go.Mesh3d(
|
||||
# x=verts[:,0], y=verts[:,1], z=verts[:,2],
|
||||
# i=faces[:,0], j=faces[:,1], k=faces[:,2],
|
||||
# color='lightpink', opacity=0.50),
|
||||
# go.Mesh3d(
|
||||
# x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
|
||||
# i=F[:,0], j=F[:,1], k=F[:,2],
|
||||
# color='lightblue', opacity=0.50),
|
||||
# go.Scatter3d(
|
||||
# x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
|
||||
# mode='markers', marker=dict(size=2, color='red')),
|
||||
# go.Scatter3d(
|
||||
# x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
|
||||
# mode='markers', marker=dict(size=2, color='blue')),
|
||||
# ])
|
||||
# # make it cuve and limit to -1 and -1
|
||||
# fig.update_layout(scene = dict(
|
||||
# xaxis = dict(range=[-1,1],),
|
||||
# yaxis = dict(range=[-1,1],),
|
||||
# zaxis = dict(range=[-1,1],),),)
|
||||
# fig.show()
|
||||
|
||||
# Define the meshes to optimize
|
||||
mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
|
||||
mesh2 = Meshes(verts=star_body, faces=F[None,:,:])
|
||||
mesh2 = Meshes(verts=star_body[None,:,:], faces=F[None,:,:])
|
||||
|
||||
mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
|
||||
|
||||
avg_edge_length = mesh_edge_lengths(mesh1).mean()
|
||||
|
||||
# Define the optimizer and learning rate
|
||||
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.01)
|
||||
mesh1.verts_packed().requires_grad = True
|
||||
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
|
||||
# optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
|
||||
|
||||
# Define the loss functions
|
||||
def laplacian_smoothing_loss(mesh):
|
||||
return mesh_laplacian_smoothing(mesh, method="uniform")
|
||||
return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
|
||||
|
||||
def edge_loss(mesh1, mesh2):
|
||||
return mesh_edge_loss(mesh1, mesh2)
|
||||
def edge_loss(mesh):
|
||||
return mesh_edge_loss(mesh, target_length=avg_edge_length)
|
||||
|
||||
def normal_consistency_loss(mesh1, mesh2):
|
||||
return mesh_normal_consistency(mesh1, mesh2)
|
||||
def normal_consistency_loss(mesh):
|
||||
return mesh_normal_consistency(mesh)
|
||||
|
||||
# Optimization loop
|
||||
for i in range(100):
|
||||
for i in range(10000):
|
||||
optimizer.zero_grad()
|
||||
|
||||
# find boundaries with the find holes function?
|
||||
|
||||
# Forward pass
|
||||
mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
|
||||
loss = chamfer_distance(mesh1_out, mesh2)
|
||||
loss += laplacian_smoothing_loss(mesh1_out)
|
||||
loss += edge_loss(mesh1_out)
|
||||
loss += normal_consistency_loss(mesh1_out)
|
||||
ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
|
||||
bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
|
||||
# crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
|
||||
crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
|
||||
loss: torch.Tensor = ch_loss * 10 + (bound_ch_loss + crotch_ch_loss) * 5
|
||||
# loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
|
||||
loss += laplacian_smoothing_loss(mesh1_out) * 0.05
|
||||
loss += edge_loss(mesh1_out) * 100
|
||||
# loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
|
||||
loss += normal_consistency_loss(mesh1_out) * 0.1
|
||||
|
||||
# Backward pass
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# Print the loss
|
||||
print(f"Loss: {loss.item()}")
|
||||
if i % 10 == 0:
|
||||
print(f"Loss: {loss.item()} chamfer: {ch_loss.item()} ")
|
||||
|
||||
# create a plotly plot with the body and the mesh1_out
|
||||
verts = mesh1_out.verts_packed().detach().cpu().numpy()
|
||||
star_body = star_body.detach().cpu().numpy()
|
||||
faces = faces.detach().cpu().numpy()
|
||||
F = F.detach().cpu().numpy()
|
||||
mesh1_boundaries = mesh1_boundaries.detach().cpu().numpy()
|
||||
mesh2_boundaries = mesh2_boundaries.detach().cpu().numpy()
|
||||
fig = go.Figure(data=[
|
||||
go.Mesh3d(
|
||||
x=verts[:,0], y=verts[:,1], z=verts[:,2],
|
||||
i=faces[:,0], j=faces[:,1], k=faces[:,2],
|
||||
color='lightpink', opacity=0.80),
|
||||
go.Mesh3d(
|
||||
x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
|
||||
i=F[:,0], j=F[:,1], k=F[:,2],
|
||||
color='lightblue', opacity=0.50),
|
||||
go.Scatter3d(
|
||||
x=verts[mesh1_boundaries,0], y=verts[mesh1_boundaries,1], z=verts[mesh1_boundaries,2],
|
||||
mode='markers', marker=dict(size=2, color='red')),
|
||||
go.Scatter3d(
|
||||
x=star_body[mesh2_boundaries,0], y=star_body[mesh2_boundaries,1], z=star_body[mesh2_boundaries,2],
|
||||
mode='markers', marker=dict(size=2, color='blue')),
|
||||
go.Scatter3d(
|
||||
x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
|
||||
mode='markers', marker=dict(size=2, color='red')),
|
||||
go.Scatter3d(
|
||||
x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
|
||||
mode='markers', marker=dict(size=2, color='blue')),
|
||||
])
|
||||
# make it cuve and limit to -1 and -1
|
||||
fig.update_layout(scene = dict(
|
||||
xaxis = dict(range=[-1,1],),
|
||||
yaxis = dict(range=[-1,1],),
|
||||
zaxis = dict(range=[-1,1],),),)
|
||||
fig.update_layout(scene_aspectmode='cube')
|
||||
fig.show()
|
||||
|
||||
pass
|
||||
|
||||
if __name__ == "__main__":
|
||||
folder = '/data-net/dataset/cloth3d/train'
|
||||
# folder = r'D:\Work\Dataset\cloth3d\train'
|
||||
folder = '/data-net/code/cloth3d/train'
|
||||
# folder = r'D:\Work\code\cloth3d\train'
|
||||
# frames = [f'00761_{i}' for i in range(300)]
|
||||
# folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
|
||||
frames = get_frames(folder, 'frames_all.txt')
|
||||
@@ -203,5 +457,5 @@ if __name__ == "__main__":
|
||||
|
||||
|
||||
register(folder, frame)
|
||||
|
||||
break
|
||||
pass
|
||||
|
||||
Reference in New Issue
Block a user