208 lines
6.5 KiB
Python
208 lines
6.5 KiB
Python
import os
|
|
import sys
|
|
import random
|
|
import numpy as np
|
|
import plotly.graph_objects as go
|
|
import torch
|
|
from torch.nn.functional import mse_loss
|
|
from pytorch3d.structures import Meshes
|
|
from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
from data.cloth3d_io import (
|
|
read_garment_mask,
|
|
read_garment_template,
|
|
read_posed_garment_frame,
|
|
loadInfo,
|
|
read_posed_body_frame,
|
|
read_posed_body,
|
|
read_posed_garment,
|
|
read_curvatures,
|
|
read_curvature_frame,
|
|
writeOBJ,
|
|
readOBJ,
|
|
quads2tris
|
|
)
|
|
from smpl.smpl_torch import SMPL
|
|
# from custom_utils.custom_utils import _maskFaces
|
|
|
|
smpl = SMPL('cuda', 'smpl')
|
|
|
|
|
|
def _maskFaces(V, F, mask, return_unfixed=False):
|
|
V_mask = mask
|
|
F_mask = np.take(V_mask, F.reshape(-1), axis=0)
|
|
F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
|
|
|
|
V_ind = np.zeros(V_mask.shape)
|
|
V_ind[V_mask] = np.arange(V_mask.sum())
|
|
F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
|
|
V_new = V[V_mask, :]
|
|
|
|
if return_unfixed:
|
|
return V_new, F_new, F[F_mask, :]
|
|
else:
|
|
return V_new, F_new
|
|
|
|
def get_frames(folder, frames_file):
|
|
frames = []
|
|
with open(os.path.join(folder, frames_file), 'r') as f:
|
|
for line in f:
|
|
frames.append(line.strip())
|
|
return frames
|
|
|
|
|
|
def plot(mesh, faces):
|
|
fig = go.Figure()
|
|
|
|
fig.add_trace(go.Mesh3d(
|
|
x=mesh[:,0],
|
|
y=mesh[:,1],
|
|
z=mesh[:,2],
|
|
i = faces[:,0],
|
|
j = faces[:,1],
|
|
k = faces[:,2]
|
|
),)
|
|
|
|
fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
|
|
)
|
|
helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
|
|
yaxis = dict(nticks=4, range=[-1,1],),
|
|
zaxis = dict(nticks=4, range=[-1,1],),)
|
|
fig.update_layout(scene = helper,
|
|
)
|
|
|
|
fig.show()
|
|
|
|
|
|
def build_lookup_table(repeated):
|
|
# iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
|
|
|
|
lookup = {}
|
|
for i in range(len(repeated)):
|
|
for j in range(len(repeated[i])):
|
|
if len(repeated[i]) == 1:
|
|
continue
|
|
lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
|
|
return lookup
|
|
|
|
|
|
from data.cloth3d_io import readPC2Frame
|
|
import scipy.io as sio
|
|
import scipy
|
|
def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
|
|
pc16_path = os.path.join(folder, id, garment + '.pc16')
|
|
verts = readPC2Frame(pc16_path, frame, True)
|
|
|
|
map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
|
|
|
|
if deorient:
|
|
R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
|
|
# R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
|
|
verts = R.apply(verts)
|
|
|
|
return verts, map
|
|
|
|
def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
|
|
pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
|
|
verts = readPC2Frame(pc16_path, frame, True)
|
|
|
|
if deorient:
|
|
R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
|
|
# R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
|
|
verts = R.apply(verts)
|
|
|
|
return verts
|
|
|
|
def register(folder, frame):
|
|
id, frame_nr, garment = frame.split()
|
|
frame_nr = int(frame_nr)
|
|
info = loadInfo(os.path.join(folder, id, 'info.mat'))
|
|
raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
|
|
mask = raw_mask > 0
|
|
mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
|
|
body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
|
|
template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
|
|
faces = quads2tris(faces)
|
|
pose = torch.zeros(72)
|
|
shape = torch.zeros(11)
|
|
shape[-1] = 1
|
|
# for star pose
|
|
pose[:,5] = 3.1415/4
|
|
pose[:,8] = -3.1415/4
|
|
star_body = smpl(pose.unsqueeze(0), shape.unsqueeze(0))[0][0].cpu().numpy()
|
|
|
|
is_skirt = (garment == 'Dress' or garment == 'Skirt')
|
|
skirt_mod = '_skirt' if is_skirt else ''
|
|
F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
|
|
|
|
plot(template, faces)
|
|
# v = mesh_posed[mask,:]
|
|
|
|
# Define the meshes to optimize
|
|
mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
|
|
mesh2 = Meshes(verts=star_body, faces=F[None,:,:])
|
|
|
|
# Define the optimizer and learning rate
|
|
optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.01)
|
|
|
|
# Define the loss functions
|
|
def laplacian_smoothing_loss(mesh):
|
|
return mesh_laplacian_smoothing(mesh, method="uniform")
|
|
|
|
def edge_loss(mesh1, mesh2):
|
|
return mesh_edge_loss(mesh1, mesh2)
|
|
|
|
def normal_consistency_loss(mesh1, mesh2):
|
|
return mesh_normal_consistency(mesh1, mesh2)
|
|
|
|
# Optimization loop
|
|
for i in range(100):
|
|
optimizer.zero_grad()
|
|
|
|
# Forward pass
|
|
mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
|
|
loss = chamfer_distance(mesh1_out, mesh2)
|
|
loss += laplacian_smoothing_loss(mesh1_out)
|
|
loss += edge_loss(mesh1_out)
|
|
loss += normal_consistency_loss(mesh1_out)
|
|
|
|
# Backward pass
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# Print the loss
|
|
print(f"Loss: {loss.item()}")
|
|
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
folder = '/data-net/dataset/cloth3d/train'
|
|
# folder = r'D:\Work\Dataset\cloth3d\train'
|
|
# frames = [f'00761_{i}' for i in range(300)]
|
|
# folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
|
|
frames = get_frames(folder, 'frames_all.txt')
|
|
root = '/'.join(folder.split('/')[:-1])
|
|
type = folder.split('/')[-1]
|
|
|
|
use_no_orient = False
|
|
no_shape = False
|
|
is_skirt = False
|
|
|
|
skirt_modifier = '_skirt' if is_skirt else ''
|
|
orient_mod = '_no_orient' if use_no_orient else ''
|
|
shape_mod = '_no_shape' if no_shape else ''
|
|
|
|
# mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
|
|
# faces = np.load('/data-net/dataset/00761/faces/00761.npy')
|
|
|
|
for frame in frames:
|
|
frame = random.choice(frames)
|
|
# frame = '00761 0 Tshirt'
|
|
frame = '01893 256 Jumpsuit'
|
|
|
|
|
|
register(folder, frame)
|
|
|
|
pass
|