align-to-score.py

import json
import argparse, os, sys
from pprint import pprint
import numpy as np
import pandas as pd
import librosa.display
import matplotlib.pyplot as plt
import scipy.interpolate
import pretty_midi
from libfmp.b.b_plot import plot_signal, plot_chromagram
from libfmp.c3.c3s2_dtw_plot import plot_matrix_with_points
import libfmp.c1

from synctoolbox.dtw.mrmsdtw import sync_via_mrmsdtw
from synctoolbox.dtw.utils import compute_optimal_chroma_shift, shift_chroma_vectors, make_path_strictly_monotonic, evaluate_synchronized_positions
from synctoolbox.feature.csv_tools import read_csv_to_df, df_to_pitch_features, df_to_pitch_onset_features
from synctoolbox.feature.chroma import pitch_to_chroma, quantize_chroma, quantized_chroma_to_CENS
from synctoolbox.feature.dlnco import pitch_onset_features_to_DLNCO
from synctoolbox.feature.pitch import audio_to_pitch_features
from synctoolbox.feature.pitch_onset import audio_to_pitch_onset_features
from synctoolbox.feature.utils import estimate_tuning

# synctool defaults
Fs = 22050
feature_rate = 50
step_weights = np.array([1.5, 1.5, 2.0])
threshold_rec = 10 ** 6

# other hardwired defaults
midiTempo = 120 # as generated by Verovio

def bulk_align(files, ref_ix):
    audios = [librosa.load(f, sr=Fs)[0] for f in files]
    #features_df = pd.DataFrame({"audio_ix":{}, "chroma":{}, "onsets":{}})
    features = []
    # generate reference annotations
   # ref_isochronic_annotations = np.arange(0, librosa.get_duration(audios[ref_ix]), 0.1)
    ref_isochronic_annotations = np.arange(0, librosa.get_duration(audios[ref_ix]), 0.02)
    import libfmp.c2

    tuning_offset1 = estimate_tuning(audios[0], Fs)
    tuning_offset2 = estimate_tuning(audios[1], Fs)

    chroma1, onsets1 = get_features_from_audio(audios[0], tuning_offset1, Fs, feature_rate)
    chroma2, onsets2 = get_features_from_audio(audios[1], tuning_offset2, Fs, feature_rate)
    wp = sync_via_mrmsdtw(f_chroma1=chroma1, f_onset1=onsets1, f_chroma2=chroma2, f_onset2=onsets2, input_feature_rate=feature_rate, step_weights=step_weights, threshold_rec=threshold_rec, verbose=False, dtw_implementation="librosa")

    for ix, audio in enumerate(audios):
        #generate tuning offset df
        tuning_offset = estimate_tuning(audio, Fs)
#        tuning_offsets_df.insert(ix, "audio_" + str(ix), [tuning_offset])
        # calculate chroma and onset features
        f_chroma_quantized, f_DLNCO = get_features_from_audio(audio, tuning_offset, Fs, feature_rate)
        features.append({'chroma':f_chroma_quantized, 'onsets':f_DLNCO})
        #'features_df.loc[ix] = ["audio_"+str(ix), f_chroma_quantized_audio, f_DLNCO_audio]
    # generate all warp paths to reference
#    ref_chroma = features_df.iloc[ref_ix]['chroma']
#    ref_onsets = features_df.iloc[ref_ix]['onsets']
    ref_chroma = features[ref_ix]['chroma']
    ref_onsets = features[ref_ix]['onsets']
    annotations_map = dict()
    for ix, _ in enumerate(features):
        chroma = features[ix]['chroma']
        onsets = features[ix]['onsets']
        wp = sync_via_mrmsdtw(f_chroma1=ref_chroma, f_onset1=ref_onsets, f_chroma2=chroma, f_onset2=onsets, input_feature_rate=feature_rate, step_weights=step_weights, threshold_rec=threshold_rec, verbose=False, dtw_implementation="librosa")
        # make wp strictly monotonic (better for transferring annotations according to synctoolbox)
        wp = make_path_strictly_monotonic(wp)
        # transfer reference annotations
        transferred_anno = scipy.interpolate.interp1d(wp[0] / feature_rate, wp[1] / feature_rate, kind='linear')(ref_isochronic_annotations)
        annotations_map[files[ix]] = np.ndarray.tolist(transferred_anno)
    return annotations_map


def get_features_from_audio(audio, tuning_offset, Fs, feature_rate, visualize=False):
    f_pitch = audio_to_pitch_features(f_audio=audio, Fs=Fs, tuning_offset=tuning_offset, feature_rate=feature_rate, verbose=visualize)
    f_chroma = pitch_to_chroma(f_pitch=f_pitch)
    f_chroma_quantized = quantize_chroma(f_chroma=f_chroma)
#    if visualize:
#        plot_chromagram(f_chroma_quantized, title='Quantized chroma features - Audio', Fs=feature_rate, figsize=(9,3))

    f_pitch_onset = audio_to_pitch_onset_features(f_audio=audio, Fs=Fs, tuning_offset=tuning_offset, verbose=visualize)
    f_DLNCO = pitch_onset_features_to_DLNCO(f_peaks=f_pitch_onset, feature_rate=feature_rate, feature_sequence_length=f_chroma_quantized.shape[1], visualize=visualize)
    return f_chroma_quantized, f_DLNCO

def get_features_from_annotation(df_annotation, feature_rate):
    f_pitch = df_to_pitch_features(df_annotation, feature_rate=feature_rate)
    f_chroma = pitch_to_chroma(f_pitch=f_pitch)
    f_chroma_quantized = quantize_chroma(f_chroma=f_chroma)
    f_pitch_onset = df_to_pitch_onset_features(df_annotation)
    f_DLNCO = pitch_onset_features_to_DLNCO(f_peaks=f_pitch_onset,
                                            feature_rate=feature_rate,
                                            feature_sequence_length=f_chroma_quantized.shape[1],
                                            visualize=False)
    
    return f_chroma_quantized, f_DLNCO

def score_align(audio, midi):
    # convert MIDI to libfmp format
    
    # load audio
    audio, _ = librosa.load(audio, Fs)

    # create annotation dataframe for score
    midi_data = pretty_midi.PrettyMIDI(midi)
    midi_list = []

    for instrument in midi_data.instruments:
        for note in instrument.notes:
            start = note.start
            duration = note.end - note.start
            pitch = note.pitch
            velocity = note.velocity
            midi_list.append([start, duration, pitch, velocity, instrument.name])
            
    midi_list = sorted(midi_list, key=lambda x: (x[0], x[2]))

    score_annotation = pd.DataFrame(midi_list, columns=['start', 'duration', 'pitch', 'velocity', 'instrument'])
    score_annotation = score_annotation[score_annotation['pitch'] != 0] # Filter out percussion. TODO is there a better way that doesn't lose this information?
    
    # Estimate tuning
    tuning_offset = estimate_tuning(audio, Fs)

    # get audio features
    chroma_audio, onsets_audio = get_features_from_audio(audio, tuning_offset, Fs, feature_rate)
                                        
    pprint(score_annotation)
    # get annotation features
    chroma_anno, onsets_anno = get_features_from_annotation(score_annotation, feature_rate)

    # determine optimal chroma shift (in case performance is in different key to score)
    f_cens_1hz_audio = quantized_chroma_to_CENS(chroma_audio, 201, 50, feature_rate)[0]
    f_cens_1hz_annotation = quantized_chroma_to_CENS(chroma_anno, 201, 50, feature_rate)[0]
    opt_chroma_shift = compute_optimal_chroma_shift(f_cens_1hz_audio, f_cens_1hz_annotation)

    f_chroma_quantized_annotation = shift_chroma_vectors(chroma_anno, opt_chroma_shift)
    onsets_anno = shift_chroma_vectors(onsets_anno, opt_chroma_shift)


    # do the warp
    wp = sync_via_mrmsdtw(f_chroma1=chroma_audio, 
                        f_onset1=onsets_audio, 
                        f_chroma2=chroma_anno, 
                        f_onset2=onsets_anno, 
                        input_feature_rate=feature_rate, 
                        step_weights=step_weights, 
                        threshold_rec=threshold_rec, 
                        verbose=False)

    # make warping path strictly monotonic (better for transferring annotations accordign to synctoolbox)
    wp = make_path_strictly_monotonic(wp)

    score_annotation["warped-end"] = score_annotation["start"] + score_annotation["duration"]
    score_annotation[['warped-start', 'warped-end']] = scipy.interpolate.interp1d(wp[1] / feature_rate, 
                            wp[0] / feature_rate, kind='linear', fill_value="extrapolate")(score_annotation[['start', 'end']])
    score_annotation["warped-duration"] = score_annotation["warped-end"] - score_annotation["warped-start"]

    # deduplicate
    onsets = score_annotation[["start", "warped-start"]].drop_duplicates()
    offsets = score_annotation[["end", "warped-end"]].drop_duplicates()

    return {
        "score_onset": onsets["start"].tolist(),
        "ref_onset": onsets["warped-start"].tolist(),
        "score_offset": offsets["end"].tolist(),
        "ref_offset": offsets["warped-end"].tolist()
    }

def alignment_ix_to_score_time(alignment_grids, ref_name, ix):
    return alignment_grids[ref_name][ix] / ((60 / midiTempo) * 1000)

def calculate_tempi(alignment_grids, ref_name):
    tempi = dict()
    for file_name, grid in alignment_grids.items():
        print("## CALCULATING TEMPO FOR: " + file_name)
        tempi[file_name] = list()
        for ix, perf_time in enumerate(grid):
            if ix == 0:
                tempi[file_name].append(0) # 0 bmp at very start of track
            else:
                delta_perf_time = perf_time - grid[ix-1]
                delta_score_time = alignment_ix_to_score_time(alignment_grids, ref_name, ix) - alignment_ix_to_score_time(alignment_grids, ref_name, ix-1)
                tempi[file_name].append((delta_score_time / delta_perf_time) * 60000) # beats per minute
    return tempi
            


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--audioDirectory', help="Directory containing recordings to be matched", required=True)
    parser.add_argument('-r', '--referenceAudio', help="Filename of reference recording in directory. If omited, I'll use the first file I find", required=False)
    parser.add_argument('-m', '--referenceMidi', help="Filename of reference MIDI in directory.", required=True)
    parser.add_argument('-u', '--meiUri', help="URI of MEI file used to generate reference MIDI.", required=True)
    parser.add_argument('-o', '--output', help="Filename for output file", required=True)
    args = parser.parse_args()
   
    audio_files = [f for f in os.listdir(args.audioDirectory) if f.endswith('.wav') or f.endswith('.mp3')]
    ref_index = 0
    if len(audio_files) < 2:
        sys.exit("Specified audio directory must contain at least two audio (.wav or .mp3) files")
    if args.referenceAudio:  
        if args.referenceAudio not in audio_files:
            sys.exit("Cannot find specified reference audio in specified audio directory")
        else: 
            ref_index = audio_files.index(args.referenceAudio)
    files = [os.path.join(args.audioDirectory, f) for f in audio_files]
    print(files)
    annotations_map = dict()
    annotations_map["body"] = dict()
    annotations_map["body"]["audio"] = bulk_align(files, ref_index)
    annotations_map["body"]["score"] = score_align(os.path.join(args.audioDirectory, audio_files[ref_index]), args.referenceMidi)
    annotations_map["header"] = dict()
    annotations_map["header"]["ref"] = os.path.join(args.audioDirectory, audio_files[ref_index])
    annotations_map["header"]["meiUri"] = args.meiUri
    #annotations_map["header"]["tempi"] = calculate_tempi(annotations_map["body"]["audio"], files[ref_index])
    with open(args.output, 'w',encoding="utf-8") as out:
        json.dump(annotations_map, out, indent = 2)
    print(files)