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") # parser.add_argument("-n", "--name", help = "end of range") args = parser.parse_args() os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu 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 from data.utils import Frame from smpl.smpl_torch import SMPL from custom_utils.custom_utils import _maskFaces def clean_dataset(folder): ids = os.listdir(folder) for id in tqdm(ids): if not os.path.isdir(os.path.join(folder, id)): continue files = os.listdir(os.path.join(folder, id)) for file in files: if file.startswith('glmarks') or file.find('nricp') != -1: os.remove(os.path.join(folder, id, file)) def generate_frames_file(folder, file_name): ids = os.listdir(folder) frames = [] for id in tqdm(ids): if not os.path.isdir(os.path.join(folder, id)): continue garments = [f.replace('.obj','') for f in os.listdir(os.path.join(folder, id)) if f.endswith('.obj')] info = loadInfo(os.path.join(folder, id, 'info.mat')) for garment in garments: for i in range(info['poses'].shape[1]): frames.append(Frame(id, i, garment)) with open(os.path.join(folder, file_name), 'w') as f: for i, frame in enumerate(frames): f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n') def split_data(folder, ratio): with open(os.path.join(folder, 'frames_all.txt'), "r") as f: frames = f.readlines() ids = sorted(list(set([frame.split()[0] for frame in frames]))) train_last_id = round(len(ids) * ratio) train_frames = [] valid_frames = [] for frame in tqdm(frames): id = frame.split()[0] if id in ids[:train_last_id]: train_frames.append(frame) else: valid_frames.append(frame) with open(os.path.join(folder, f'frames_train_{ratio}.txt'), 'w') as f: f.writelines(train_frames) with open(os.path.join(folder, f'frames_valid_{round(1-ratio,1)}.txt'), 'w') as f: f.writelines(valid_frames) def genereated_registered_and_mask(folder, is_skirt = False): ids = 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('.pc16','') for f in os.listdir(garment_dir) if f.find('_') == -1 and not f.endswith('.mat') and not f.endswith('.obj') and not f.startswith('Combined')] assert len(garments) == 1 if os.path.exists(os.path.join(garment_dir, garments[0] + '_registered.pc16')): continue garment = garments[0] verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'] map = sio.loadmat(os.path.join(garment_dir, garment + '_map.mat'))['map'].reshape((-1)) verts = verts[:, map, :] # ---------------> This is the garment after registration raw_mask = sio.loadmat(os.path.join(garment_dir, garment + '_mask.mat'))['mask'].reshape((-1)) mask = raw_mask > 0 # boundaries = raw_mask == 2 writePC2(os.path.join(garment_dir, garment + '_registered.pc16'), verts, True) np.save(os.path.join(garment_dir, garment + '_mask.npy'), mask) 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 = 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) 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'] 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 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 = 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 = 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 = 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'] pose = np.transpose(pose, [1, 0]).astype(np.float32) pose[:, :3] = 0 beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32) 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 = 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') 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']) _, 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 verify_registration_ans_mask(folder, out_folder): ids = 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') 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 if __name__ == '__main__': 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) # 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)