python/demos/demo_spectrogram.py

   1 #! /usr/bin/env python
   2
   3 import sys, os.path
   4 from aubio import pvoc, source
   5 from numpy import array, arange, zeros, log10, vstack
   6 import matplotlib.pyplot as plt
   7
   8 def get_spectrogram(filename, samplerate = 0):
   9     win_s = 512                                        # fft window size
  10     hop_s = win_s / 2                                  # hop size
  11     fft_s = win_s / 2 + 1                              # spectrum bins
  12
  13     a = source(filename, samplerate, hop_s)            # source file
  14     if samplerate == 0: samplerate = a.samplerate
  15     pv = pvoc(win_s, hop_s)                            # phase vocoder
  16     specgram = zeros([0, fft_s], dtype='float32')      # numpy array to store spectrogram
  17
  18     # analysis
  19     while True:
  20         samples, read = a()                              # read file
  21         specgram = vstack((specgram,pv(samples).norm))   # store new norm vector
  22         if read < a.hop_size: break
  23
  24     # plotting
  25     fig = plt.imshow(log10(specgram.T + .001), origin = 'bottom', aspect = 'auto', cmap=plt.cm.gray_r)
  26     ax = fig.axes
  27     ax.axis([0, len(specgram), 0, len(specgram[0])])
  28     # show axes in Hz and seconds
  29     time_step = hop_s / float(samplerate)
  30     total_time = len(specgram) * time_step
  31     print "total time: %0.2fs" % total_time,
  32     print ", samplerate: %.2fkHz" % (samplerate / 1000.)
  33     n_xticks = 10
  34     n_yticks = 10
  35
  36     def get_rounded_ticks( top_pos, step, n_ticks ):
  37         top_label = top_pos * step
  38         # get the first label
  39         ticks_first_label = top_pos * step / n_ticks
  40         # round to the closest .1
  41         ticks_first_label = round ( ticks_first_label * 10. ) / 10.
  42         # compute all labels from the first rounded one
  43         ticks_labels = [ ticks_first_label * n for n in range(n_ticks) ] + [ top_label ]
  44         # get the corresponding positions
  45         ticks_positions = [ ticks_labels[n] / step for n in range(n_ticks) ] + [ top_pos ]
  46         # convert to string
  47         ticks_labels = [  "%.1f" % x for x in ticks_labels ]
  48         # return position, label tuple to use with x/yticks
  49         return ticks_positions, ticks_labels
  50
  51     # apply to the axis
  52     x_ticks, x_labels = get_rounded_ticks ( len(specgram), time_step, n_xticks )
  53     y_ticks, y_labels = get_rounded_ticks ( len(specgram[0]), (samplerate / 1000. / 2.) / len(specgram[0]), n_yticks )
  54     ax.set_xticks( x_ticks )
  55     ax.set_yticks ( y_ticks )
  56     ax.set_xticklabels( x_labels )
  57     ax.set_yticklabels ( y_labels )
  58     ax.set_ylabel('Frequency (kHz)')
  59     ax.set_xlabel('Time (s)')
  60     ax.set_title(os.path.basename(soundfile))
  61     for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] +
  62             ax.get_xticklabels() + ax.get_yticklabels()):
  63         item.set_fontsize('x-small')
  64     return fig
  65
  66 if __name__ == '__main__':
  67     if len(sys.argv) < 2:
  68         print "Usage: %s <filename>" % sys.argv[0]
  69   else:
  70       for soundfile in sys.argv[1:]:
  71           fig = get_spectrogram(soundfile)
  72           # display graph
  73           fig.show()
  74           #outimage = os.path.basename(soundfile) + '.png'
  75           #print ("writing: " + outimage)
  76           #plt.savefig(outimage)
  77           plt.close()