()
| 10 | |
| 11 | |
| 12 | def main(): |
| 13 | # parse options |
| 14 | cfg = parse_args(phase="test") # parse config file |
| 15 | cfg.TRAIN.STAGE = "token" |
| 16 | cfg.TRAIN.BATCH_SIZE = 1 |
| 17 | |
| 18 | # set seed |
| 19 | pl.seed_everything(cfg.SEED_VALUE) |
| 20 | |
| 21 | # gpu setting |
| 22 | if cfg.ACCELERATOR == "gpu": |
| 23 | os.environ["PYTHONWARNINGS"] = "ignore" |
| 24 | os.environ["TOKENIZERS_PARALLELISM"] = "false" |
| 25 | |
| 26 | # create dataset |
| 27 | datasets = build_data(cfg, phase='token') |
| 28 | print("datasets module initialized") |
| 29 | output_dir = os.path.join(datasets.hparams.data_root, cfg.DATASET.CODE_PATH) |
| 30 | |
| 31 | os.makedirs(output_dir, exist_ok=True) |
| 32 | |
| 33 | # create model |
| 34 | model = build_model(cfg, datasets) |
| 35 | if hasattr(model, "motion_vae"): |
| 36 | model.vae = model.motion_vae |
| 37 | print("model loaded") |
| 38 | |
| 39 | # Strict load vae model |
| 40 | assert cfg.TRAIN.PRETRAINED_VAE is not None |
| 41 | state_dict = torch.load(cfg.TRAIN.PRETRAINED_VAE, |
| 42 | map_location="cpu")['state_dict'] |
| 43 | print(f"Loading pretrain vae from {cfg.TRAIN.PRETRAINED_VAE}") |
| 44 | |
| 45 | if cfg.ACCELERATOR == "gpu": |
| 46 | model = model.to('cuda') |
| 47 | |
| 48 | for batch in tqdm(datasets.train_dataloader(), |
| 49 | desc=f'motion tokenize'): |
| 50 | name = batch['text'] |
| 51 | |
| 52 | pose = batch['motion'] |
| 53 | pose = pose.cuda().float() |
| 54 | |
| 55 | if pose.shape[1] == 0: |
| 56 | continue |
| 57 | target, _ = model.vae.encode(pose) |
| 58 | target = target.to('cpu').numpy() |
| 59 | |
| 60 | target_path = os.path.join(output_dir, name[0] + '.npy') |
| 61 | Path(target_path).parent.mkdir(parents=True, exist_ok=True) |
| 62 | np.save(target_path, target) |
| 63 | |
| 64 | print( |
| 65 | f'Motion tokenization done, the motion tokens are saved to {output_dir}' |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | if __name__ == "__main__": |
no test coverage detected