Spaces:
Runtime error
Runtime error
| import librosa | |
| import numpy as np | |
| from pathlib import Path | |
| import json | |
| import os.path | |
| import sys | |
| import argparse | |
| THIS_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| ROOT_DIR = os.path.abspath(os.path.join(THIS_DIR, os.pardir)) | |
| sys.path.append(ROOT_DIR) | |
| from audio_feature_utils import * | |
| from utils import distribute_tasks | |
| parser = argparse.ArgumentParser(description="Preprocess audio data") | |
| parser.add_argument("data_path", type=str, help="Directory contining Beat Saber level folders") | |
| parser.add_argument("--feature_names", metavar='', type=str, default="mel", help="mel, chroma, multi_mel, spectralflux, madmombeats. Comma separated") | |
| parser.add_argument("--combined_feature_name", metavar='', type=str, default=None, help="name for the combined features, if several") | |
| parser.add_argument("--audio_format", type=str, default="mp3") | |
| parser.add_argument("--mel_feature_size", metavar='', type=int, default=None) | |
| # parser.add_argument("--step_size", metavar='', type=float, default=0.01666666666) | |
| parser.add_argument("--fps", metavar='', type=float, default=60) | |
| parser.add_argument("--sampling_rate", metavar='', type=float, default=96000) | |
| parser.add_argument("--replace_existing", action="store_true") | |
| parser.add_argument("--notranspose", action="store_true") | |
| args = parser.parse_args() | |
| # makes arugments into global variables of the same name, used later in the code | |
| globals().update(vars(args)) | |
| step_size=1.0/fps | |
| data_path = Path(data_path) | |
| feature_names = feature_names.split(",") | |
| if len(feature_names) > 1 and combined_feature_name is None: | |
| combined_feature_name = "_".join(feature_names) | |
| elif len(feature_names) == 1: | |
| combined_feature_name = feature_names[0] | |
| ## distributing tasks accross nodes ## | |
| from mpi4py import MPI | |
| comm = MPI.COMM_WORLD | |
| rank = comm.Get_rank() | |
| size = comm.Get_size() | |
| print(rank) | |
| print("creating {} of size {}".format(",".join(feature_names), mel_feature_size)) | |
| #assuming mp3 for now. | |
| candidate_audio_files = sorted(data_path.glob('**/*.'+audio_format), key=lambda path: path.parent.__str__()) | |
| tasks = distribute_tasks(candidate_audio_files,rank,size) | |
| for i in tasks: | |
| path = candidate_audio_files[i] | |
| song_file_path = path.__str__() | |
| # feature files are going to be saved as numpy files | |
| if feature_names == ["mel"] or feature_names == ["multi_mel"]: | |
| features_file = song_file_path+"_"+combined_feature_name+"_"+str(mel_feature_size)+".npy" | |
| else: | |
| features_file = song_file_path+"_"+combined_feature_name+".npy" | |
| if replace_existing or not os.path.isfile(features_file): | |
| print("creating feature file",i) | |
| featuress = [] | |
| for feature_name in feature_names: | |
| # get song | |
| y_wav, sr = librosa.load(song_file_path, sr=sampling_rate) | |
| sr = sampling_rate | |
| hop = int(round(sr * step_size)) | |
| # hop = int(sr * step_size) | |
| #get feature | |
| if feature_name == "chroma": | |
| features = extract_features_hybrid(y_wav,sr,hop).transpose(1,0)[1:] | |
| elif feature_name == "mel": | |
| features = extract_features_mel(y_wav,sr,hop,mel_dim=mel_feature_size).transpose(1,0)[1:] | |
| elif feature_name == "envelope": | |
| features = extract_features_envelope(y_wav,sr,hop)[1:] | |
| elif feature_name == "multi_mel": | |
| features = extract_features_multi_mel(y_wav, sr=sampling_rate, hop=hop, nffts=[1024,2048,4096], mel_dim=mel_feature_size) | |
| elif feature_name == "spectralflux": #actually this is the same as envelope I think | |
| features = extract_features_spectral_flux(song_file_path,fps) | |
| elif feature_name == "madmombeats": | |
| features = extract_features_madmombeat(song_file_path,fps) | |
| featuress.append(features) | |
| shortest_length = 99999999999 | |
| for feat in featuress: | |
| if feat.shape[0] < shortest_length: | |
| shortest_length = feat.shape[0] | |
| for i in range(len(featuress)): | |
| featuress[i] = featuress[i][:shortest_length] | |
| featuress = np.concatenate(featuress,1) | |
| np.save(features_file,featuress) | |