mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-13 23:53:02 +00:00
Paul Sajna
f86b44d8fc
* 48KHz softvc * add some notes * almost forgot the inference script * fix readme link * requirements.txt
67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
from pathlib import Path
|
|
import logging
|
|
import argparse
|
|
|
|
import torch
|
|
import numpy as np
|
|
from sklearn.cluster import KMeans
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def cluster(args):
|
|
with open(args.subset) as file:
|
|
subset = [line.strip() for line in file]
|
|
|
|
logger.info(f"Loading features from {args.in_dir}")
|
|
features = []
|
|
for path in subset:
|
|
in_path = args.in_dir / path
|
|
features.append(np.load(in_path.with_suffix(".npy")))
|
|
features = np.concatenate(features, axis=0)
|
|
|
|
logger.info(f"Clustering features of shape: {features.shape}")
|
|
kmeans = KMeans(n_clusters=args.n_clusters).fit(features)
|
|
|
|
checkpoint_path = args.checkpoint_dir / f"kmeans_{args.n_clusters}.pt"
|
|
checkpoint_path.parent.mkdir(exist_ok=True, parents=True)
|
|
torch.save(
|
|
checkpoint_path,
|
|
{
|
|
"n_features_in_": kmeans.n_features_in_,
|
|
"_n_threads": kmeans._n_threads,
|
|
"cluster_centers_": kmeans.cluster_centers_,
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Cluster speech features features.")
|
|
parser.add_argument(
|
|
"in_dir",
|
|
metavar="in-dir",
|
|
help="path to the encoded dataset",
|
|
type=Path,
|
|
)
|
|
parser.add_argument(
|
|
"subset",
|
|
matavar="subset",
|
|
help="path to the .txt file containing the list of files to cluster",
|
|
type=Path,
|
|
)
|
|
parser.add_argument(
|
|
"checkpoint_dir",
|
|
metavar="checkpoint-dir",
|
|
help="path to the checkpoint directory",
|
|
type=Path,
|
|
)
|
|
parser.add_argument(
|
|
"--n-clusters",
|
|
help="number of clusters",
|
|
type=int,
|
|
default=100,
|
|
)
|
|
args = parser.parse_args()
|
|
cluster(args)
|