113 lines
4.2 KiB
Python
113 lines
4.2 KiB
Python
import os
|
|
import sys
|
|
import numpy as np
|
|
import trimesh as tm
|
|
import scipy.io as sio
|
|
import plotly.graph_objects as go
|
|
import plotly.express as px
|
|
from tqdm import tqdm
|
|
import torch
|
|
from scipy.spatial.transform import Rotation as R
|
|
import argparse
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
from data.cloth3d_io import readOBJ, readPC2
|
|
from data.utils import Frame
|
|
from smpl.smpl_torch import SMPL
|
|
|
|
|
|
def calculate_normals(vertices, mesh_file, poses):
|
|
_, faces, _, _ = readOBJ(mesh_file)
|
|
rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)
|
|
|
|
normals = []
|
|
for i in range(vertices.shape[0]):
|
|
r = R.from_matrix(rodr[i, 0, :, :])
|
|
# assert r.as_matrix() == rodr[frame, 0, :, :]
|
|
out = r.apply(vertices[i], inverse=True)
|
|
|
|
# create a plotly mesh plot of out and vertices [i]
|
|
# fig = go.Figure(data=[
|
|
# go.Mesh3d(
|
|
# x=out[:, 0], y=out[:, 2], z=out[:, 1],
|
|
# # i, j and k give the vertices of triangles
|
|
# i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
|
|
# name='y',
|
|
# showscale=False
|
|
# ),
|
|
# go.Mesh3d(
|
|
# x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
|
|
# # i, j and k give the vertices of triangles
|
|
# i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
|
|
# name='x',
|
|
# showscale=False
|
|
# )
|
|
# ])
|
|
# fig.show()
|
|
|
|
mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
|
|
normals.append(mesh.vertex_normals)
|
|
return np.stack(normals)
|
|
|
|
|
|
def get_frames(folder):
|
|
uv_static_folder = os.path.join(folder, 'pose_shape')
|
|
frames = sorted(os.listdir(uv_static_folder))
|
|
|
|
frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
|
|
frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
|
|
return frames
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("-f", "--folder", help = "end of range", default='/data-net/dataset/00761')
|
|
parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/data-net/dataset/cloth3d/train')
|
|
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
|
parser.add_argument("-e", "--end", help = "end of range", default=10)
|
|
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
|
args = parser.parse_args()
|
|
|
|
obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
|
|
folder = args.folder #'/data-local/127/train'
|
|
|
|
os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
|
|
frames = get_frames(folder)
|
|
smpl = SMPL('cpu', 'smpl')
|
|
|
|
ids = {frame.id for frame in frames}
|
|
ids = sorted(list(ids))[int(args.start):int(args.end)]
|
|
|
|
for id in tqdm(ids):
|
|
# find obj files in the folder
|
|
garment_obj_folder = os.path.join(obj_folder, id[:5])
|
|
garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]
|
|
|
|
if len(id) == 5:
|
|
garment = garments[0]
|
|
elif len(id) == 6:
|
|
garment = garments[1]
|
|
else:
|
|
raise ValueError('Invalid frame id')
|
|
|
|
mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
|
|
vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
|
|
poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
|
|
normals = calculate_normals(vertices, mesh_file, poses)
|
|
|
|
map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
|
|
mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
|
|
|
|
mesh_normals = normals[:,map,:]
|
|
mesh_normals = mesh_normals * mask[None,:, None]
|
|
|
|
for i in range(vertices.shape[0]):
|
|
# np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
|
|
np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])
|
|
|
|
# pass
|
|
|
|
# print(frame.id, frame.frame_nr)
|
|
|