From ff9771a41578f8e38dbffcc89b69fd849fe1b8b3 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Jun 2023 00:12:59 +0200 Subject: [PATCH] all preprocessing in one place --- data/preprocessing.py | 400 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 395 insertions(+), 5 deletions(-) diff --git a/data/preprocessing.py b/data/preprocessing.py index 53cd1c5..098285c 100644 --- a/data/preprocessing.py +++ b/data/preprocessing.py @@ -1,22 +1,39 @@ 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 +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 @@ -132,12 +149,385 @@ def create_extras_file(folder): 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 = '/home/workspace/data/dataset/cloth3d/train' + + 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) \ No newline at end of file + # 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) \ No newline at end of file