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 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 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) 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 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 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) # 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)