source: gs3-extensions/mars-src/trunk/bin/script/essentia-hpcp.py@ 34411

#!/usr/bin/env python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
import os

import essentia
import essentia.standard  as esstandard
import essentia.streaming as esstreaming

# https://stackoverflow.com/questions/9622163/save-plot-to-image-file-instead-of-displaying-it-using-matplotlib
# Statements need to be in this order for MacOS operating a venv python2, otherwise an error message results
# related to the version being used not being a system OS 'Framework'
import matplotlib
#  Select the Anti-Grain Geometry C++ backend for rendering to a file
matplotlib.use('Agg') 
import matplotlib.pyplot as pyplot

argc = len(sys.argv)

input_audio_filename = None
output_filename_root = None

if argc <= 1:
    print("Usage: "+sys.argv[0] +" input_file [output_file]\n",file=sys.stderr)    
    sys.exit(1)
else:
    input_audio_filename = sys.argv[1]
    if argc == 2:
        output_filename_root = os.path.splitext(input_audio_filename)[0]
    else:
#        output_filename_root = sys.argv[2]
        output_filename_root = os.path.splitext(sys.argv[2])[0]



# Initialize algorithms we will use
loader = esstreaming.MonoLoader(filename=input_audio_filename)

framecutter   = esstreaming.FrameCutter(frameSize=4096, hopSize=2048, silentFrames='noise')
windowing     = esstreaming.Windowing(type='blackmanharris62')
spectrum      = esstreaming.Spectrum()
spectralpeaks = esstreaming.SpectralPeaks(orderBy='magnitude',
                                          magnitudeThreshold=0.00001,
                                          minFrequency=20,
                                          maxFrequency=3500,
                                          maxPeaks=60)


# Use default HPCP parameters for plots, however we will need higher resolution
# and custom parameters for better Key estimation

hpcp     = esstreaming.HPCP()
hpcp_key = esstreaming.HPCP(size=36,                # we will need higher resolution for Key estimation
                            referenceFrequency=440, # assume tuning frequency is 44100.
                            bandPreset=False,
                            minFrequency=20,
                            maxFrequency=3500,
                            weightType='cosine',
                            nonLinear=False,
                            windowSize=1.)

key       = esstreaming.Key(profileType='edma',     # Use profile for electronic music
                            numHarmonics=4,
                            pcpSize=36,
                            slope=0.6,
                            usePolyphony=True,
                            useThreeChords=True)

# Use pool to store data
pool = essentia.Pool()

# Connect streaming algorithms
loader.audio              >> framecutter.signal
framecutter.frame         >> windowing.frame        >> spectrum.frame
spectrum.spectrum         >> spectralpeaks.spectrum
spectralpeaks.magnitudes  >> hpcp.magnitudes
spectralpeaks.frequencies >> hpcp.frequencies

spectralpeaks.magnitudes  >> hpcp_key.magnitudes
spectralpeaks.frequencies >> hpcp_key.frequencies
hpcp_key.hpcp             >> key.pcp

hpcp.hpcp     >> (pool, 'tonal.hpcp')

key.key       >> (pool, 'tonal.key_key')
key.scale     >> (pool, 'tonal.key_scale')
key.strength  >> (pool, 'tonal.key_strength')

# Run streaming network
essentia.run(loader)


output_filename_hpcp = output_filename_root + ".png"

# Plot HPCP
# data vals get normalized to be [0..1] 
pyplot.imshow(pool['tonal.hpcp'].T, aspect='auto', origin='lower', interpolation='none')

pyplot.axis('off')
pyplot.savefig(output_filename_hpcp , bbox_inches="tight", pad_inches=0) 


output_filename_hpcp_json = output_filename_root + ".json"
output = esstandard.YamlOutput(filename=output_filename_hpcp_json, format='json') 
output(pool)

# print("Estimated key and scale:", pool['tonal.key_key'] + " " + pool['tonal.key_scale'])