save all

scripts/cvae.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 if args['debug']:
-    from configs.iccvw.abl_all_5_var import config
+    from configs.iccvw.abl_all_5 import config
 else:
     config_module = importlib.import_module(args['config'])
     config = config_module.config
@@ -45,7 +45,7 @@ from data.dataset import Dataset
 from models.cond_encoder import AutoEncoder
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
-from utils.utils import text_to_html
+from utils.utils import text_to_html, save_stl, create_temp_frames
 from utils.postprocess import fix_collision
 from losses.losses import collision_loss, normal_loss
 from pytorch_msssim import MS_SSIM
@@ -352,7 +352,52 @@ def sample():
             break
 
 
+def generate():
+    from data.utils import Frame
+    load_model(model, f'checkpoints/iccvw/{config["name"]}')
 
+    frame = '01852_70_Tshirt'
+    # frame = '03849_27_Dress' # 03849_27_Dress 00054_1_Dress
+    frame = Frame(frame_str=frame)
+    create_temp_frames(config['train_dir'], frame)
+
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file='frames_temp.txt', use_cache=config['dataset_in_memory'])
+    train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
+
+    batch = next(iter(train_dataloader))
+
+    if config['use_skinning']:
+        mesh_orig, uv_orig = batch['mesh_unposed'], batch['uv_unposed']
+    else:
+        mesh_orig, uv_orig = batch['mesh_posed'], batch['uv_posed']
+
+    # with record_function("forward"):
+    model.eval()
+    with torch.inference_mode():
+        pred, _, _ = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_normals'], batch['uv_mask']])
+    mesh_pred = uv2mesh(pred)
+
+    if config['use_skinning']:
+        shape = batch['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])])
+        # shape_offsets = torch.stack([torch.matmul(shape[i, :10], smpl.shapedirs[gender[i]][:, 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(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0]
+    else:
+        mesh_pred_posed = mesh_pred
+
+    fix_data = np.load(Path(dataset_train.folder) / 'mesh_fixed' / f'{frame.id}_{frame.garment}.npz')
+    mask = fix_data['mask']
+    faces = fix_data['faces']
+
+    pred = mesh_pred_posed.cpu().numpy()[0, mask, :]
+    body = batch['mesh_body_posed'].cpu().numpy()[0]
+    body_faces = smpl.f
+    save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_mesh.stl', pred, faces)
+    save_stl(Path('temp') / 'baseline' / f'{frame.id}_{frame.frame_nr}_{frame.garment}_baseline_body.stl', body, body_faces)
+    pass
 
 
 def train():
@@ -410,6 +455,9 @@ if __name__ == '__main__':
 
     model = CVAE(config)
 
+    generate()
+    exit()
+
     sample()
     exit()