| 11 | |
| 12 | @hydra.main(config_path="config/encode.yaml") |
| 13 | def encode_dataset(cfg): |
| 14 | out_dir = Path(utils.to_absolute_path(cfg.out_dir)) / cfg.dataset.path / "test" |
| 15 | out_dir.mkdir(exist_ok=True, parents=True) |
| 16 | |
| 17 | root_path = Path(utils.to_absolute_path("datasets")) / cfg.dataset.path |
| 18 | with open(root_path / "test.json") as file: |
| 19 | metadata = json.load(file) |
| 20 | |
| 21 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 22 | |
| 23 | encoder = hydra.utils.instantiate(cfg.model.encoder) |
| 24 | encoder.to(device) |
| 25 | |
| 26 | print("Load checkpoint from: {}:".format(cfg.checkpoint)) |
| 27 | checkpoint_path = utils.to_absolute_path(cfg.checkpoint) |
| 28 | checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) |
| 29 | encoder.load_state_dict(checkpoint["encoder"]) |
| 30 | |
| 31 | encoder.eval() |
| 32 | |
| 33 | for _, _, _, path in tqdm(metadata): |
| 34 | path = root_path.parent / path |
| 35 | mel = torch.from_numpy(np.load(path.with_suffix(".mel.npy"))).unsqueeze(0).to(device) |
| 36 | with torch.no_grad(): |
| 37 | z, indices = encoder.encode(mel) |
| 38 | |
| 39 | z = z.squeeze().cpu().numpy() |
| 40 | |
| 41 | out_path = out_dir / path.stem |
| 42 | with open(out_path.with_suffix(".txt"), "w") as file: |
| 43 | np.savetxt(file, z, fmt="%.16f") |
| 44 | |
| 45 | |
| 46 | if __name__ == "__main__": |