sigsep / bsseval Goto Github PK

import musdb
import museval

import numpy as np


def estimate_and_evaluate1(track):
    """ Simple baseline system using mixture as estimate """
    estimates = {}
    estimates['vocals'] = 0.25 * track.audio
    estimates['accompaniment'] = 0.75 * track.audio

    scores = museval.eval_mus_track(track, estimates, output_dir='.')
    print('Score for `estimate_and_evaluate1`:')
    print(scores)

    return estimates


def estimate_and_evaluate2(track):#
    """ Modified baseline system, which sets the second half of `vocals` to zero """
    estimates = {}
    estimates['vocals'] = 0.25 * track.audio
    estimates['accompaniment'] = 0.75 * track.audio

    estimates['vocals'] *= np.vstack((np.ones((track.audio.shape[0] // 2, 2)),
                                      np.zeros((track.audio.shape[0] - track.audio.shape[0] // 2, 2))))

    scores = museval.eval_mus_track(track, estimates, output_dir='.')
    print('Score for `estimate_and_evaluate2`:')
    print(scores)

    return estimates


def estimate_and_evaluate3(track):#
    """ Modified baseline system, which sets the first half of `vocals` to zero """
    estimates = {}
    estimates['vocals'] = 0.25 * track.audio
    estimates['accompaniment'] = 0.75 * track.audio

    estimates['vocals'] *= np.vstack((np.zeros((track.audio.shape[0] // 2, 2)),
                                      np.ones((track.audio.shape[0] - track.audio.shape[0] // 2, 2))))

    scores = museval.eval_mus_track(track, estimates, output_dir='.')
    print('Score for `estimate_and_evaluate3`:')
    print(scores)

    return estimates


mus = musdb.DB(root_dir='/speech/db/mul/separ4/sisec/data2018/', is_wav=True)
mus.run(estimate_and_evaluate1, estimates_dir=".", tracks=[mus.load_mus_tracks(subsets='test')[0]])
mus.run(estimate_and_evaluate2, estimates_dir=".", tracks=[mus.load_mus_tracks(subsets='test')[0]])
mus.run(estimate_and_evaluate3, estimates_dir=".", tracks=[mus.load_mus_tracks(subsets='test')[0]])