save

scripts/cvae_pyramid.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.pyr_64_cape_f_skin import config
+    from configs.iccvw.pyr_64_128_256_512 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -501,6 +501,86 @@ def speed_step(levels_input, conds):
     return model.sample([levels_input, conds], training=False)
 
 
+def sample():
+    from data.utils import Frame
+    from utils.utils import save_stl
+    from pathlib import Path
+    import trimesh as tm
+
+    def create_temp_frames(folder: str, frame: Frame):
+        with open(Path(folder) / 'frames_temp.txt', 'w') as f:
+            f.writelines([f'{frame}\n' for _ in range(4)])
+
+    model.eval()
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
+    count = 0
+
+    frames = ['01852_70_Tshirt']
+    frames = [Frame(frame_str=frame) for frame in frames]
+
+
+
+    for frame in frames:
+        create_temp_frames(config['train_dir'], frame)
+        dataset = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=True, resolutions=[64, 128, 256, 512], is_cape=config['is_cape'])
+        dataloader = DataLoader(dataset, batch_size=4, drop_last=True, shuffle=True, )
+        batch = next(iter(dataloader))
+
+        uv_static = batch['uv_static']
+        uv_body_posed = batch['uv_body_posed']
+        uv_mask = { 512: batch['uv_mask'], 256: batch['uv_mask_256'], 128: batch['uv_mask_128'], 64: batch['uv_mask_64'] }
+        uv_normals = batch['uv_normals']
+
+        batch_size = uv_static.size(0)
+
+
+        for i in range(batch_size):
+            counter = 0
+            frame = Frame(batch['id'][i], batch['frame_nr'][i], batch['garment'][i])
+            fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+            mask = fix_data['mask']
+            faces = fix_data['faces']
+
+            gt = batch['mesh_posed'][i].cpu().numpy()[mask]
+            save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_gt.stl', gt, faces)
+
+            uv_static_tiled = uv_static[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_body_posed_tiled = uv_body_posed[i][None,:,:,:].repeat(4, 1, 1, 1)
+            uv_mask_tiled = { 512: uv_mask[512][i][None,:,:].repeat(4, 1, 1), 256: uv_mask[256][i][None,:,:].repeat(4, 1, 1), 128: uv_mask[128][i][None,:,:].repeat(4, 1, 1), 64: uv_mask[64][i][None,:,:].repeat(4, 1, 1)}
+            uv_normals_tiled = uv_normals[i][None,:,:,:].repeat(4, 1, 1, 1)
+            pose_shape = batch['pose_shape'][i][None,:].repeat(4, 1)
+            body_posed = batch['mesh_body_posed'][i][None,:,:].repeat(4, 1, 1)
+            mask_tiled = torch.from_numpy(mask[None,:]).to(body_posed.device).repeat(4, 1)
+        
+            for _ in range(4):
+                with torch.inference_mode():
+                    pred = model.sample(uv_static_tiled, uv_body_posed_tiled, uv_normals_tiled, uv_mask_tiled)
+                
+                mesh_pred = uv2mesh(pred)
+
+                shape = pose_shape[:, 72:]
+                gender = torch.greater_equal(shape[:, 10] , 0.5).to(torch.long)
+                shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])])
+
+                mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None]
+                mesh_pred_posed = smpl(pose_shape[:, :72], pose_shape[:, 72:], mesh_pred_shaped)[0]
+
+                mesh_posed_fixed = fix_collision(mesh_pred_posed, body_posed, mask_tiled).cpu().numpy()
+
+                for k in range(4):
+                    _mesh_posed_fixed = mesh_posed_fixed[k, mask, :]
+                    m = tm.Trimesh(_mesh_posed_fixed, faces, process=False, validate=False)
+                    tm.smoothing.filter_humphrey(m, iterations=2)
+                    verts = m.vertices
+                    save_stl(Path('temp') / 'sampling_pyramid' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_{counter}.stl', verts, faces)
+                    counter += 1
+
+            pass
+
+
+        if count > 100:
+            break
+
 
 def train():
     global g_balancing_coeff
@@ -560,6 +640,9 @@ if __name__ == '__main__':
 
     model: Pyramid = Pyramid(config)
 
+    sample()
+    exit()
+
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
     ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)