Modularizing code to compute features

src/compute_audio_features.py

@@ -1,15 +1,22 @@
 """Script to compute audio features from the
 original Harmonix audio files.
+
+Created by Oriol Nieto.
 """
 
 
-import os
+import argparse
 import glob
-import librosa
-import numpy as np
 import json
+import os
+import time
+import numpy as np
+
 from joblib import Parallel, delayed
 
+import librosa
+
+
 INPUT_DIR = "mp3s"
 OUTPUT_DIR = "audio_features"
 OUT_JSON = "info.json"
@@ -35,8 +42,8 @@ def compute_melspecs(audio):
  fmax=MEL_FMAX)
 
 
-def compute_all_features(mp3_file):
- """Computes all the audio features"""
+def compute_all_librosa_features(mp3_file, output_dir):
+ """Computes all the audio features with librosa"""
  # Decode and read mp3
  audio, _ = librosa.load(mp3_file, sr=SR)
 
@@ -45,13 +52,13 @@ def compute_all_features(mp3_file):
 
  # Save
  out_file = os.path.join(
- OUTPUT_DIR, os.path.basename(mp3_file).replace(".mp3", "-mel.npy"))
+ output_dir, os.path.basename(mp3_file).replace(".mp3", "-mel.npy"))
  np.save(out_file, mel)
 
 
-def save_params():
+def save_params(output_dir):
  """Saves the parameters to a JSON file."""
- out_json = os.path.join(OUTPUT_DIR, OUT_JSON)
+ out_json = os.path.join(output_dir, OUT_JSON)
  out_dict = {
  "librosa_version": librosa.__version__,
  "numpy_version": np.__version__,
@@ -67,16 +74,44 @@ def save_params():
 
 
 if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="Computes audio features for the Harmonix set.",
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+ parser.add_argument("-i",
+ "--input_dir",
+ default=INPUT_DIR,
+ action="store",
+ help="Path to the Harmonix set audio.")
+ parser.add_argument("-o",
+ "--output_dir",
+ default=OUTPUT_DIR,
+ action="store",
+ help="Output directory.")
+ parser.add_argument("-j",
+ "--n_jobs",
+ default=N_JOBS,
+ action="store",
+ type=int,
+ help="Number of jobs to run in parallel.")
+
+ args = parser.parse_args()
+ start_time = time.time()
+
  # Create output dir if doesn't exist
- if not os.path.exists(OUTPUT_DIR):
- os.makedirs(OUTPUT_DIR)
+ if not os.path.exists(args.output_dir):
+ os.makedirs(args.output_dir)
 
  # Read mp3s
- mp3s = glob.glob(os.path.join(INPUT_DIR, "*.mp3"))
+ mp3s = glob.glob(os.path.join(args.input_dir, "*.mp3"))
 
  # Compute features for each mp3 in parallel
- Parallel(n_jobs=N_JOBS)(
- delayed(compute_all_features)(mp3_file) for mp3_file in mp3s)
+ Parallel(n_jobs=args.n_jobs)(
+ delayed(compute_all_librosa_features)(mp3_file, args.output_dir)
+ for mp3_file in mp3s)
 
  # Save parameters
- save_params()
+ save_params(args.output_dir)
+
+ # Done!
+ print("Done! Took %.2f seconds." % (time.time() - start_time))