librosa_vggish.py

#!/usr/bin/env python
'''Compute VGGish features for a batch of files'''

import argparse

import os
import sys

import librosa
import pandas as pd
import numpy as np
import tensorflow as tf
from tqdm import tqdm

import vggish_input
import vggish_slim
import vggish_postprocess
import vggish_params


MODEL_PARAMS = 'vggish_model.ckpt'
PCA_PARAMS = 'vggish_pca_params.npz'


def load_input(filename):

    y, sr = librosa.load(filename, sr=vggish_params.SAMPLE_RATE, mono=True)
    y = librosa.util.normalize(y)

    return vggish_input.waveform_to_examples(y, sr)


def run_model(files_in, outpath):

    pproc = vggish_postprocess.Postprocessor(PCA_PARAMS)

    with tf.Graph().as_default(), tf.Session() as sess:

        vggish_slim.define_vggish_slim(training=False)
        vggish_slim.load_vggish_slim_checkpoint(sess, MODEL_PARAMS)
        features_tensor = sess.graph.get_tensor_by_name(vggish_params.INPUT_TENSOR_NAME)
        embedding_tensor = sess.graph.get_tensor_by_name(vggish_params.OUTPUT_TENSOR_NAME)

        for file_in in tqdm(files_in):

            file_out = os.path.join(outpath, os.path.extsep.join([os.path.basename(file_in), 'npz']))
            input_data = load_input(file_in)

            [embedding] = sess.run([embedding_tensor], feed_dict={features_tensor: input_data})

            emb_pca = pproc.postprocess(embedding)

            np.savez(file_out, features=embedding, features_z=emb_pca)


def process_args(args):

    parser = argparse.ArgumentParser(description='VGGish feature extractor')

    parser.add_argument(dest='input_list', action='store',
                        type=str, help='Path to input CSV file')

    parser.add_argument(dest='output_path', type=str, action='store',
                        help='Path to store output files in NPZ format')

    return parser.parse_args(args)


def load_files_in(input_list):

    files_in = pd.read_table(input_list, header=None)
    return list(files_in[0])


if __name__ == '__main__':
    args = process_args(sys.argv[1:])

    files_in = load_files_in(args.input_list)

    run_model(files_in, args.output_path)