import sys
import librosa
import numpy
import scipy
import matplotlib.pyplot as plt

if len(sys.argv) < 2:
  print "python spectrum.py file.wav"
  sys.exit()

FILENAME = sys.argv[1]

N_FFT = 2048

def read_audio_spectum(filename):
    fs, x = scipy.io.wavfile.read(filename)
    S = librosa.stft(numpy.transpose(x)[0], N_FFT)
    S = numpy.log1p(numpy.abs(S))
    return S, fs

a_spectrum, fs = read_audio_spectum(FILENAME)
a_spectrum = numpy.flipud(a_spectrum)

print a_spectrum.shape

fig = plt.figure(figsize=(5, int(len(a_spectrum[0])/100)))

img = plt.imshow(a_spectrum)
img.set_cmap('hot')
plt.set_cmap('hot')
plt.axis('off')
plt.savefig(FILENAME + '.png', transparent=True, bbox_inches='tight', pad_inches=0)