meng2003/music2dance

1import librosa2import numpy as np3from pathlib import Path4import json5import os.path6import sys7import argparse8 9'''10Compute transforms which r4quire to be fitted on all the data at once rather than sequential (so they dont 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)16from audio_feature_utils import extract_features_hybrid, extract_features_mel, extract_features_multi_mel17from utils import distribute_tasks18#from scripts.feature_extraction.utils import distribute_tasks19 20parser = argparse.ArgumentParser(description="Preprocess songs data")21 22parser.add_argument("data_path", type=str, help="Directory contining Beat Saber level folders")23parser.add_argument("--feature_name", metavar='', type=str, default="mel", help="mel, chroma, multi_mel")24parser.add_argument("--transforms", metavar='', type=str, default="scaler", help="comma-separated lists of transforms to extract (scaler,pca_transform)")25args = parser.parse_args()26 27# makes arugments into global variables of the same name, used later in the code28globals().update(vars(args))29data_path = Path(data_path)30 31## distributing tasks accross nodes ##32from mpi4py import MPI33comm = MPI.COMM_WORLD34rank = comm.Get_rank()35size = comm.Get_size()36print(rank)37assert size == 138candidate_files = sorted(data_path.glob('**/*'+feature_name+'.npy'), key=lambda path: path.parent.__str__())39tasks = range(len(candidate_files))40 41from sklearn import decomposition, preprocessing42features = None43for i in tasks:44    path = candidate_files[i]45    feature_file = path.__str__()46    if i == 0:47        features = np.load(feature_file)48    else:49        feature = np.load(feature_file)50        features = np.concatenate([features,feature],0)51 52import pickle53transforms = transforms.split(",")54for transform in transforms:55    if transform == "2moments":56        if len(features.shape) == 3:57            features = features[:,0,:]58        C = np.dot(features.T,features)/features.shape[0]59        m = np.mean(features,0)60        pickle.dump((m,C), open(data_path.joinpath(feature_name+'_2moments.pkl'), 'wb'))61    elif transform == "2moments_ext":62        if len(features.shape) == 3:63            features = features[:,0,:]64        if features.shape[0] % 3 != 0:65            features = features[:-(features.shape[0]%3)]66        features = np.reshape(features,(-1,3*features.shape[1]))67        C = np.dot(features.T,features)/features.shape[0]68        m = np.mean(features,0)69        pickle.dump((m,C), open(data_path.joinpath(feature_name+'_2moments_ext.pkl'), 'wb'))70    elif transform == "scaler":71        scaler = preprocessing.StandardScaler().fit(features)72        pickle.dump(scaler, open(data_path.joinpath(feature_name+'_scaler.pkl'), 'wb'))73    elif transform == "pca_transform":74        feature_size = features.shape[1]75        pca = decomposition.PCA(n_components=feature_size)76        pca_transform = pca.fit(features)77        pickle.dump(pca_transform, open(data_path.joinpath(feature_name+'_pca_transform.pkl'), 'wb'))78    else:79        raise NotImplementedError("Transform type "+transform+" not implemented")80