save all

data/dataset.py

@@ -276,7 +276,7 @@ if __name__ == '__main__':
                   'uv_mask',
                   }
     
-    dataset = Dataset('/data-local/127/cloth3d_processed', signature, resolutions=[512], device=device, frames_file='frames_train.txt')
+    dataset = Dataset('/data-local/127_2/cloth3d_processed_test', signature, resolutions=[512], device=device, frames_file='frames_temp.txt')
     train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=False)
     smpl = SMPL(orig_smpl=False, path='smpl', device='cuda')
 
@@ -293,6 +293,16 @@ if __name__ == '__main__':
     step = 0
     for batch in train_dataloader:
         batch
+
+        from custom_utils.custom_utils import save_stl
+        fix_data = np.load(os.path.join(dataset.folder, 'mesh_fixed', f'{batch["id"][0]}_{batch["garment"][0]}.npz'))
+        faces = fix_data['faces']
+        mask = fix_data['mask']
+
+        save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_posed.stl', batch['mesh_posed'][0].cpu().numpy()[mask], faces)
+        save_stl(f'temp/{batch["id"][0]}_{batch["frame_nr"][0].cpu().numpy()}_{batch["garment"][0]}_mesh_unposed.stl', batch['mesh_unposed'][0].cpu().numpy()[mask], faces)
+        
+
         print(batch['id'][0])
         with profiler:
             m = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], batch['mesh_unposed'])[0]