luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

all preprocessing in one place

Browse Source
This commit is contained in:
luni
2023-06-30 00:12:59 +02:00
parent c7590519b3
commit ff9771a415
1 changed files with 395 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/preprocessing.py
+395 -5
View File
@@ -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)
# 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)