import os import sys import argparse parser = argparse.ArgumentParser() # parser.add_argument("-f", "--folder", help = "end of range") parser.add_argument("-s", "--start", help = "start of range", type=int, default=0) parser.add_argument("-e", "--end", help = "end of range", type=int, default=1) parser.add_argument("-g", "--gpu", help = "gpu id", default='0') parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true') args = parser.parse_args() os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu os.environ["CUDA_LAUNCH_BLOCKING"] = '1' import math import shutil import warnings import tables as tb import numpy as np import scipy.io as sio import trimesh as tm from tqdm import tqdm from tables import NaturalNameWarning warnings.filterwarnings('ignore', category=NaturalNameWarning) import plotly.graph_objects as go import torch from scipy.spatial.transform import Rotation as R sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris, readPC2Frame from data.utils import Frame from smpl.smpl_torch import SMPL from custom_utils.custom_utils import _maskFaces, remove_long_edges def create_extras_file(folder): # create a new pytables file root = '/'.join(folder.split('/')[:-1]) type = folder.split('/')[-1] h5file = tb.open_file(os.path.join(root, f'{type}_extras_test.h5'), mode='w', title='Extras') frames = [] with open(os.path.join(folder, 'frames_all.txt'), 'r') as f: lines = f.readlines() lines = [line.strip() for line in lines] frames = lines for frame in tqdm(frames[:1000]): # add a new empty group to the root group for each frame id, frame_nr, garment = frame.split() if not id in h5file.root: h5file.create_group('/', id) if not frame_nr in h5file.root[id]: h5file.create_group(f'/{id}', frame_nr) h5file.create_group(f'/{id}/{frame_nr}', f'{garment}') h5file.close() def deorient(verts, pose): rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).cpu().numpy() rotated = np.zeros_like(verts) for frame in range(verts.shape[0]): r = R.from_matrix(rodr[frame, 0, :, :]) rotated[frame] = r.apply(verts[frame], inverse=True) return rotated def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder): ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True) os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True) ids = ['00021'] for id in tqdm(ids): # id = '01467' garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] info = loadInfo(os.path.join(folder, id, 'info.mat')) for garment in garments: # if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): # continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] # if info['poses'].shape[1] == verts.shape[0]: # continue reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat')) verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1 mask = np.zeros((14475), dtype=bool) mask[mask_idx] = True if info['poses'].shape[1] != verts.shape[0]: assert False, 'Somethign wrong here' info['poses'] = info['poses'][:, :verts.shape[0]] verts = deorient(verts, info['poses']) # # plot as scatter plot with plotly # import plotly.graph_objects as go # frame = 5 # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))]) # fig.show() writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True) np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask) def create_pose_shape(folder, out_folder): ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True) os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True) for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue info = loadInfo(os.path.join(folder, id, 'info.mat')) shape = np.concatenate([info['shape'], np.asarray([info['gender']])]).astype(np.float32) np.save(os.path.join(out_folder, 'pose', f'{id}.npy'), info['poses'].T.astype(np.float32)) np.save(os.path.join(out_folder, 'shape', f'{id}.npy'), shape) def create_mesh_static(folder, out_folder): ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True) for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): continue verts_static, _, _, _ = readOBJ(os.path.join(garment_dir, f'{garment}.obj')) reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat')) verts_static = (verts_static[reg_data['F'], :] * reg_data['W'][:, :, None]).sum(axis=1).astype(np.float32) # rotate 90 degrees around x axis verts_static = np.matmul(verts_static, np.array([[1,0,0],[0,0,-1],[0,1,0]])) # fig = go.Figure(data=[go.Scatter3d(x=verts_static[:,0], y=verts_static[:,2], z=verts_static[:,1], mode='markers', marker=dict(size=1))]) # fig.update_layout(scene_aspectmode='cube') # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() np.save(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'), verts_static.astype(np.float32)) def create_mesh_body_posed_no_orient(folder, out_folder): ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True) for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) pose, shape, gender = info['poses'], info['shape'], info['gender'] # if pose.shape[0] == 72: pose = np.transpose(pose, [1, 0]).astype(np.float32) pose[:, :3] = 0 beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0]) T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0] T = T.cpu().numpy() writePC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), T, True) # verts = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'] # # plot as scatter plot with plotly # import plotly.graph_objects as go # frame = 5 # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))]) # fig.show() pass def calculate_normals(vertices, faces): normals = [] for i in range(vertices.shape[0]): mesh = tm.Trimesh(vertices=vertices[i], faces=faces, process=False, validate=False) normals.append(mesh.vertex_normals) return np.stack(normals) def create_mesh_normals(folder, out_folder): ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True) # F = np.load('smpl/smpl_faces.npy') # F_skirt = np.load('smpl/smpl_faces_skirt.npy') # ids = ['00438'] for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] info = sio.loadmat(os.path.join(folder, id, 'info.mat')) for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): continue verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] verts = deorient(verts, info['poses'][:, :verts.shape[0]]) _, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj')) faces = quads2tris(faces) normals = calculate_normals(verts, faces.astype(np.int32)) reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat')) normals = (normals[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) writePC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), normals, True) # normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V'] # colors = np.clip(np.abs(normals), 0, 1) # verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) # # plot as scatter plot with plotly # import plotly.graph_objects as go # frame = 5 # fig = go.Figure(data=[go.Scatter3d( # x=verts[frame,:,0], # y=verts[frame,:,2], # z=verts[frame,:,1], # mode='markers', # marker=dict( # size=1, # color=normals[frame,:,:], # colorscale='Viridis', # opacity=0.8 # ) # )]) # moved_verts = verts + 0.05 * normals # # Create the scatter plot for the moved vertices # fig.add_trace(go.Scatter3d( # x=moved_verts[frame,:,0], # y=moved_verts[frame,:,2], # z=moved_verts[frame,:,1], # mode='markers', # marker=dict( # size=1, # color='red', # opacity=0.8 # ) # )) # fig.show() # # create a plotly mesh plot for verts and faces # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) # fig.show() pass def create_mesh_normals_cape(folder, out_folder): frames = os.listdir(os.path.join(folder, 'mesh_posed')) os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True) F = np.load('smpl/smpl_faces_orig.npy') # F_skirt = np.load('smpl/smpl_faces_skirt.npy') for frame in tqdm(frames): id, frame_nr = frame.split('.')[:2] verts = np.load(os.path.join(folder, 'mesh_posed', f'{id}.{frame_nr}.npz'))['mesh_posed'] normals = calculate_normals(verts[None,:,:], F)[0] np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}.{frame_nr}.npz'), mesh_normals=normals) # normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V'] # colors = np.clip(np.abs(normals), 0, 1) # verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2) # plot as scatter plot with plotly # import plotly.graph_objects as go # frame = 5 # fig = go.Figure(data=[go.Scatter3d( # x=verts[:,0], # y=verts[:,2], # z=verts[:,1], # mode='markers', # marker=dict( # size=1, # color=normals[:,:], # colorscale='Viridis', # opacity=0.8 # ) # )]) # moved_verts = verts + 0.05 * normals # # Create the scatter plot for the moved vertices # fig.add_trace(go.Scatter3d( # x=moved_verts[:,0], # y=moved_verts[:,2], # z=moved_verts[:,1], # mode='markers', # marker=dict( # size=1, # color='red', # opacity=0.8 # ) # )) # fig.show() # # create a plotly mesh plot for verts and faces # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) # fig.show() pass def verify_registration_ans_mask(folder, out_folder): ids = sorted(os.listdir(folder)) for id in ids: garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] for garment in garments: vertices = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'] mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) if mask.sum() == vertices.shape[1]: print(id, garment) pass def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1): ids = sorted(os.listdir(folder))[start:end] os.makedirs(os.path.join(out_folder, 'mesh_unposed_no_shape'), exist_ok=True) F = np.load('smpl/smpl_faces.npy') # F_skirt = np.load('smpl/smpl_faces_skirt.npy') for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] # info = sio.loadmat(os.path.join(folder, id, 'info.mat')) pose = np.load(os.path.join(out_folder, 'pose', f'{id}.npy')) shape = np.tile(np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:], [pose.shape[0],1]) pose = torch.tensor(pose, dtype=torch.float32, device='cuda') shape = torch.tensor(shape, dtype=torch.float32, device='cuda') smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0]) for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): continue mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy')) mesh_posed = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'] mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) # mask_idx = np.where(mask)[0] pose[:,:3] = 0 # pose = pose[:10] # shape = shape[:10] # mesh_posed = mesh_posed[:10] unposed = np.zeros((14475, 3), dtype=np.float32) unposed[mask] = mesh_static unposed = torch.tensor(np.tile(unposed[None,:,:], [mesh_posed.shape[0],1,1]), dtype=torch.float32, device='cuda') unposed.requires_grad = True mesh_posed = torch.tensor(mesh_posed, dtype=torch.float32, device='cuda') optimizer = torch.optim.Adam([unposed], lr=0.01) # verts_1 = unposed[0].cpu().detach().numpy()[mask] # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) # fig = go.Figure(data=[go.Mesh3d(x=verts_1[:,0], y=verts_1[:,2], z=verts_1[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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() # verts_2 = mesh_posed[0].cpu().detach().numpy() # fig = go.Figure(data=[go.Mesh3d(x=verts_2[:,0], y=verts_2[:,2], z=verts_2[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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() # zero_body = smpl(torch.zeros_like(pose), torch.zeros_like(shape))[0] # verts_3 = smpl(pose, shape, zero_body)[0][0].cpu().detach().numpy() # # verts_3 = zero_body[0].cpu().detach().numpy() # faces_3 = F # fig = go.Figure(data=[go.Mesh3d(x=verts_3[:,0], y=verts_3[:,2], z=verts_3[:,1], i=faces_3[:,0], j=faces_3[:,1], k=faces_3[:,2], color='lightpink', opacity=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() for i in range(5000): optimizer.zero_grad() pred = smpl(pose, shape, unposed)[0] loss : torch.Tensor = ((pred[:,mask,:] - mesh_posed)**2).sum() # Backward pass loss.backward() optimizer.step() with torch.inference_mode(): euclidean_error = torch.sqrt(((pred[:,mask,:] - mesh_posed)**2).sum(-1)).mean() # Print the loss if i % 50 == 0: print(f"it: {i} loss: {loss:.8f} euclidean_error: {euclidean_error:.8f}") if euclidean_error < 0.0000001: break if euclidean_error < 0.0000005 and i > 2000: break if euclidean_error < 0.000001 and i > 4000: break unposed = unposed.detach().cpu().numpy() gender = int(shape[0].cpu().numpy()[-1] > 0) shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]]) shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:] unshaped = (unposed - shape_offsets)[:,mask,:] writePC2(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True) # verts = unposed[:,mask,:] # verts = unshaped # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask) # frame = 9 # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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() # verts = pred.detach().cpu().numpy()[:,mask,:] # frame = 9 # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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 F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda') F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda') F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [1,1,1]), device='cuda') F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [1,1,1]), device='cuda') from scipy.spatial import KDTree from pytorch3d.structures import Meshes from torch_cluster import knn def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor): # flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function batch_size = V.shape[0] vertices_count = V.shape[1] batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(vertices_count) batch_y = torch.arange(batch_size, device=V.get_device()).repeat_interleave(14475) V = V.reshape(-1, V.shape[-1]) B = B.reshape(-1, B.shape[-1]) idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_y, batch_y=batch_x)[1] idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count) return idx def collision_loss(V: torch.Tensor, B: torch.Tensor, thr: float = .004): F_norm_local = F_norm mesh = Meshes(B, F_norm_local) B_vn = torch.stack(mesh.verts_normals_list()) # compute correspondences idx = _nearest_neighbour(V, B) # compute loss D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])]) 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) thresholded = torch.minimum(dot - thr, torch.zeros_like(dot)) loss = torch.sum(thresholded ** 2, axis=-1) return loss def area_loss(pred, gt, faces): """ Calculates the area loss between predicted and ground truth meshes. Args: pred (torch.Tensor): Predicted mesh vertices of shape (batch_size, num_vertices, 3). gt (torch.Tensor): Ground truth mesh vertices of shape (batch_size, num_vertices, 3). faces (torch.Tensor): Mesh faces of shape (num_faces, 3). Returns: torch.Tensor: Area loss between predicted and ground truth meshes. """ # Calculate the area of each face for predicted and ground truth meshes pred_faces = pred[:, faces, :] gt_faces = gt[:, faces, :] pred_edges1 = pred_faces[:, :, 1, :] - pred_faces[:, :, 0, :] pred_edges2 = pred_faces[:, :, 2, :] - pred_faces[:, :, 0, :] gt_edges1 = gt_faces[:, :, 1, :] - gt_faces[:, :, 0, :] gt_edges2 = gt_faces[:, :, 2, :] - gt_faces[:, :, 0, :] pred_cross = torch.cross(pred_edges1, pred_edges2, dim=2) gt_cross = torch.cross(gt_edges1, gt_edges2, dim=2) pred_area = 0.5 * torch.norm(pred_cross, dim=2) gt_area = 0.5 * torch.norm(gt_cross, dim=2) # Calculate the area loss area_loss = torch.mean(torch.abs(pred_area - gt_area)) return area_loss def create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=0, end=1): ids = sorted(os.listdir(folder))[start:end] os.makedirs(os.path.join(out_folder, 'mesh_static_unposed_no_shape'), exist_ok=True) F = np.load('smpl/smpl_faces.npy') # F_skirt = np.load('smpl/smpl_faces_skirt.npy') ids = ['00005'] for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] shape = np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:] shape = torch.tensor(shape, dtype=torch.float32, device='cuda') smpl = SMPL(path='smpl', device='cuda', batch_size=1) pose = np.zeros((24, 3), np.float32) pose[1,2] = .15 pose[2,2] = -.15 pose = torch.tensor(pose.flatten()[None,:], dtype=torch.float32, device='cuda') for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): continue mesh_static, faces, _, _ = readOBJ(os.path.join(folder, id, f'{garment}.obj')) faces = quads2tris(faces) mesh_static = np.matmul(mesh_static, np.array([[1,0,0],[0,0,-1],[0,1,0]])) unposed = mesh_static unposed = torch.tensor(unposed[None,:,:], dtype=torch.float32, device='cuda') unposed.requires_grad = True mesh_static = torch.tensor(mesh_static[None,:,:], dtype=torch.float32, device='cuda') optimizer = torch.optim.Adam([unposed], lr=0.0001) for i in range(10000): optimizer.zero_grad() body = smpl(pose, shape)[0] c_loss = collision_loss(unposed, body) a_loss = area_loss(unposed, mesh_static, faces) loss : torch.Tensor = c_loss + a_loss # Backward pass loss.backward() optimizer.step() pose = pose / 1.0005 # Print the loss if i % 50 == 0: pose_sum = np.abs(pose.detach().cpu().numpy()).sum().item() loss_str = loss.detach().cpu().numpy()[0] c_loss_str = c_loss.detach().cpu().numpy()[0] a_loss_str = a_loss.detach().cpu().numpy() print(f"it: {i} loss: {loss_str:.10f} c_loss: {c_loss_str:.10f}, a_loss: {a_loss_str:.8f}, pose: {pose_sum:.8f}") unposed = unposed.detach().cpu().numpy() body = body.detach().cpu().numpy() gender = int(shape[0].cpu().numpy()[-1] > 0) shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]]) shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:] unshaped = (unposed - shape_offsets) writePC2(os.path.join(out_folder, 'mesh_static_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True) verts = unposed # verts = unshaped frame = 0 fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) # add body mesh fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,2], z=body[frame,:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=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() # verts = pred.detach().cpu().numpy()[:,mask,:] # frame = 9 # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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 def create_frames_file(folder, out_folder, split=0.9): ids = os.listdir(folder) frames_non_skirt = [] frames_skirt = [] ids_non_skirt = set() ids_skirt = set() for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] frames = info['poses'].shape[1] assert info['poses'].shape[0] == 72 for garment in garments: for i in range(frames): if garment == 'Skirt' or garment == 'Dress': frames_skirt.append(Frame(id, i, garment)) ids_skirt.add(id) else: frames_non_skirt.append(Frame(id, i, garment)) ids_non_skirt.add(id) nr_train_non_skirt = int(len(ids_non_skirt) * split) nr_train_skirt = int(len(ids_skirt) * split) ids_skirt = list(ids_skirt) ids_non_skirt = list(ids_non_skirt) with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f: for frame in frames_non_skirt: if frame.id in ids_non_skirt[:nr_train_non_skirt]: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f: for frame in frames_skirt: if frame.id in ids_skirt[:nr_train_skirt]: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f: for frame in frames_non_skirt: if frame.id in ids_non_skirt[nr_train_non_skirt:]: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f: for frame in frames_skirt: if frame.id in ids_skirt[nr_train_skirt:]: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') def create_frames_file_test(folder, out_folder): ids = os.listdir(folder) frames_non_skirt = [] frames_skirt = [] ids_non_skirt = set() ids_skirt = set() for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] frames = info['poses'].shape[1] assert info['poses'].shape[0] == 72 for garment in garments: for i in range(frames): if garment == 'Skirt' or garment == 'Dress': frames_skirt.append(Frame(id, i, garment)) ids_skirt.add(id) else: frames_non_skirt.append(Frame(id, i, garment)) ids_non_skirt.add(id) ids_skirt = list(ids_skirt) ids_non_skirt = list(ids_non_skirt) with open(os.path.join(out_folder, 'frames_test.txt'), 'w') as f: for frame in frames_non_skirt: if frame.id in ids_non_skirt: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') with open(os.path.join(out_folder, 'frames_test_skirt.txt'), 'w') as f: for frame in frames_skirt: if frame.id in ids_skirt: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1): ids = os.listdir(folder) frames_non_skirt = [] frames_skirt = [] ids_non_skirt = set() ids_skirt = set() for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue info = sio.loadmat(os.path.join(folder, id, 'info.mat')) garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] frames = info['poses'].shape[1] assert info['poses'].shape[0] == 72 for garment in garments: for i in range(frames): if garment == 'Skirt' or garment == 'Dress': frames_skirt.append(Frame(id, i, garment)) ids_skirt.add(id) else: frames_non_skirt.append(Frame(id, i, garment)) ids_non_skirt.add(id) nr_train_non_skirt = int(len(ids_non_skirt) * split * part) nr_train_skirt = int(len(ids_skirt) * split * part) ids_skirt = list(ids_skirt) ids_non_skirt = list(ids_non_skirt) with open(os.path.join(out_folder, f'frames_train_{int(part*100)}p.txt'), 'w') as f: for frame in frames_non_skirt: if frame.id in ids_non_skirt[:nr_train_non_skirt]: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') with open(os.path.join(out_folder, f'frames_train_skirt_{int(part*100)}p.txt'), 'w') as f: for frame in frames_skirt: if frame.id in ids_skirt[:nr_train_skirt]: f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n') def check_irregular_poses(folder, folder_reg, out_folder): for i in range(100, 200, 5): frame: Frame = Frame('01512', i, 'Dress') mesh_posed = readPC2Frame(os.path.join(out_folder, 'mesh_posed_no_orient', f'{frame.id}_{frame.garment}.pc16'), frame.frame_nr, True) # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy') mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy') mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V'] mesh_body_posed = mesh_body_posed[::2][frame.frame_nr] F = np.load('smpl/smpl_faces_skirt.npy') _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask) fig = go.Figure(data=[go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)]) fig.add_trace(go.Mesh3d(x=mesh_body_posed[:,0], y=mesh_body_posed[:,2], z=mesh_body_posed[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=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 def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'): from custom_utils.custom_utils import UV2Mesh import plotly.express as ex import plotly.graph_objects as go frames = [] with open(os.path.join(out_folder, frames_file), 'r') as f: for line in f: line = line.strip() id, frame_nr, garment = line.split(' ') frames.append(Frame(id, int(frame_nr), garment)) uv_sum = np.zeros((512, 256, 3), dtype=np.float32) mask_sum = np.zeros((512, 256), dtype=np.float32) for frame in tqdm(frames): uv = np.load(os.path.join(out_folder, 'uv_unposed_no_shape', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed'] mask = np.load(os.path.join(out_folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz'))['uv_mask'] uv_sum += uv mask_sum += mask mask_sum[mask_sum == 0] = 1 avg_garment = uv_sum / mask_sum[:, :, None] avg_mesh = UV2Mesh()(torch.from_numpy(avg_garment.transpose(2,0,1))[None,:,:,:]).numpy()[0] F = np.load('smpl/smpl_faces.npy') fig = go.Figure(data=[go.Mesh3d(x=avg_mesh[:,0], y=avg_mesh[:,2], z=avg_mesh[:,1], i=F[:,0], j=F[:,1], k=F[:,2], color='lightpink', opacity=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() np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment) def remove_irregular_poses(folder, folder_reg, out_folder): ids = sorted(os.listdir(folder)) for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] info = loadInfo(os.path.join(folder, id, 'info.mat')) for garment in garments: verts_count = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'].shape[0] if info['poses'].shape[1] != verts_count: shutil.rmtree(os.path.join(folder, id)) if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')): os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')) if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')): os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')) if os.path.exists(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')): os.remove(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')) if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')): os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.npy')) if len(garments) == 2: if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')): os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')) if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')): os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')) if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.pc16')): os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.npy')) break # os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garments[1]}.npy')) def double_check_irregular_poses(folder, folder_reg, out_folder): # ids = sorted(os.listdir(folder)) # for id in tqdm(ids): # garment_dir = os.path.join(folder, id) # if not os.path.isdir(garment_dir): # continue # garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] # info = loadInfo(os.path.join(folder, id, 'info.mat')) # for garment in garments: # body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0] # verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0] # assert body_count == verts_count frames = sorted(os.listdir(out_folder + '/mesh_posed_no_orient')) frames = [(id.split('_')[0], id.split('_')[1].split('.')[0]) for id in frames] for frame in tqdm(frames): id = frame[0] garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garment = frame[1] body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0] verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0] assert body_count == verts_count def get_ids(): ids = [] with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train.txt'), 'r') as f: for line in f: ids.append(line.split()[0]) with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train_skirt.txt'), 'r') as f: for line in f: ids.append(line.split()[0]) ids = sorted(list(set(ids))) with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train_ids.txt'), 'w') as f: for id in ids: f.write(f'{id}\n') ids = [] with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test.txt'), 'r') as f: for line in f: ids.append(line.split()[0]) with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test_skirt.txt'), 'r') as f: for line in f: ids.append(line.split()[0]) ids = sorted(list(set(ids))) with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test_ids.txt'), 'w') as f: for id in ids: f.write(f'{id}\n') pass def generate_fixed_masks(folder, out_folder): ids = sorted(os.listdir(folder)) os.makedirs(os.path.join(out_folder, 'mesh_fixed'), exist_ok=True) F = np.load('smpl/smpl_faces.npy') F_skirt = np.load('smpl/smpl_faces_skirt.npy') import random # import trimesh import pyvista as pv random.shuffle(ids) for id in tqdm(ids): garment_dir = os.path.join(folder, id) if not os.path.isdir(garment_dir): continue garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')] for garment in garments: if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')): continue is_skirt = garment in ['Dress', 'Skirt'] mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy')) mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')) _, faces = _maskFaces(np.zeros((14475, 3)), F_skirt if is_skirt else F, mask) _, new_mask = remove_long_edges(mesh_static, faces.astype(np.int32), 0.2 if is_skirt else 0.1) new_mask_full = np.zeros((14475), dtype=bool) new_mask_full[mask] = new_mask _, new_faces = _maskFaces(np.zeros((14475, 3)), F_skirt if is_skirt else F, new_mask_full) mesh_new = mesh_static[new_mask] # fig = go.Figure(data=[go.Mesh3d(x=mesh_new[:,0], y=mesh_new[:,1], z=mesh_new[:,2], i=new_faces[:,0], j=new_faces[:,1], k=new_faces[:,2])]) # fig.update_layout(scene_aspectmode='cube') # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() flattened_faces = np.asarray([np.concatenate((np.array([3]), f)) for f in new_faces.astype(np.int32)]).flatten() m = pv.PolyData(mesh_new, flattened_faces) m.fill_holes(0.03, inplace=True, progress_bar=False) # m = trimesh.Trimesh(vertices=mesh_new, faces=new_faces, process=False, validate=False) # trimesh.repair.fill_holes(m) mesh_new = np.asarray(m.points) new_faces = np.asarray(m.faces).reshape([-1, 4])[:, 1:] m = tm.Trimesh(vertices=mesh_new, faces=new_faces, process=False, validate=False) tm.repair.fix_winding(m) mesh_new = np.asarray(m.vertices) new_faces = np.asarray(m.faces) # fig = go.Figure(data=[go.Mesh3d(x=mesh_new[:,0], y=mesh_new[:,1], z=mesh_new[:,2], i=new_faces[:,0], j=new_faces[:,1], k=new_faces[:,2])]) # fig.update_layout(scene_aspectmode='cube') # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() # fig = go.Figure(data=[go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,1], z=mesh_static[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2])]) # fig.update_layout(scene_aspectmode='cube') # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) # fig.show() # print((~new_mask).sum()) np.savez_compressed(os.path.join(out_folder, 'mesh_fixed', f'{id}_{garment}.npz'), mask=new_mask_full, mas_convert=new_mask, faces=new_faces) pass if __name__ == '__main__': folder = '/data-net/dataset/cloth3d/test' folder_reg = '/data-local/127/cloth3d_registration/test' out_folder = '/data-local/127_2/cloth3d_processed_test' # folder = '/data-net/dataset/cloth3d/train' # folder_reg = '/data-local/127/cloth3d_registration/train' # out_folder = '/data-local/127/cloth3d_processed' smpl = SMPL(path='smpl', device='cuda') # clean_dataset(folder) # generate_frames_file(folder, 'frames_all.txt') # split_data(folder, 0.9) # create_extras_file(folder) # check_irregular_poses(folder, folder_reg, out_folder) # remove_irregular_poses(folder, folder_reg, out_folder) # double_check_irregular_poses(folder, folder_reg, out_folder) create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder) # verify_registration_ans_mask(folder, out_folder) # create_pose_shape(folder, out_folder) # create_mesh_static(folder, out_folder) # create_mesh_body_posed_no_orient(folder, out_folder) # create_mesh_normals(folder, out_folder) # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end) # create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=args.start, end=args.end) # calculate_average(folder, out_folder) # create_frames_file(folder, out_folder) # get_ids() # create_frames_file_test(folder, out_folder) # create_frames_file_partial(folder, out_folder, part=0.5) # folder = '/data-local/127_2/cape_processed/test/female' # create_mesh_normals_cape(folder, None) # generate_fixed_masks(folder, out_folder)