import os import sys import random import numpy as np import plotly.graph_objects as go 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 ( read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame, writeOBJ, readOBJ, quads2tris ) from smpl.smpl_torch import SMPL # from custom_utils.custom_utils import _maskFaces 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): V_mask = mask F_mask = np.take(V_mask, F.reshape(-1), axis=0) F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3 V_ind = np.zeros(V_mask.shape) V_ind[V_mask] = np.arange(V_mask.sum()) F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3)) V_new = V[V_mask, :] if return_unfixed: return V_new, F_new, F[F_mask, :] else: return V_new, F_new def get_frames(folder, frames_file): frames = [] with open(os.path.join(folder, frames_file), 'r') as f: for line in f: frames.append(line.strip()) return frames def plot(mesh, faces): fig = go.Figure() fig.add_trace(go.Mesh3d( x=mesh[:,0], y=mesh[:,1], z=mesh[:,2], i = faces[:,0], j = faces[:,1], k = faces[:,2] ),) fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), ) helper = dict(xaxis = dict(nticks=4, range=[-1,1],), yaxis = dict(nticks=4, range=[-1,1],), zaxis = dict(nticks=4, range=[-1,1],),) fig.update_layout(scene = helper, ) fig.show() def build_lookup_table(repeated): # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element lookup = {} for i in range(len(repeated)): for j in range(len(repeated[i])): if len(repeated[i]) == 1: continue lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j return lookup 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) map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1)) if deorient: 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 def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False): pc16_path = os.path.join(folder, id, 'Body_posed.pc16') verts = readPC2Frame(pc16_path, frame, True) if deorient: 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 # try to optimize the already registered garment ?????????? 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) boundaries = raw_mask == 2 registered_boundaries = torch.tensor(np.where(boundaries[mask])[0]).to(device) # 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_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) ##################### 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) registered = template[map][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) template_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) registered = torch.tensor(registered).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=registered[None,:,:], faces=F[None,:,:]) mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:]) mesh1_boundaries = registered_boundaries mesh2_boundaries = template_boundaries # mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries) avg_edge_length = mesh_edge_lengths(mesh1).mean() # Define the optimizer and learning rate 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") # cotcurv uniform cot def edge_loss(mesh): return mesh_edge_loss(mesh, target_length=avg_edge_length) def normal_consistency_loss(mesh): return mesh_normal_consistency(mesh) # Optimization loop 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()]) 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 * 3 + (bound_ch_loss) * 20 # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1 lap_loss = laplacian_smoothing_loss(mesh1_out) e_loss = edge_loss(mesh1_out) # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100 norm_loss = normal_consistency_loss(mesh1_out) loss += lap_loss * 0.05 + e_loss * 100 + norm_loss * 0.1 # Backward pass loss.backward() optimizer.step() # Print the loss if i % 10 == 0: print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} bound_chamfer: {bound_ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}") # create a plotly plot with the body and the mesh1_out registered = mesh1_out.verts_packed().detach().cpu().numpy() template = mesh2.verts_packed().detach().cpu().numpy() faces = faces.detach().cpu().numpy() F = F.detach().cpu().numpy() # faces, F = F, faces # verts, star_body = star_body, verts mesh1_boundaries, mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy(), mesh2_boundaries.detach().cpu().numpy() # mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy() fig = go.Figure(data=[ go.Mesh3d( x=template[:,0], y=template[:,1], z=template[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=0.80), go.Mesh3d( x=registered[:,0], y=registered[:,1], z=registered[:,2], i=F[:,0], j=F[:,1], k=F[:,2], color='lightpink', opacity=0.50), go.Scatter3d( x=registered[mesh1_boundaries,0], y=registered[mesh1_boundaries,1], z=registered[mesh1_boundaries,2], mode='markers', marker=dict(size=2, color='red')), go.Scatter3d( x=template[mesh2_boundaries,0], y=template[mesh2_boundaries,1], z=template[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() # create a plotly wireframe plot of registered fig = go.Figure() for edge in mesh1_out.edges_packed().detach().cpu().numpy(): fig.add_trace(go.Scatter3d( x=registered[edge,0], y=registered[edge,1], z=registered[edge,2], mode='lines', line=dict(color='red', width=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/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') root = '/'.join(folder.split('/')[:-1]) type = folder.split('/')[-1] use_no_orient = False no_shape = False is_skirt = False skirt_modifier = '_skirt' if is_skirt else '' orient_mod = '_no_orient' if use_no_orient else '' shape_mod = '_no_shape' if no_shape else '' # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy') # faces = np.load('/data-net/dataset/00761/faces/00761.npy') for frame in frames: frame = random.choice(frames) # frame = '00761 0 Tshirt' frame = '01893 256 Jumpsuit' register(folder, frame) break pass