animate_pitchgram_window.py

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import moviepy.editor as mpy
from moviepy.video.io.bindings import mplfig_to_npimage
import numpy as np
from scipy.signal import medfilt
import tfr

# --- parameters ---
audio_file = '../data/green-sleeves-intro.flac'
video_file = "green-sleeves-intro_pitchgram_window.mp4"
fps = 30

# ------

signal_frames = tfr.SignalFrames(audio_file, frame_size=4096, hop_size=512)

fs = signal_frames.sample_rate
output_frame_size = fs / fps

X_pitchgram = tfr.pitchgram(signal_frames, output_frame_size=output_frame_size, magnitudes='power_db_normalized')

# X_pitchgram_harmonic = medfilt(X_pitchgram, (15, 1))

window_size = 200

data = np.vstack([
    np.zeros((window_size//2, X_pitchgram.shape[1])),
    X_pitchgram,
    np.zeros((window_size//2, X_pitchgram.shape[1]))])
print(data.shape)

frame_count = len(X_pitchgram)
duration = frame_count / fps
print(frame_count, duration)

# tone_labels = np.array(['C', 'Db', 'D', 'Eb', 'E', 'F', 'Gb', 'G', 'Ab', 'A', 'Bb', 'B'])

# key = 3
# fifths = True
# step = 7 if fifths else 1
# idx = ((step * np.arange(12) + key + 12)) % 12
# relative_idx = (step * np.arange(12)) % 12
#
# print(idx, relative_idx)

fig = plt.figure(figsize=(10, 6), facecolor='white')
ax = fig.add_subplot(111)
imshow_plot = ax.imshow(data[0:window_size].T,
    interpolation='nearest', cmap='gray', vmin=0, vmax=1, origin='lower')
ax.axvline(window_size//2, color='r')
ax.set_xlabel('time')
ax.set_ylabel('pitch')
ax.tick_params(axis='x', which='both', bottom='off', top='off', labelbottom='off')
# ax.set_xticks(np.arange(12))
# ax.set_xticklabels(tone_labels[relative_idx])

def make_frame_mpl(t):
    # round() to prevent floating point errors that would lead to skipping back
    # a frame
    i = int(round(t * fps))
    imshow_plot.set_data(data[i:i+window_size].T)
    return mplfig_to_npimage(fig)

animation = mpy.VideoClip(make_frame_mpl, duration=duration)
animation.audio = mpy.AudioFileClip(audio_file)
animation.write_videofile(video_file, fps=fps)