meng2003/music2dance

1import librosa2import numpy as np3from pathlib import Path4import json5import os.path6import sys7import argparse8 9'''10Compute transforms which can be computed sequentially (so they implement the `partial_fit` function)11'''12 13THIS_DIR = os.path.dirname(os.path.abspath(__file__))14ROOT_DIR = os.path.abspath(os.path.join(THIS_DIR, os.pardir))15sys.path.append(ROOT_DIR)16 17parser = argparse.ArgumentParser(description="Preprocess songs data")18 19parser.add_argument("data_path", type=str, help="Directory contining Beat Saber level folders")20parser.add_argument("--feature_name", metavar='', type=str, default="mel", help="mel, chroma, multi_mel")21parser.add_argument("--transforms", metavar='', type=str, default="scaler", help="comma-separated lists of transforms to extract (scaler,pca_transform)")22args = parser.parse_args()23 24# makes arugments into global variables of the same name, used later in the code25globals().update(vars(args))26data_path = Path(data_path)27 28## distributing tasks accross nodes ##29from mpi4py import MPI30comm = MPI.COMM_WORLD31rank = comm.Get_rank()32size = comm.Get_size()33print(rank)34assert size == 135candidate_files = sorted(data_path.glob('**/*'+feature_name+'.npy'), key=lambda path: path.parent.__str__())36tasks = range(len(candidate_files))37 38from sklearn import decomposition, preprocessing39import pickle40transforms = transforms.split(",")41transforms_dict = {}42for transform in transforms:43    if transform == "scaler":44        scaler = preprocessing.StandardScaler()45        transforms_dict["scaler"] = scaler46    elif transform == "pca_transform":47        features = np.load(candidate_files[0].__str__())48        feature_size = features.shape[1]49        pca = decomposition.PCA(n_components=feature_size)50        transforms_dict["pca_transform"] = pca51    else:52        raise NotImplementedError("Transform type "+transform+" not implemented")53for i in tasks:54    path = candidate_files[i]55    feature_file = path.__str__()56    features = np.load(feature_file)57    for transform in transforms:58        if len(features.shape) == 3:59            features = features[:,0,:]60        transforms_dict[transform].partial_fit(features)61 62for transform in transforms:63    pickle.dump(transforms_dict[transform], open(data_path.joinpath(feature_name+'_'+transform+'.pkl'), 'wb'))64