animate_octave_chroma_matrix.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_chromgram_harmonic_in_pitch.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))

octave_count = 115 // 12
X_octave_chromagram_harmonic = X_pitchgram_harmonic[:, :octave_count*12].reshape(-1, octave_count, 12)

# X_octave_chromagram_harmonic = medfilt(X_octave_chromagram, (15, 1, 1))
# X_octave_chromagram_percussive = medfilt(X_octave_chromagram, (1, 15))

data = X_octave_chromagram_harmonic
print(data.shape)

frame_count = len(data)
duration = frame_count / fps

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], interpolation='nearest', cmap='gray', vmin=0, vmax=1)
ax.set_xticks(np.arange(12))
ax.set_xticklabels(tone_labels[relative_idx])

def make_frame_mpl(t):
    i = int(t * fps)
    imshow_plot.set_data(data[i, :, idx].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)