TastyPiano / src /music /pipeline /audio2piano_solo_prob.py
Cédric Colas
initial commit
e775f6d
raw
history blame
1.78 kB
import numpy as np
import librosa
import sys
sys.path.append('../../../data/')
from src.music.utilities.processing_models import piano_detection_model
from src.music.config import CHKPT_PATH_PIANO_EVAL
PIANO_SOLO_DETECTOR = piano_detection_model.PianoSoloDetector(CHKPT_PATH_PIANO_EVAL)
exclude_playlist_folders = ['synth_audio_recorded', 'from_url']
def clean_start_and_end_blanks(probs):
if len(probs) > 20:
# clean up to 10s in each direction
n_zeros_start = 0
for i in range(10):
if probs[i] <= 0.001:
n_zeros_start += 1
else:
break
n_zeros_end = 0
for i in range(10):
if probs[-(i + 1)] <= 0.001:
n_zeros_end += 1
else:
break
if n_zeros_end == 0:
return probs[n_zeros_start:]
else:
return probs[n_zeros_start:-n_zeros_end]
else:
return probs
def calculate_piano_solo_prob(audio_path, verbose=False):
"""Calculate the piano solo probability of all downloaded mp3s, and append
the probability to the meta csv file. Code from https://github.com/bytedance/GiantMIDI-Piano
"""
try:
error_msg = 'Error in audio loading?'
(audio, _) = librosa.core.load(audio_path, sr=piano_detection_model.SR, mono=True)
error_msg += ' Nope. Error in solo prediction?'
probs = PIANO_SOLO_DETECTOR.predict(audio)
# probs = clean_start_and_end_blanks(probs) # remove blanks at start and end (<=10s each way). If not piano, the rest of the song will be enough to tell.
piano_solo_prob = np.mean(probs)
error_msg += ' Nope. '
return piano_solo_prob, ''
except:
return None, error_msg + 'Yes.'