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

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/preprocessing.py
+184 -10
View File
@@ -12,6 +12,7 @@ args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import math
import shutil
@@ -76,6 +77,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_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)
ids = ['00021']
for id in tqdm(ids):
# id = '01467'
@@ -87,8 +90,8 @@ def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
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
# 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]:
@@ -488,6 +491,176 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
pass
F_coll = torch.tensor(np.load("smpl/smpl_faces.npy"), device='cuda')
F_coll_skirt = torch.tensor(np.load("smpl/smpl_faces_skirt.npy"), device='cuda')
F_norm = torch.tensor(np.tile(np.load("smpl/smpl_faces.npy")[None,:,:], [1,1,1]), device='cuda')
F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[None,:,:], [1,1,1]), device='cuda')
from scipy.spatial import KDTree
from pytorch3d.structures import Meshes
from torch_cluster import knn
def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
vertices_count = V.shape[1]
batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(vertices_count)
batch_y = torch.arange(batch_size, device=V.get_device()).repeat_interleave(14475)
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(B.to(V.dtype), V, k=1, batch_x=batch_y, batch_y=batch_x)[1]
idx = (idx - batch_x * vertices_count).reshape(batch_size, vertices_count)
return idx
def collision_loss(V: torch.Tensor, B: torch.Tensor, thr: float = .004):
F_norm_local = F_norm
mesh = Meshes(B, F_norm_local)
B_vn = torch.stack(mesh.verts_normals_list())
# compute correspondences
idx = _nearest_neighbour(V, B)
# compute loss
D = V - torch.stack([B[i,idx[i],:] for i in range(V.shape[0])])
b_vn = torch.stack([B_vn[i,idx[i],:] for i in range(V.shape[0])])
dot = torch.einsum('abc,abc->ab', D, b_vn)
thresholded = torch.minimum(dot - thr, torch.zeros_like(dot))
loss = torch.sum(thresholded ** 2, axis=-1)
return loss
def area_loss(pred, gt, faces):
"""
Calculates the area loss between predicted and ground truth meshes.
Args:
pred (torch.Tensor): Predicted mesh vertices of shape (batch_size, num_vertices, 3).
gt (torch.Tensor): Ground truth mesh vertices of shape (batch_size, num_vertices, 3).
faces (torch.Tensor): Mesh faces of shape (num_faces, 3).
Returns:
torch.Tensor: Area loss between predicted and ground truth meshes.
"""
# Calculate the area of each face for predicted and ground truth meshes
pred_faces = pred[:, faces, :]
gt_faces = gt[:, faces, :]
pred_edges1 = pred_faces[:, :, 1, :] - pred_faces[:, :, 0, :]
pred_edges2 = pred_faces[:, :, 2, :] - pred_faces[:, :, 0, :]
gt_edges1 = gt_faces[:, :, 1, :] - gt_faces[:, :, 0, :]
gt_edges2 = gt_faces[:, :, 2, :] - gt_faces[:, :, 0, :]
pred_cross = torch.cross(pred_edges1, pred_edges2, dim=2)
gt_cross = torch.cross(gt_edges1, gt_edges2, dim=2)
pred_area = 0.5 * torch.norm(pred_cross, dim=2)
gt_area = 0.5 * torch.norm(gt_cross, dim=2)
# Calculate the area loss
area_loss = torch.mean(torch.abs(pred_area - gt_area))
return area_loss
def create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=0, end=1):
ids = sorted(os.listdir(folder))[start:end]
os.makedirs(os.path.join(out_folder, 'mesh_static_unposed_no_shape'), exist_ok=True)
F = np.load('smpl/smpl_faces.npy')
# F_skirt = np.load('smpl/smpl_faces_skirt.npy')
ids = ['00005']
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')]
shape = np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:]
shape = torch.tensor(shape, dtype=torch.float32, device='cuda')
smpl = SMPL(path='smpl', device='cuda', batch_size=1)
pose = np.zeros((24, 3), np.float32)
pose[1,2] = .15
pose[2,2] = -.15
pose = torch.tensor(pose.flatten()[None,:], 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, faces, _, _ = readOBJ(os.path.join(folder, id, f'{garment}.obj'))
faces = quads2tris(faces)
mesh_static = np.matmul(mesh_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
unposed = mesh_static
unposed = torch.tensor(unposed[None,:,:], dtype=torch.float32, device='cuda')
unposed.requires_grad = True
mesh_static = torch.tensor(mesh_static[None,:,:], dtype=torch.float32, device='cuda')
optimizer = torch.optim.Adam([unposed], lr=0.0001)
for i in range(10000):
optimizer.zero_grad()
body = smpl(pose, shape)[0]
c_loss = collision_loss(unposed, body)
a_loss = area_loss(unposed, mesh_static, faces)
loss : torch.Tensor = c_loss + a_loss
# Backward pass
loss.backward()
optimizer.step()
pose = pose / 1.0005
# Print the loss
if i % 50 == 0:
pose_sum = np.abs(pose.detach().cpu().numpy()).sum().item()
loss_str = loss.detach().cpu().numpy()[0]
c_loss_str = c_loss.detach().cpu().numpy()[0]
a_loss_str = a_loss.detach().cpu().numpy()
print(f"it: {i} loss: {loss_str:.10f} c_loss: {c_loss_str:.10f}, a_loss: {a_loss_str:.8f}, pose: {pose_sum:.8f}")
unposed = unposed.detach().cpu().numpy()
body = body.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)
writePC2(os.path.join(out_folder, 'mesh_static_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
verts = unposed
# verts = unshaped
frame = 0
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)])
# add body mesh
fig.add_trace(go.Mesh3d(x=body[frame,:,0], y=body[frame,:,2], z=body[frame,:,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()
# 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)
@@ -879,13 +1052,13 @@ def generate_fixed_masks(folder, out_folder):
if __name__ == '__main__':
# folder = '/data-net/dataset/cloth3d/test'
# folder_reg = '/data-local/127/cloth3d_registration/test'
# out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/test'
folder_reg = '/data-local/127/cloth3d_registration/test'
out_folder = '/data-local/127_2/cloth3d_processed_test'
folder = '/data-net/dataset/cloth3d/train'
folder_reg = '/data-local/127/cloth3d_registration/train'
out_folder = '/data-local/127/cloth3d_processed'
# 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')
@@ -898,13 +1071,14 @@ if __name__ == '__main__':
# remove_irregular_poses(folder, folder_reg, out_folder)
# double_check_irregular_poses(folder, folder_reg, out_folder)
# create_mesh_posed_no_orient_and_mask(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)
# create_mesh_unposed_no_shape_static(folder, folder_reg, out_folder, start=args.start, end=args.end)
# calculate_average(folder, out_folder)
# create_frames_file(folder, out_folder)
@@ -915,4 +1089,4 @@ if __name__ == '__main__':
# folder = '/data-local/127_2/cape_processed/test/female'
# create_mesh_normals_cape(folder, None)
generate_fixed_masks(folder, out_folder)
# generate_fixed_masks(folder, out_folder)