diff --git a/data/dataset.py b/data/dataset.py index 7ebc6d5..4da75f7 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -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