pop2piano / preprocess /pop_align.py
sweetcocoa's picture
initial test
88490a8
raw
history blame
11.1 kB
import librosa
import soundfile as sf
import glob
import os
import copy
import sys
import numpy as np
import pyrubberband as pyrb
import pretty_midi
from omegaconf import OmegaConf
from tqdm.auto import tqdm
from synctoolbox.dtw.mrmsdtw import sync_via_mrmsdtw
from synctoolbox.dtw.utils import (
compute_optimal_chroma_shift,
shift_chroma_vectors,
make_path_strictly_monotonic,
)
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
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
print(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.dsp import normalize, get_stereo
from midiaudiopair import MidiAudioPair
Fs = 22050
feature_rate = 50
step_weights = np.array([1.5, 1.5, 2.0])
threshold_rec = 10 ** 6
def save_delayed_song(
sample,
dry_run,
):
import warnings
warnings.filterwarnings(action="ignore")
song_audio, _ = librosa.load(sample.original_song, Fs)
midi_pm = pretty_midi.PrettyMIDI(sample.original_midi)
if np.power(song_audio, 2).sum() < 1: # low energy: invalid file
print("invalid audio :", sample.original_song)
sample.delete_files_myself()
return
rd = get_aligned_results(midi_pm=midi_pm, song_audio=song_audio)
mix_song = rd["mix_song"]
song_pitch_shifted = rd["song_pitch_shifted"]
midi_warped_pm = rd["midi_warped_pm"]
pitch_shift_for_song_audio = rd["pitch_shift_for_song_audio"]
tuning_offset_song = rd["tuning_offset_song"]
tuning_offset_piano = rd["tuning_offset_piano"]
try:
if dry_run:
print("write audio files: ", sample.song)
else:
sf.write(
file=sample.song,
data=song_pitch_shifted,
samplerate=Fs,
format="wav",
)
except:
print("Fail : ", sample.song)
try:
if dry_run:
print("write warped midi :", sample.midi)
else:
midi_warped_pm.write(sample.midi)
except:
midi_warped_pm._tick_scales = midi_pm._tick_scales
try:
if dry_run:
print("write warped midi2 :", sample.midi)
else:
midi_warped_pm.write(sample.midi)
except:
print("ad-hoc failed midi : ", sample.midi)
print("ad-hoc midi : ", sample.midi)
sample.yaml.song.pitch_shift = pitch_shift_for_song_audio.item()
sample.yaml.song.tuning_offset = tuning_offset_song.item()
sample.yaml.piano.tuning_offset = tuning_offset_piano.item()
OmegaConf.save(sample.yaml, sample.yaml_path)
def get_aligned_results(midi_pm, song_audio):
piano_audio = midi_pm.fluidsynth(Fs)
song_audio = normalize(song_audio)
# The reason for estimating tuning ::
# https://www.audiolabs-erlangen.de/resources/MIR/FMP/C3/C3S1_TranspositionTuning.html
tuning_offset_1 = estimate_tuning(song_audio, Fs)
tuning_offset_2 = estimate_tuning(piano_audio, Fs)
# DLNCO features (Sebastian Ewert, Meinard Müller, and Peter Grosche: High Resolution Audio Synchronization Using Chroma Onset Features, In Proceedings of IEEE International Conference on Acoustics, Speech, and Signal Processing (ICASSP): 1869–1872, 2009.):
# helpful to increase synchronization accuracy, especially for music with clear onsets.
# Quantized and smoothed chroma : CENS features
# Because, MrMsDTW Requires CENS.
f_chroma_quantized_1, f_DLNCO_1 = get_features_from_audio(
song_audio, tuning_offset_1
)
f_chroma_quantized_2, f_DLNCO_2 = get_features_from_audio(
piano_audio, tuning_offset_2
)
# Shift chroma vectors :
# Otherwise, different keys of two audio leads to degradation of alignment.
opt_chroma_shift = compute_optimal_chroma_shift(
quantized_chroma_to_CENS(f_chroma_quantized_1, 201, 50, feature_rate)[0],
quantized_chroma_to_CENS(f_chroma_quantized_2, 201, 50, feature_rate)[0],
)
f_chroma_quantized_2 = shift_chroma_vectors(f_chroma_quantized_2, opt_chroma_shift)
f_DLNCO_2 = shift_chroma_vectors(f_DLNCO_2, opt_chroma_shift)
wp = sync_via_mrmsdtw(
f_chroma1=f_chroma_quantized_1,
f_onset1=f_DLNCO_1,
f_chroma2=f_chroma_quantized_2,
f_onset2=f_DLNCO_2,
input_feature_rate=feature_rate,
step_weights=step_weights,
threshold_rec=threshold_rec,
verbose=False,
)
wp = make_path_strictly_monotonic(wp)
pitch_shift_for_song_audio = -opt_chroma_shift % 12
if pitch_shift_for_song_audio > 6:
pitch_shift_for_song_audio -= 12
if pitch_shift_for_song_audio != 0:
song_audio_shifted = pyrb.pitch_shift(
song_audio, Fs, pitch_shift_for_song_audio
)
else:
song_audio_shifted = song_audio
time_map_second = wp / feature_rate
midi_pm_warped = copy.deepcopy(midi_pm)
midi_pm_warped = simple_adjust_times(
midi_pm_warped, time_map_second[1], time_map_second[0]
)
piano_audio_warped = midi_pm_warped.fluidsynth(Fs)
song_audio_shifted = normalize(song_audio_shifted)
stereo_sonification_piano = get_stereo(song_audio_shifted, piano_audio_warped)
rd = dict(
mix_song=stereo_sonification_piano,
song_pitch_shifted=song_audio_shifted,
midi_warped_pm=midi_pm_warped,
pitch_shift_for_song_audio=pitch_shift_for_song_audio,
tuning_offset_song=tuning_offset_1,
tuning_offset_piano=tuning_offset_2,
)
return rd
def simple_adjust_times(pm, original_times, new_times):
"""
most of these codes are from original pretty_midi
https://github.com/craffel/pretty-midi/blob/main/pretty_midi/pretty_midi.py
"""
for instrument in pm.instruments:
instrument.notes = [
copy.deepcopy(note)
for note in instrument.notes
if note.start >= original_times[0] and note.end <= original_times[-1]
]
# Get array of note-on locations and correct them
note_ons = np.array(
[note.start for instrument in pm.instruments for note in instrument.notes]
)
adjusted_note_ons = np.interp(note_ons, original_times, new_times)
# Same for note-offs
note_offs = np.array(
[note.end for instrument in pm.instruments for note in instrument.notes]
)
adjusted_note_offs = np.interp(note_offs, original_times, new_times)
# Correct notes
for n, note in enumerate(
[note for instrument in pm.instruments for note in instrument.notes]
):
note.start = (adjusted_note_ons[n] > 0) * adjusted_note_ons[n]
note.end = (adjusted_note_offs[n] > 0) * adjusted_note_offs[n]
# After performing alignment, some notes may have an end time which is
# on or before the start time. Remove these!
pm.remove_invalid_notes()
def adjust_events(event_getter):
"""This function calls event_getter with each instrument as the
sole argument and adjusts the events which are returned."""
# Sort the events by time
for instrument in pm.instruments:
event_getter(instrument).sort(key=lambda e: e.time)
# Correct the events by interpolating
event_times = np.array(
[
event.time
for instrument in pm.instruments
for event in event_getter(instrument)
]
)
adjusted_event_times = np.interp(event_times, original_times, new_times)
for n, event in enumerate(
[
event
for instrument in pm.instruments
for event in event_getter(instrument)
]
):
event.time = adjusted_event_times[n]
for instrument in pm.instruments:
# We want to keep only the final event which has time ==
# new_times[0]
valid_events = [
event
for event in event_getter(instrument)
if event.time == new_times[0]
]
if valid_events:
valid_events = valid_events[-1:]
# Otherwise only keep events within the new set of times
valid_events.extend(
event
for event in event_getter(instrument)
if event.time > new_times[0] and event.time < new_times[-1]
)
event_getter(instrument)[:] = valid_events
# Correct pitch bends and control changes
adjust_events(lambda i: i.pitch_bends)
adjust_events(lambda i: i.control_changes)
return pm
def get_features_from_audio(audio, tuning_offset, 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)
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 main(samples, dry_run):
import multiprocessing
from joblib import Parallel, delayed
Parallel(n_jobs=multiprocessing.cpu_count() // 2)(
delayed(save_delayed_song)(sample=sample, dry_run=dry_run)
for sample in tqdm(samples)
)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="piano cover downloader")
parser.add_argument(
"data_dir",
type=str,
default=None,
help="""directory contains {id}/{song_filename.wav}
""",
)
parser.add_argument(
"--dry_run", default=False, action="store_true", help="whether dry_run"
)
args = parser.parse_args()
def getfiles():
meta_files = sorted(glob.glob(args.data_dir + "/*.yaml"))
print("meta ", len(meta_files))
samples = list()
for meta_file in tqdm(meta_files):
m = MidiAudioPair(meta_file, auto_remove_no_song=True)
if m.error_code != MidiAudioPair.NO_SONG:
aux_txt = os.path.join(
m.audio_dir,
m.yaml.piano.ytid,
f"{m.yaml.piano.title[:50]}___{m.yaml.song.title[:50]}.txt",
)
with open(aux_txt, "w") as f:
f.write(".")
samples.append(m)
print(f"files available {len(samples)}")
return samples
samples = getfiles()
main(samples=samples, dry_run=args.dry_run)