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

small datset fix

Browse Source
This commit is contained in:
luni
2023-07-04 18:21:39 +02:00
parent ffaf1a9a49
commit 33782848f2
1 changed files with 89 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
data/dataset.py
+89 -49
View File
@@ -140,7 +140,7 @@ class DatasetCloth3D(torch.utils.data.Dataset):
self.resolutions = resolutions self.resolutions = resolutions
self.transforms = tv.transforms.ToTensor() 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) 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 item['uv_unposed' + res_string] = uv_unposed_no_shape
if 'uv_body_posed' + res_string in self.signature: 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 item['uv_body_posed' + res_string] = uv_body_posed_no_orient
if 'uv_posed' + res_string in self.signature: if 'uv_posed' + res_string in self.signature:
@@ -244,6 +244,9 @@ if __name__ == '__main__':
import plotly.express as ex import plotly.express as ex
from plotly.subplots import make_subplots from plotly.subplots import make_subplots
from custom_utils.custom_utils import _maskFaces 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 = np.load('smpl/smpl_faces.npy')
F_skirt = np.load('smpl/smpl_faces_skirt.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') 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) 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: for batch in train_dataloader:
batch batch
b = 0 print(batch['id'][0])
mesh_mask = batch['mesh_mask'][b].cpu().numpy() with profiler:
faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0]
_, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) profiler.step()
step += 1
if step > 25:
break
mesh_static = batch['mesh_static'][b].cpu().numpy()[mesh_mask] profiler.stop()
mesh_posed = batch['mesh_posed'][b].cpu().numpy()[mesh_mask] print(profiler.key_averages().table(sort_by="cpu_time_total", row_limit=25))
mesh_unposed = batch['mesh_unposed'][b].cpu().numpy()[mesh_mask] print(profiler.key_averages().table(sort_by="cuda_time_total", row_limit=25))
mesh_body_posed = batch['mesh_body_posed'][b].cpu().numpy() print(profiler.key_averages(group_by_stack_n=5).table(sort_by='self_cpu_time_total', row_limit=10))
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() # b = 0
uv_mask = np.repeat(uv_mask[np.newaxis, :, :], 3, axis=0) # mesh_mask = batch['mesh_mask'][b].cpu().numpy()
uv_static = (batch['uv_static'][b].cpu().numpy() + 1) / 2 # faces = F_skirt if batch['garment'][b] in ['Skirt', 'Dress'] else F
uv_posed = (batch['uv_posed'][b].cpu().numpy() + 1) / 2 # _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask)
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) # c = collision_loss(batch['mesh_posed'], batch['mesh_body_posed'], batch['mesh_mask'])
img = np.transpose(img, (1, 2, 0)) # 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),
ex.imshow(img).show() # 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 pass