Merge remote-tracking branch 'origin/main' into main
This commit is contained in:
+89
-49
@@ -140,7 +140,7 @@ class DatasetCloth3D(torch.utils.data.Dataset):
|
||||
self.resolutions = resolutions
|
||||
|
||||
self.transforms = tv.transforms.ToTensor()
|
||||
self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=True)
|
||||
self.frames = read_frames_file(os.path.join(folder, frames_file), shuffle=False)
|
||||
|
||||
self.static_cache = StaticCache(folder, signature, is_cape=is_cape, resolutions=resolutions, device=device)
|
||||
|
||||
@@ -208,7 +208,7 @@ class DatasetCloth3D(torch.utils.data.Dataset):
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
|
||||
if 'uv_body_posed' + res_string in self.signature:
|
||||
uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string, f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed']
|
||||
uv_body_posed_no_orient = np.load(os.path.join(self.folder, 'uv_body_posed_no_orient' + res_string[1:], f'{frame.id}_{frame.frame_nr}.npz')); uv_body_posed_no_orient = uv_body_posed_no_orient['uv_body_posed']
|
||||
item['uv_body_posed' + res_string] = uv_body_posed_no_orient
|
||||
|
||||
if 'uv_posed' + res_string in self.signature:
|
||||
@@ -244,6 +244,9 @@ if __name__ == '__main__':
|
||||
import plotly.express as ex
|
||||
from plotly.subplots import make_subplots
|
||||
from custom_utils.custom_utils import _maskFaces
|
||||
from losses.losses import collision_loss, normal_loss
|
||||
from smpl.smpl_torch import SMPL
|
||||
from torch.profiler import profile, record_function, ProfilerActivity
|
||||
|
||||
F = np.load('smpl/smpl_faces.npy')
|
||||
F_skirt = np.load('smpl/smpl_faces_skirt.npy')
|
||||
@@ -267,61 +270,98 @@ if __name__ == '__main__':
|
||||
|
||||
dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
|
||||
train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
|
||||
smpl = SMPL(orig_smpl=False, path='smpl', device='cuda')
|
||||
|
||||
# create a profiler object
|
||||
profiler = torch.profiler.profile(
|
||||
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
|
||||
schedule=torch.profiler.schedule(wait=2, warmup=1, active=3, repeat=2),
|
||||
on_trace_ready=torch.profiler.tensorboard_trace_handler("log_dir"),
|
||||
with_stack=True,
|
||||
profile_memory=True,
|
||||
)
|
||||
|
||||
|
||||
step = 0
|
||||
for batch in train_dataloader:
|
||||
batch
|
||||
b = 0
|
||||
mesh_mask = batch['mesh_mask'][b].cpu().numpy()
|
||||
faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
|
||||
print(batch['id'][0])
|
||||
with profiler:
|
||||
m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0]
|
||||
profiler.step()
|
||||
step += 1
|
||||
if step > 25:
|
||||
break
|
||||
|
||||
mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask]
|
||||
mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask]
|
||||
mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask]
|
||||
mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy()
|
||||
|
||||
fig = make_subplots(rows=2, cols=2,
|
||||
specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}],
|
||||
[{'type': 'mesh3d'}, {'type': 'mesh3d'}]])
|
||||
|
||||
fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2)
|
||||
fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=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], opacity=1), row=2, col=2)
|
||||
|
||||
fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'),
|
||||
scene2=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'),
|
||||
scene3=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'),
|
||||
scene4=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
yaxis=dict(range=[-1, 1], autorange=False),
|
||||
zaxis=dict(range=[-1, 1], autorange=False),
|
||||
aspectmode='cube'))
|
||||
|
||||
fig.show()
|
||||
profiler.stop()
|
||||
print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25))
|
||||
print(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=25))
|
||||
print(profiler.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=10))
|
||||
|
||||
|
||||
uv_mask = batch['uv_mask'][b].cpu().numpy()
|
||||
uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0)
|
||||
uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2
|
||||
uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2
|
||||
uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2
|
||||
uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2
|
||||
uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3
|
||||
# b = 0
|
||||
# mesh_mask = batch['mesh_mask'][b].cpu().numpy()
|
||||
# faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F
|
||||
# _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
|
||||
|
||||
img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2)
|
||||
img = np.transpose(img, (1, 2, 0))
|
||||
ex.imshow(img).show()
|
||||
# c = collision_loss(batch['mesh_posed'], batch['mesh_body_posed'], batch['mesh_mask'])
|
||||
# fig = go.Figure(data=[go.Mesh3d(x=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,0], y=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,1], z=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,2], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1),
|
||||
# go.Mesh3d(x=batch['mesh_body_posed'].cpu().numpy()[0,:,0], y=batch['mesh_body_posed'].cpu().numpy()[0,:,1], z=batch['mesh_body_posed'].cpu().numpy()[0,:,2], i=F[:,0], j=F[:,1], k=F[:,2], color='lightblue', opacity=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()
|
||||
|
||||
pass
|
||||
|
||||
|
||||
|
||||
# mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask]
|
||||
# mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask]
|
||||
# mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask]
|
||||
# mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy()
|
||||
|
||||
# fig = make_subplots(rows=2, cols=2,
|
||||
# specs=[[{'type': 'mesh3d'}, {'type': 'mesh3d'}],
|
||||
# [{'type': 'mesh3d'}, {'type': 'mesh3d'}]])
|
||||
|
||||
# fig.add_trace(go.Mesh3d(x=mesh_static[:,0], y=mesh_static[:,2], z=mesh_static[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=1)
|
||||
# fig.add_trace(go.Mesh3d(x=mesh_posed[:,0], y=mesh_posed[:,2], z=mesh_posed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=1, col=2)
|
||||
# fig.add_trace(go.Mesh3d(x=mesh_unposed[:,0], y=mesh_unposed[:,2], z=mesh_unposed[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], opacity=1), row=2, col=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], opacity=1), row=2, col=2)
|
||||
|
||||
# fig.update_layout(scene=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
# yaxis=dict(range=[-1, 1], autorange=False),
|
||||
# zaxis=dict(range=[-1, 1], autorange=False),
|
||||
# aspectmode='cube'),
|
||||
# scene2=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
# yaxis=dict(range=[-1, 1], autorange=False),
|
||||
# zaxis=dict(range=[-1, 1], autorange=False),
|
||||
# aspectmode='cube'),
|
||||
# scene3=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
# yaxis=dict(range=[-1, 1], autorange=False),
|
||||
# zaxis=dict(range=[-1, 1], autorange=False),
|
||||
# aspectmode='cube'),
|
||||
# scene4=dict(xaxis=dict(range=[-1, 1], autorange=False),
|
||||
# yaxis=dict(range=[-1, 1], autorange=False),
|
||||
# zaxis=dict(range=[-1, 1], autorange=False),
|
||||
# aspectmode='cube'))
|
||||
|
||||
# fig.show()
|
||||
|
||||
|
||||
# uv_mask = batch['uv_mask'][b].cpu().numpy()
|
||||
# uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0)
|
||||
# uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2
|
||||
# uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2
|
||||
# uv_unposed = (batch['uv_unposed'][b].cpu().numpy() + 1) / 2
|
||||
# uv_body_posed = (batch['uv_body_posed'][b].cpu().numpy() + 1) / 2
|
||||
# uv_normals = (batch['uv_normals'][b].cpu().numpy() + 1.5) / 3
|
||||
|
||||
# img = np.concatenate([uv_mask, uv_static, uv_posed, uv_unposed, uv_body_posed, uv_normals], axis=2)
|
||||
# img = np.transpose(img, (1, 2, 0))
|
||||
# ex.imshow(img).show()
|
||||
|
||||
# pass
|
||||
|
||||
pass
|
||||
|
||||
|
||||
+5
-13
@@ -82,7 +82,7 @@ class StaticCache():
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_static' + res_string in self.signature:
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static
|
||||
uv_static = np.load(os.path.join(self.folder, 'uv_static' + self.filled + res_string, id + '.npz')); uv_static = uv_static['uv_static' + self.filled]; data['uv_static' + res_string] = uv_static.astype(np.float32)
|
||||
if 'uv_mask' + res_string in self.signature:
|
||||
if self.is_cape:
|
||||
data['uv_mask' + res_string] = self.uv_masks[res]
|
||||
@@ -158,6 +158,7 @@ class DatasetCape(torch.utils.data.Dataset):
|
||||
item['id'] = frame.id
|
||||
if 'frame_nr' in self.signature:
|
||||
item['frame_nr'] = frame.frame_nr
|
||||
item['garment'] = 'cape'
|
||||
if 'pose_shape' in self.signature:
|
||||
pose_shape = np.load(os.path.join(self.folder, 'pose_shape', frame.id + separator + frame.frame_nr + '.npy')).astype(np.float32)
|
||||
pose_shape[:3] = 0
|
||||
@@ -183,27 +184,18 @@ class DatasetCape(torch.utils.data.Dataset):
|
||||
if res != 512:
|
||||
res_string = '_' + str(res)
|
||||
|
||||
if 'uv_curvature' + res_string in self.signature:
|
||||
uv_curvature = np.load(os.path.join(self.folder, 'uv_curvature' + res_string, frame.id + separator + frame.frame_nr + '.npz'));
|
||||
if res_string != '':
|
||||
uv_curvature = uv_curvature['uv_curvature']
|
||||
else:
|
||||
uv_curvature = uv_curvature['curvature_uv']
|
||||
uv_curvature = torch.nn.functional.one_hot(torch.tensor(uv_curvature, device=self.device, dtype=torch.long), 7).float().to(self.device)
|
||||
item['uv_curvature' + res_string] = uv_curvature
|
||||
|
||||
if 'uv_unposed' + res_string in self.signature:
|
||||
uv_unposed_no_shape = np.load(os.path.join(self.folder, 'uv_unposed_no_shape' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_unposed_no_shape = uv_unposed_no_shape['uv_unposed_no_shape']
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape
|
||||
item['uv_unposed' + res_string] = uv_unposed_no_shape.astype(np.float32)
|
||||
|
||||
if 'uv_body_posed' + res_string in self.signature:
|
||||
uv_posed_body_no_orient = np.load(os.path.join(self.folder, uv_posed_body_no_orient_folder_name + res_string, frame.id + separator + frame.frame_nr + '.npz'))
|
||||
uv_posed_body_no_orient = uv_posed_body_no_orient['uv_posed_body_no_orient']
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient
|
||||
item['uv_body_posed' + res_string] = uv_posed_body_no_orient.astype(np.float32)
|
||||
|
||||
if 'uv_posed' + res_string in self.signature:
|
||||
uv_posed = np.load(os.path.join(self.folder, 'uv_posed_no_orient' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_posed = uv_posed['uv_posed_no_orient']
|
||||
item['uv_posed' + res_string] = uv_posed
|
||||
item['uv_posed' + res_string] = uv_posed.astype(np.float32)
|
||||
|
||||
if 'uv_normals' + res_string in self.signature:
|
||||
uv_normals = np.load(os.path.join(self.folder, 'uv_normals' + res_string, frame.id + separator + frame.frame_nr + '.npz')); uv_normals = uv_normals['uv_normals']
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
import argparse
|
||||
import os
|
||||
import tqdm
|
||||
import cv2
|
||||
import scipy.interpolate as spi
|
||||
import numpy as np
|
||||
import plotly.graph_objects as go
|
||||
import plotly.express as ex
|
||||
import torch
|
||||
from torchvision.transforms import Resize
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
||||
|
||||
|
||||
def get_img_idxs():
|
||||
img_idxs = dict()
|
||||
img_coords = dict()
|
||||
uv_masks = dict()
|
||||
uv_mask = np.load('custom_data/garment_uv_mask_full.npy')
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - uv_mask))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((uv_mask))), (1,0))
|
||||
img_idxs[512] = img_idx
|
||||
img_coords[512] = coords
|
||||
uv_masks[512] = uv_mask.astype(np.int32)
|
||||
for res in [32, 64, 128, 256]:
|
||||
down = resizers_nearest[res](torch.from_numpy(uv_mask[None,:,:]).to(device)).cpu().numpy()
|
||||
down = np.squeeze(down.astype(bool))
|
||||
coords = np.transpose(np.stack(np.nonzero((1 - down))), (1,0))
|
||||
img_idx = np.transpose(np.stack(np.nonzero((down))), (1,0))
|
||||
img_idxs[res] = img_idx
|
||||
img_coords[res] = coords
|
||||
uv_masks[res] = down.astype(np.int32)
|
||||
|
||||
|
||||
uv_masks_body = dict()
|
||||
uv_mask_body = np.load('custom_utils/mask_256x512.npy')
|
||||
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
|
||||
uv_masks_body[512] = uv_mask_body
|
||||
for res in [32, 64, 128, 256]:
|
||||
down = resizers_nearest[res](torch.from_numpy(uv_mask_body[None,:,:]).to(device)).cpu().numpy()
|
||||
down = np.squeeze(down.astype(bool))
|
||||
uv_masks_body[res] = down
|
||||
|
||||
return img_idxs, img_coords, uv_masks, uv_masks_body
|
||||
|
||||
|
||||
def fill(img, mask):
|
||||
coords = img_coords[img.shape[0]]
|
||||
img_idx = img_idxs[img.shape[0]]
|
||||
# coords = np.transpose(np.stack(np.nonzero((1 - uv_masks[img.shape[0]]))), (1,0))
|
||||
# img_idx = np.transpose(np.stack(np.nonzero((uv_masks[img.shape[0]]))), (1,0))
|
||||
img_data = img[img_idx[:,0], img_idx[:,1], :]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256,3), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1], :] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1], :] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def fill_label(img, mask):
|
||||
coords = img_coords[img.shape[0]]
|
||||
img_idx = img_idxs[img.shape[0]]
|
||||
img_data = img[img_idx[:,0], img_idx[:,1]]
|
||||
|
||||
interpolator = spi.NearestNDInterpolator(x=img_idx, y=img_data)
|
||||
img_2 = interpolator(coords)
|
||||
res = np.zeros((512,256), dtype=np.float32)
|
||||
res[coords[:,0], coords[:,1]] = img_2
|
||||
res[img_idx[:,0], img_idx[:,1]] = img_data
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def generate_downscaled_uv_map(folder, id, size):
|
||||
img = np.load(os.path.join(folder, 'uv_mask', id + '.npz')); img = img['uv_mask'].astype(np.float32)
|
||||
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(bool))
|
||||
np.savez_compressed(os.path.join(folder, 'uv_mask' + '_' + str(size), id + '.npz'), uv_mask = down)
|
||||
|
||||
|
||||
class UVMaskStore():
|
||||
def __init__(self, folder, size):
|
||||
self.size = size
|
||||
|
||||
def __getitem__(self, id) -> dict:
|
||||
return uv_masks[self.size]
|
||||
|
||||
|
||||
def downscale_uv_maps_posed_body(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_posed_body_no_orient']
|
||||
img = fill(img, uv_masks_body[512])
|
||||
for size in sizes:
|
||||
mask = uv_masks_body[size]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_posed_body_no_orient = down)
|
||||
|
||||
|
||||
def downscale_uv_normals(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_normals']
|
||||
img = fill(img, uv_masks_body[512])
|
||||
for size in sizes:
|
||||
mask = uv_masks_body[size]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_normals = down)
|
||||
|
||||
|
||||
|
||||
def downscale_uv_maps_unposed(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_unposed_no_shape']
|
||||
mask_full = uv_mask_store[512][None]
|
||||
img = fill(img, mask_full)
|
||||
for size in sizes:
|
||||
mask = uv_mask_store[size][None]
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_unposed_no_shape = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_static(folder, type, sizes):
|
||||
for size in sizes:
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_static']
|
||||
mask_full = uv_mask_store[512][None]
|
||||
img = fill(img, mask_full)
|
||||
|
||||
for size in sizes:
|
||||
down = resizers[size](torch.from_numpy(img.transpose((2,0,1))).to(device)).cpu().numpy()
|
||||
mask = uv_mask_store[size][None]
|
||||
down = down * mask[None,:,:]
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_static = down)
|
||||
|
||||
|
||||
def downscale_uv_maps_mask(folder, type, size):
|
||||
os.makedirs(os.path.join(folder, type + '_' + str(size)), exist_ok=True)
|
||||
|
||||
frames = os.listdir(os.path.join(folder, type))[int(args.start):int(args.end)]
|
||||
|
||||
for frame in tqdm.tqdm(frames):
|
||||
img = np.load(os.path.join(folder, type, frame)); img = img['uv_mask'].astype(np.float32)
|
||||
down = tf.image.resize(img[:,:,tf.newaxis], size=(size, int(size/2)), method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
|
||||
down = np.squeeze(down.numpy().astype(bool))
|
||||
np.savez_compressed(os.path.join(folder, type + '_' + str(size), frame), uv_mask = down)
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127_2/cape_processed/train/male')
|
||||
parser.add_argument("-s", "--start", help = "start of range", default=0)
|
||||
parser.add_argument("-e", "--end", help = "end of range", default=100000)
|
||||
parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
|
||||
parser.add_argument("-d", "--debug", help = "end of range", default=False, action='store_true')
|
||||
args = parser.parse_args()
|
||||
|
||||
device = 'cuda'
|
||||
|
||||
sizes = [32, 64, 128, 256]
|
||||
folder = args.folder
|
||||
|
||||
resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BICUBIC).to(device) for res in sizes }
|
||||
resizers_nearest = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.NEAREST_EXACT).to(device) for res in sizes }
|
||||
|
||||
img_idxs, img_coords, uv_masks, uv_masks_body = get_img_idxs()
|
||||
|
||||
uv_mask_store = { size: UVMaskStore(folder, size=size) for size in sizes }
|
||||
uv_mask_store[512] = UVMaskStore(folder, size=512)
|
||||
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=32)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=64)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=128)
|
||||
# downscale_uv_maps_mask(folder=folder, type='uv_mask', size=256)
|
||||
# downscale_uv_maps_static(folder=folder, type='uv_static', sizes=sizes)
|
||||
|
||||
downscale_uv_maps_unposed(folder=folder, type='uv_unposed_no_shape', sizes=sizes)
|
||||
downscale_uv_normals(folder=folder, type='uv_normals', sizes=sizes)
|
||||
downscale_uv_maps_posed_body(folder=folder, type='uv_posed_body_no_orient', sizes=sizes)
|
||||
|
||||
|
||||
print()
|
||||
+177
-108
@@ -36,94 +36,6 @@ 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])
|
||||
@@ -261,14 +173,15 @@ def create_mesh_body_posed_no_orient(folder, out_folder):
|
||||
continue
|
||||
|
||||
info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
|
||||
|
||||
|
||||
pose, shape, gender = info['poses'], info['shape'], info['gender']
|
||||
# if pose.shape[0] == 72:
|
||||
pose = np.transpose(pose, [1, 0]).astype(np.float32)
|
||||
|
||||
pose[:, :3] = 0
|
||||
beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
|
||||
|
||||
smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0])
|
||||
|
||||
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()
|
||||
|
||||
@@ -456,6 +369,9 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
|
||||
|
||||
pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
|
||||
shape = torch.tensor(shape, dtype=torch.float32, device='cuda')
|
||||
|
||||
smpl = SMPL(path='smpl', device='cuda', batch_size=pose.shape[0])
|
||||
|
||||
|
||||
for garment in garments:
|
||||
if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
|
||||
@@ -627,6 +543,49 @@ def create_frames_file(folder, out_folder, split=0.9):
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
|
||||
def create_frames_file_test(folder, out_folder):
|
||||
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)
|
||||
|
||||
ids_skirt = list(ids_skirt)
|
||||
ids_non_skirt = list(ids_non_skirt)
|
||||
|
||||
with open(os.path.join(out_folder, 'frames_test.txt'), 'w') as f:
|
||||
for frame in frames_non_skirt:
|
||||
if frame.id in ids_non_skirt:
|
||||
f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
|
||||
|
||||
with open(os.path.join(out_folder, 'frames_test_skirt.txt'), 'w') as f:
|
||||
for frame in frames_skirt:
|
||||
if frame.id in ids_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)
|
||||
|
||||
@@ -674,25 +633,26 @@ def create_frames_file_partial(folder, out_folder, split=0.9, part=0.1):
|
||||
|
||||
|
||||
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')
|
||||
for i in range(100, 200, 5):
|
||||
frame: Frame = Frame('01512', i, '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]
|
||||
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[::2][frame.frame_nr]
|
||||
|
||||
F = np.load('smpl/smpl_faces_skirt.npy')
|
||||
_, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
|
||||
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()
|
||||
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
|
||||
|
||||
|
||||
@@ -734,13 +694,118 @@ def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'):
|
||||
np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment)
|
||||
|
||||
|
||||
def remove_irregular_poses(folder, folder_reg, out_folder):
|
||||
ids = sorted(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('.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:
|
||||
|
||||
verts_count = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V'].shape[0]
|
||||
if info['poses'].shape[1] != verts_count:
|
||||
shutil.rmtree(os.path.join(folder, id))
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.pc16'))
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'))
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'))
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
|
||||
# os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.npy'))
|
||||
|
||||
if len(garments) == 2:
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_static', f'{id}_{garments[1]}.pc16'))
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garments[1]}.pc16'))
|
||||
if os.path.exists(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.pc16')):
|
||||
os.remove(os.path.join(out_folder, 'mesh_mask', f'{id}_{garments[1]}.npy'))
|
||||
break
|
||||
# os.remove(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garments[1]}.npy'))
|
||||
|
||||
|
||||
def double_check_irregular_poses(folder, folder_reg, out_folder):
|
||||
# ids = sorted(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('.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:
|
||||
# body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0]
|
||||
# verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0]
|
||||
|
||||
# assert body_count == verts_count
|
||||
|
||||
frames = sorted(os.listdir(out_folder + '/mesh_posed_no_orient'))
|
||||
frames = [(id.split('_')[0], id.split('_')[1].split('.')[0]) for id in frames]
|
||||
|
||||
for frame in tqdm(frames):
|
||||
id = frame[0]
|
||||
garment_dir = os.path.join(folder, id)
|
||||
if not os.path.isdir(garment_dir):
|
||||
continue
|
||||
|
||||
garment = frame[1]
|
||||
body_count = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V'].shape[0]
|
||||
verts_count = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V'].shape[0]
|
||||
|
||||
assert body_count == verts_count
|
||||
|
||||
|
||||
def get_ids():
|
||||
ids = []
|
||||
with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train.txt'), 'r') as f:
|
||||
for line in f:
|
||||
ids.append(line.split()[0])
|
||||
|
||||
with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train_skirt.txt'), 'r') as f:
|
||||
for line in f:
|
||||
ids.append(line.split()[0])
|
||||
|
||||
ids = sorted(list(set(ids)))
|
||||
|
||||
with open(os.path.join('/data-local/127/cloth3d_processed', 'frames_train_ids.txt'), 'w') as f:
|
||||
for id in ids:
|
||||
f.write(f'{id}\n')
|
||||
|
||||
ids = []
|
||||
with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test.txt'), 'r') as f:
|
||||
for line in f:
|
||||
ids.append(line.split()[0])
|
||||
|
||||
with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test_skirt.txt'), 'r') as f:
|
||||
for line in f:
|
||||
ids.append(line.split()[0])
|
||||
|
||||
ids = sorted(list(set(ids)))
|
||||
|
||||
with open(os.path.join('/data-local/127_2/cloth3d_processed_test', 'frames_test_ids.txt'), 'w') as f:
|
||||
for id in ids:
|
||||
f.write(f'{id}\n')
|
||||
|
||||
pass
|
||||
|
||||
|
||||
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'
|
||||
# out_folder = '/data-local/127_2/cloth3d_processed_test'
|
||||
|
||||
folder = '/data-net/dataset/cloth3d/train'
|
||||
folder_reg = '/data-local/127/cloth3d_registration/train'
|
||||
@@ -754,6 +819,8 @@ if __name__ == '__main__':
|
||||
# create_extras_file(folder)
|
||||
|
||||
# check_irregular_poses(folder, folder_reg, out_folder)
|
||||
# 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)
|
||||
# verify_registration_ans_mask(folder, out_folder)
|
||||
@@ -762,9 +829,11 @@ if __name__ == '__main__':
|
||||
# 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)
|
||||
# calculate_average(folder, out_folder)
|
||||
|
||||
# create_frames_file(folder, out_folder)
|
||||
get_ids()
|
||||
# create_frames_file_test(folder, out_folder)
|
||||
# create_frames_file_partial(folder, out_folder, part=0.5)
|
||||
|
||||
# folder = '/data-local/127_2/cape_processed/test/female'
|
||||
|
||||
Reference in New Issue
Block a user