771 lines
30 KiB
Python
771 lines
30 KiB
Python
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 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 = 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']
|
|
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 = 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')
|
|
|
|
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_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):
|
|
frame: Frame = Frame('01512', 50, '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[int(mesh_body_posed.shape[0]/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)
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
|
|
# folder = '/data-net/dataset/cloth3d/test'
|
|
# folder_reg = '/data-local/127/cloth3d_registration/test'
|
|
# out_folder = '/data-local/127/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)
|
|
|
|
# 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)
|
|
# 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) |