CLI tool for building a tempogram of audio file

tempogram.py

import numpy as np
import librosa
import librosa.display
import matplotlib
import matplotlib.pyplot as plt
import time

def draw_tempogram(audio_in, audio_out, offset, duration):
  """ Make PNG file for specified MP3 file using LibRosa tools.
  """

  checkpoint = start_time = time.time()
  print("Loading %s..." % audio_in)
  y, sr = librosa.load(audio_in, offset=offset, duration=duration)
  print("OK in %s seconds" % (time.time() - checkpoint))

  checkpoint = time.time()
  print("Building tempogram...")
  hop_length = 512
  oenv = librosa.onset.onset_strength(y=y, sr=sr, hop_length=hop_length)
  tempogram = librosa.feature.tempogram(onset_envelope=oenv, sr=sr, hop_length=hop_length)
  tempo = librosa.beat.tempo(onset_envelope=oenv, sr=sr, hop_length=hop_length)[0]
  plt.figure(figsize=(16, 4))
  librosa.display.specshow(tempogram, sr=sr, hop_length=hop_length, x_axis='time', y_axis='tempo')
  plt.axhline(tempo, color='w', linestyle='--', alpha=1, label='Estimated tempo={:g}'.format(tempo))
  plt.legend(frameon=True, framealpha=0.75)
  plt.grid()
  print("OK in %s seconds" % (time.time() - checkpoint))

  checkpoint = time.time()
  print("Saving PNG as %s..." % audio_out)
  plt.savefig(audio_out)
  print("OK in %s seconds" % (time.time() - checkpoint))

  print("Done in %s seconds." % (time.time() - start_time))

if __name__ == "__main__":
  import argparse
  parser = argparse.ArgumentParser("tempogram")
  optional = parser._action_groups.pop()
  required = parser.add_argument_group("required arguments")
  required.add_argument("-i", "--infile", type=argparse.FileType('r'), help="Audio file that needs to proceed.", required=True)
  optional.add_argument("-o", "--outfile", type=argparse.FileType('w'), help="Desired PNG name (based on infile by default).")
  optional.add_argument("-s", "--offset", type=int, default=0, help="From where to start processing in seconds.")
  optional.add_argument("-t", "--duration", type=int, help="How much seconds to process.")
  parser._action_groups.append(optional)
  args = parser.parse_args()
  infile = args.infile.name
  if args.outfile is None:
    outfile = infile + ".png"
  else:
    outfile = args.outfile.name
  draw_tempogram(infile, outfile, args.offset, args.duration)