File: spleeter/separator.py

#!/usr/bin/env python
# coding: utf8

"""
Module that provides a class wrapper for source separation.

Examples:

```python
>>> from spleeter.separator import Separator
>>> separator = Separator('spleeter:2stems')
>>> separator.separate(waveform, lambda instrument, data: ...)
>>> separator.separate_to_file(...)
```
"""

import atexit
import os
from multiprocessing import Pool
from os.path import basename, dirname, join, splitext
from typing import Any, Dict, Generator, List, Optional

# pyright: reportMissingImports=false
# pylint: disable=import-error
import numpy as np
import tensorflow as tf  # type: ignore

from . import SpleeterError
from .audio import Codec
from .audio.adapter import AudioAdapter
from .audio.convertor import to_stereo
from .model import EstimatorSpecBuilder, InputProviderFactory, model_fn
from .model.provider import ModelProvider
from .types import AudioDescriptor
from .utils.configuration import load_configuration

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


class DataGenerator(object):
    """
    Generator object that store a sample and generate it once while called.
    Used to feed a tensorflow estimator without knowing the whole data at
    build time.
    """

    def __init__(self) -> None:
        """Default constructor."""
        self._current_data = None

    def update_data(self, data) -> None:
        """Replace internal data."""
        self._current_data = data

    def __call__(self) -> Generator:
        """Generation process."""
        buffer = self._current_data
        while buffer:
            yield buffer
            buffer = self._current_data


def create_estimator(params: Dict, MWF: bool) -> tf.Tensor:
    """
    Initialize tensorflow estimator that will perform separation

    Parameters:
        params (Dict):
            A dictionary of parameters for building the model
        MWF (bool):
            Wiener filter enabled?

    Returns:
        tf.Tensor:
            A tensorflow estimator
    """
    # Load model.
    provider: ModelProvider = ModelProvider.default()
    params["model_dir"] = provider.get(params["model_dir"])
    params["MWF"] = MWF
    # Setup config
    session_config = tf.compat.v1.ConfigProto()
    session_config.gpu_options.per_process_gpu_memory_fraction = 0.7
    config = tf.estimator.RunConfig(session_config=session_config)
    # Setup estimator
    estimator = tf.estimator.Estimator(
        model_fn=model_fn, model_dir=params["model_dir"], params=params, config=config
    )
    return estimator


class Separator(object):
    """A wrapper class for performing separation."""

    def __init__(
        self,
        params_descriptor: str,
        MWF: bool = False,
        multiprocess: bool = True,
    ) -> None:
        """
        Default constructor.

        Parameters:
            params_descriptor (str):
                Descriptor for TF params to be used.
            MWF (bool):
                (Optional) `True` if MWF should be used, `False` otherwise.
            multiprocess (bool):
                (Optional) Enable multi-processing.
        """
        self._params = load_configuration(params_descriptor)
        self._sample_rate = self._params["sample_rate"]
        self._MWF = MWF
        self._tf_graph = tf.Graph()
        self._prediction_generator: Optional[Generator] = None
        self._input_provider = None
        self._builder = None
        self._features = None
        self._session = None
        if multiprocess:
            self._pool: Optional[Any] = Pool()
            atexit.register(self._pool.close)
        else:
            self._pool = None
        self._tasks: List = []
        self._data_generator = DataGenerator()

    def _get_prediction_generator(self) -> Generator:
        """
        Lazy loading access method for internal prediction generator
        returned by the predict method of a tensorflow estimator.

        Returns:
            Generator:
                Generator of prediction.
        """
        if self._prediction_generator is None:
            estimator = create_estimator(self._params, self._MWF)

            def get_dataset():
                return tf.data.Dataset.from_generator(
                    self._data_generator,
                    output_types={"waveform": tf.float32, "audio_id": tf.string},
                    output_shapes={"waveform": (None, 2), "audio_id": ()},
                )

            self._prediction_generator = estimator.predict(
                get_dataset, yield_single_examples=False
            )
        return self._prediction_generator

    def join(self, timeout: int = 200) -> None:
        """
        Wait for all pending tasks to be finished.

        Parameters:
            timeout (int):
                (Optional) Task waiting timeout.
        """
        while len(self._tasks) > 0:
            task = self._tasks.pop()
            task.get()
            task.wait(timeout=timeout)

    def _get_input_provider(self):
        if self._input_provider is None:
            self._input_provider = InputProviderFactory.get(self._params)
        return self._input_provider

    def _get_features(self):
        if self._features is None:
            provider = self._get_input_provider()
            self._features = provider.get_input_dict_placeholders()
        return self._features

    def _get_builder(self):
        if self._builder is None:
            self._builder = EstimatorSpecBuilder(self._get_features(), self._params)
        return self._builder

    def _get_session(self):
        if self._session is None:
            saver = tf.compat.v1.train.Saver()
            provider = ModelProvider.default()
            model_directory: str = provider.get(self._params["model_dir"])
            latest_checkpoint = tf.train.latest_checkpoint(model_directory)
            self._session = tf.compat.v1.Session()
            saver.restore(self._session, latest_checkpoint)
        return self._session

    def _separate_tensorflow(
        self, waveform: np.ndarray, audio_descriptor: AudioDescriptor
    ) -> Dict:
        """
        Performs source separation over the given waveform with tensorflow
        backend.

        Parameters:
            waveform (np.ndarray):
                Waveform to be separated (as a numpy array)
            audio_descriptor (AudioDescriptor):
                Audio descriptor to be used.

        Returns:
            Dict:
                Separated waveforms.
        """
        if not waveform.shape[-1] == 2:
            waveform = to_stereo(waveform)
        prediction_generator = self._get_prediction_generator()
        # NOTE: update data in generator before performing separation.
        self._data_generator.update_data(
            {"waveform": waveform, "audio_id": np.array(audio_descriptor)}
        )
        # NOTE: perform separation.
        prediction = next(prediction_generator)
        prediction.pop("audio_id")
        return prediction

    def separate(
        self, waveform: np.ndarray, audio_descriptor: Optional[str] = ""
    ) -> Dict:
        """
        Performs separation on a waveform.

        Parameters:
            waveform (np.ndarray):
                Waveform to be separated (as a numpy array)
            audio_descriptor (Optional[str]):
                (Optional) string describing the waveform (e.g. filename).

        Returns:
            Dict:
                Separated waveforms.
        """
        return self._separate_tensorflow(waveform, audio_descriptor)

    def separate_to_file(
        self,
        audio_descriptor: AudioDescriptor,
        destination: str,
        audio_adapter: Optional[AudioAdapter] = None,
        offset: float = 0,
        duration: float = 600.0,
        codec: Codec = Codec.WAV,
        bitrate: str = "128k",
        filename_format: str = "{filename}/{instrument}.{codec}",
        synchronous: bool = True,
    ) -> None:
        """
        Performs source separation and export result to file using
        given audio adapter.

        Filename format should be a Python formattable string that could
        use following parameters :

        - {instrument}
        - {filename}
        - {foldername}
        - {codec}.

        Parameters:
            audio_descriptor (AudioDescriptor):
                Describe song to separate, used by audio adapter to
                retrieve and load audio data, in case of file based
                audio adapter, such descriptor would be a file path.
            destination (str):
                Target directory to write output to.
            audio_adapter (AudioAdapter):
                (Optional) Audio adapter to use for I/O.
            offset (int):
                (Optional) Offset of loaded song.
            duration (float):
                (Optional) Duration of loaded song (default: 600s).
            codec (Codec):
                (Optional) Export codec.
            bitrate (str):
                (Optional) Export bitrate.
            filename_format (str):
                (Optional) Filename format.
            synchronous (bool):
                (Optional) True is should by synchronous.
        """
        if audio_adapter is None:
            audio_adapter = AudioAdapter.default()
        waveform, _ = audio_adapter.load(
            audio_descriptor,
            offset=offset,
            duration=duration,
            sample_rate=self._sample_rate,
        )
        sources = self.separate(waveform, audio_descriptor)
        self.save_to_file(
            sources,
            audio_descriptor,
            destination,
            filename_format,
            codec,
            audio_adapter,
            bitrate,
            synchronous,
        )

    def save_to_file(
        self,
        sources: Dict,
        audio_descriptor: AudioDescriptor,
        destination: str,
        filename_format: str = "{filename}/{instrument}.{codec}",
        codec: Codec = Codec.WAV,
        audio_adapter: Optional[AudioAdapter] = None,
        bitrate: str = "128k",
        synchronous: bool = True,
    ) -> None:
        """
        Export dictionary of sources to files.

        Parameters:
            sources (Dict):
                Dictionary of sources to be exported. The keys are the name
                of the instruments, and the values are `N x 2` numpy arrays
                containing the corresponding intrument waveform, as
                returned by the separate method
            audio_descriptor (AudioDescriptor):
                Describe song to separate, used by audio adapter to
                retrieve and load audio data, in case of file based audio
                adapter, such descriptor would be a file path.
            destination (str):
                Target directory to write output to.
            filename_format (str):
                (Optional) Filename format.
            codec (Codec):
                (Optional) Export codec.
            audio_adapter (Optional[AudioAdapter]):
                (Optional) Audio adapter to use for I/O.
            bitrate (str):
                (Optional) Export bitrate.
            synchronous (bool):
                (Optional) True is should by synchronous.
        """
        if audio_adapter is None:
            audio_adapter = AudioAdapter.default()
        foldername = basename(dirname(audio_descriptor))
        filename = splitext(basename(audio_descriptor))[0]
        generated = []
        for instrument, data in sources.items():
            path = join(
                destination,
                filename_format.format(
                    filename=filename,
                    instrument=instrument,
                    foldername=foldername,
                    codec=codec,
                ),
            )
            directory = os.path.dirname(path)
            if not os.path.exists(directory):
                os.makedirs(directory)
            if path in generated:
                raise SpleeterError(
                    (
                        f"Separated source path conflict : {path},"
                        "please check your filename format"
                    )
                )
            generated.append(path)
            if self._pool:
                task = self._pool.apply_async(
                    audio_adapter.save, (path, data, self._sample_rate, codec, bitrate)
                )
                self._tasks.append(task)
            else:
                audio_adapter.save(path, data, self._sample_rate, codec, bitrate)
        if synchronous and self._pool:
            self.join()



File: spleeter/options.py

#!/usr/bin/env python
# coding: utf8

"""This modules provides spleeter command as well as CLI parsing methods."""

from os.path import join
from tempfile import gettempdir

from typer import Argument, Exit, Option, echo
from typer.models import List, Optional

from .audio import Codec

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"

AudioInputArgument: List[str] = Argument(
    ...,
    help="List of input audio file path",
    exists=True,
    file_okay=True,
    dir_okay=False,
    readable=True,
    resolve_path=True,
)

AudioInputOption: Optional[str] = Option(
    None, "--inputs", "-i", help="(DEPRECATED) placeholder for deprecated input option"
)

AudioAdapterOption: str = Option(
    "spleeter.audio.ffmpeg.FFMPEGProcessAudioAdapter",
    "--adapter",
    "-a",
    help="Name of the audio adapter to use for audio I/O",
)

AudioOutputOption: str = Option(
    join(gettempdir(), "separated_audio"),
    "--output_path",
    "-o",
    help="Path of the output directory to write audio files in",
)

AudioOffsetOption: float = Option(
    0.0, "--offset", "-s", help="Set the starting offset to separate audio from"
)

AudioDurationOption: float = Option(
    600.0,
    "--duration",
    "-d",
    help=(
        "Set a maximum duration for processing audio "
        "(only separate offset + duration first seconds of "
        "the input file)"
    ),
)

AudioCodecOption: Codec = Option(
    Codec.WAV, "--codec", "-c", help="Audio codec to be used for the separated output"
)

AudioBitrateOption: str = Option(
    "128k", "--bitrate", "-b", help="Audio bitrate to be used for the separated output"
)

FilenameFormatOption: str = Option(
    "{filename}/{instrument}.{codec}",
    "--filename_format",
    "-f",
    help=(
        "Template string that will be formatted to generated"
        "output filename. Such template should be Python formattable"
        "string, and could use {filename}, {instrument}, and {codec}"
        "variables"
    ),
)

ModelParametersOption: str = Option(
    "spleeter:2stems",
    "--params_filename",
    "-p",
    help="JSON filename that contains params",
)


MWFOption: bool = Option(
    False, "--mwf", help="Whether to use multichannel Wiener filtering for separation"
)

MUSDBDirectoryOption: str = Option(
    ...,
    "--mus_dir",
    exists=True,
    dir_okay=True,
    file_okay=False,
    readable=True,
    resolve_path=True,
    help="Path to musDB dataset directory",
)

TrainingDataDirectoryOption: str = Option(
    ...,
    "--data",
    "-d",
    exists=True,
    dir_okay=True,
    file_okay=False,
    readable=True,
    resolve_path=True,
    help="Path of the folder containing audio data for training",
)

VerboseOption: bool = Option(False, "--verbose", help="Enable verbose logs")


def version_callback(value: bool):
    if value:
        from importlib.metadata import version

        echo(f"Spleeter Version: {version('spleeter')}")
        raise Exit()


VersionOption: bool = Option(
    None,
    "--version",
    callback=version_callback,
    is_eager=True,
    help="Return Spleeter version",
)



File: spleeter/__init__.py

#!/usr/bin/env python
# coding: utf8

"""
Spleeter is the Deezer source separation library with pretrained models.
The library is based on Tensorflow:

-   It provides already trained model for performing separation.
-   It makes it easy to train source separation model with tensorflow
    (provided you have a dataset of isolated sources).

This module allows to interact easily from command line with Spleeter
by providing train, evaluation and source separation action.
"""

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


class SpleeterError(Exception):
    """Custom exception for Spleeter related error."""

    pass



File: spleeter/types.py

#!/usr/bin/env python
# coding: utf8

""" Custom types definition. """

from typing import Any, Tuple

# pyright: reportMissingImports=false
# pylint: disable=import-error
import numpy as np

# pylint: enable=import-error

AudioDescriptor = Any
Signal = Tuple[np.ndarray, float]



File: spleeter/dataset.py

#!/usr/bin/env python
# coding: utf8

"""
Module for building data preprocessing pipeline using the tensorflow
data API. Data preprocessing such as audio loading, spectrogram
computation, cropping, feature caching or data augmentation is done
using a tensorflow dataset object that output a tuple (input_, output)
where:

-   input is a dictionary with a single key that contains the (batched)
    mix spectrogram of audio samples
-   output is a dictionary of spectrogram of the isolated tracks
    (ground truth)
"""

import os
import time
from os.path import exists
from os.path import sep as SEPARATOR
from typing import Any, Dict, List, Optional, Tuple

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore

from .audio.adapter import AudioAdapter
from .audio.convertor import db_uint_spectrogram_to_gain, spectrogram_to_db_uint
from .audio.spectrogram import (
    compute_spectrogram_tf,
    random_pitch_shift,
    random_time_stretch,
)
from .utils.logging import logger
from .utils.tensor import (
    check_tensor_shape,
    dataset_from_csv,
    set_tensor_shape,
    sync_apply,
)

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"

# Default audio parameters to use.
DEFAULT_AUDIO_PARAMS: Dict = {
    "instrument_list": ("vocals", "accompaniment"),
    "mix_name": "mix",
    "sample_rate": 44100,
    "frame_length": 4096,
    "frame_step": 1024,
    "T": 512,
    "F": 1024,
}


def get_training_dataset(
    audio_params: Dict, audio_adapter: AudioAdapter, audio_path: str
) -> Any:
    """
    Builds training dataset.

    Parameters:
        audio_params (Dict):
            Audio parameters.
        audio_adapter (AudioAdapter):
            Adapter to load audio from.
        audio_path (str):
            Path of directory containing audio.

    Returns:
        Any:
            Built dataset.
    """
    builder = DatasetBuilder(
        audio_params,
        audio_adapter,
        audio_path,
        chunk_duration=audio_params.get("chunk_duration", 20.0),
        random_seed=audio_params.get("random_seed", 0),
    )
    return builder.build(
        str(audio_params.get("train_csv")),
        cache_directory=audio_params.get("training_cache"),
        batch_size=audio_params.get("batch_size", 8),
        n_chunks_per_song=audio_params.get("n_chunks_per_song", 2),
        random_data_augmentation=False,
        convert_to_uint=True,
        wait_for_cache=False,
    )


def get_validation_dataset(
    audio_params: Dict, audio_adapter: AudioAdapter, audio_path: str
) -> Any:
    """
    Builds validation dataset.

    Parameters:
        audio_params (Dict):
            Audio parameters.
        audio_adapter (AudioAdapter):
            Adapter to load audio from.
        audio_path (str):
            Path of directory containing audio.

    Returns:
        Any:
            Built dataset.
    """
    builder = DatasetBuilder(
        audio_params, audio_adapter, audio_path, chunk_duration=12.0
    )
    return builder.build(
        str(audio_params.get("validation_csv")),
        batch_size=audio_params.get("batch_size", 8),
        cache_directory=audio_params.get("validation_cache"),
        convert_to_uint=True,
        infinite_generator=False,
        n_chunks_per_song=1,
        # should not perform data augmentation for eval:
        random_data_augmentation=False,
        random_time_crop=False,
        shuffle=False,
    )


class InstrumentDatasetBuilder(object):
    """Instrument based filter and mapper provider."""

    def __init__(self, parent: Any, instrument: Any) -> None:
        """
        Default constructor.

        Parameters:
            parent (Any):
                Parent dataset builder.
            instrument (Any):
                Target instrument.
        """
        self._parent = parent
        self._instrument = instrument
        self._spectrogram_key = f"{instrument}_spectrogram"
        self._min_spectrogram_key = f"min_{instrument}_spectrogram"
        self._max_spectrogram_key = f"max_{instrument}_spectrogram"

    def load_waveform(self, sample: Dict) -> Dict:
        """Load waveform for given sample."""
        return dict(
            sample,
            **self._parent._audio_adapter.load_waveform(
                sample[f"{self._instrument}_path"],
                offset=sample["start"],
                duration=self._parent._chunk_duration,
                sample_rate=self._parent._sample_rate,
                waveform_name="waveform",
            ),
        )

    def compute_spectrogram(self, sample: Dict) -> Dict:
        """Compute spectrogram of the given sample."""
        return dict(
            sample,
            **{
                self._spectrogram_key: compute_spectrogram_tf(
                    sample["waveform"],
                    frame_length=self._parent._frame_length,
                    frame_step=self._parent._frame_step,
                    spec_exponent=1.0,
                    window_exponent=1.0,
                )
            },
        )

    def filter_frequencies(self, sample: Dict) -> Dict:
        return dict(
            sample,
            **{
                self._spectrogram_key: sample[self._spectrogram_key][
                    :, : self._parent._F, :
                ]
            },
        )

    def convert_to_uint(self, sample: Dict) -> Dict:
        """Convert given sample from float to unit."""
        return dict(
            sample,
            **spectrogram_to_db_uint(
                sample[self._spectrogram_key],
                tensor_key=self._spectrogram_key,
                min_key=self._min_spectrogram_key,
                max_key=self._max_spectrogram_key,
            ),
        )

    def filter_infinity(self, sample: Dict) -> tf.Tensor:
        """Filter infinity sample."""
        return tf.logical_not(tf.math.is_inf(sample[self._min_spectrogram_key]))

    def convert_to_float32(self, sample: Dict) -> Dict:
        """Convert given sample from unit to float."""
        return dict(
            sample,
            **{
                self._spectrogram_key: db_uint_spectrogram_to_gain(
                    sample[self._spectrogram_key],
                    sample[self._min_spectrogram_key],
                    sample[self._max_spectrogram_key],
                )
            },
        )

    def time_crop(self, sample: Dict) -> Dict:
        def start(sample):
            """mid_segment_start"""
            return tf.cast(
                tf.maximum(
                    tf.shape(sample[self._spectrogram_key])[0] / 2
                    - self._parent._T / 2,
                    0,
                ),
                tf.int32,
            )

        return dict(
            sample,
            **{
                self._spectrogram_key: sample[self._spectrogram_key][
                    start(sample) : start(sample) + self._parent._T, :, :
                ]
            },
        )

    def filter_shape(self, sample: Dict) -> bool:
        """Filter badly shaped sample."""
        return check_tensor_shape(
            sample[self._spectrogram_key],
            (self._parent._T, self._parent._F, self._parent._n_channels),
        )

    def reshape_spectrogram(self, sample: Dict) -> Dict:
        """Reshape given sample."""
        return dict(
            sample,
            **{
                self._spectrogram_key: set_tensor_shape(
                    sample[self._spectrogram_key],
                    (self._parent._T, self._parent._F, self._parent._n_channels),
                )
            },
        )


class DatasetBuilder(object):
    MARGIN: float = 0.5
    """Margin at beginning and end of songs in seconds."""

    WAIT_PERIOD: int = 60
    """Wait period for cache (in seconds)."""

    def __init__(
        self,
        audio_params: Dict,
        audio_adapter: AudioAdapter,
        audio_path: str,
        random_seed: int = 0,
        chunk_duration: float = 20.0,
    ) -> None:
        """
        Default constructor.
        """
        # Length of segment in frames (if fs=22050 and
        # frame_step=512, then T=512 corresponds to 11.89s)
        self._T = audio_params["T"]
        # Number of frequency bins to be used (should
        # be less than frame_length/2 + 1)
        self._F = audio_params["F"]
        self._sample_rate = audio_params["sample_rate"]
        self._frame_length = audio_params["frame_length"]
        self._frame_step = audio_params["frame_step"]
        self._mix_name = audio_params["mix_name"]
        self._n_channels = audio_params["n_channels"]
        self._instruments = [self._mix_name] + audio_params["instrument_list"]
        self._instrument_builders: Optional[List] = None
        self._chunk_duration = chunk_duration
        self._audio_adapter = audio_adapter
        self._audio_params = audio_params
        self._audio_path = audio_path
        self._random_seed = random_seed

        self.check_parameters_compatibility()

    def check_parameters_compatibility(self):
        if self._frame_length / 2 + 1 < self._F:
            raise ValueError(
                "F is too large and must be set to at most frame_length/2+1. "
                "Decrease F or increase frame_length to fix."
            )

        if (
            self._chunk_duration * self._sample_rate - self._frame_length
        ) / self._frame_step < self._T:
            raise ValueError(
                "T is too large considering STFT parameters and chunk duratoin. "
                "Make sure spectrogram time dimension of chunks is larger than T "
                "(for instance reducing T or frame_step or increasing chunk duration)."
            )

    def expand_path(self, sample: Dict) -> Dict:
        """Expands audio paths for the given sample."""
        return dict(
            sample,
            **{
                f"{instrument}_path": tf.strings.join(
                    (self._audio_path, sample[f"{instrument}_path"]), SEPARATOR
                )
                for instrument in self._instruments
            },
        )

    def filter_error(self, sample: Dict) -> tf.Tensor:
        """Filter errored sample."""
        return tf.logical_not(sample["waveform_error"])

    def filter_waveform(self, sample: Dict) -> Dict:
        """Filter waveform from sample."""
        return {k: v for k, v in sample.items() if not k == "waveform"}

    def harmonize_spectrogram(self, sample: Dict) -> Dict:
        """Ensure same size for vocals and mix spectrograms."""

        def _reduce(sample):
            return tf.reduce_min(
                [
                    tf.shape(sample[f"{instrument}_spectrogram"])[0]
                    for instrument in self._instruments
                ]
            )

        return dict(
            sample,
            **{
                f"{instrument}_spectrogram": sample[f"{instrument}_spectrogram"][
                    : _reduce(sample), :, :
                ]
                for instrument in self._instruments
            },
        )

    def filter_short_segments(self, sample: Dict) -> tf.Tensor:
        """Filter out too short segment."""
        return tf.reduce_any(
            [
                tf.shape(sample[f"{instrument}_spectrogram"])[0] >= self._T
                for instrument in self._instruments
            ]
        )

    def random_time_crop(self, sample: Dict) -> Dict:
        """Random time crop of 11.88s."""
        return dict(
            sample,
            **sync_apply(
                {
                    f"{instrument}_spectrogram": sample[f"{instrument}_spectrogram"]
                    for instrument in self._instruments
                },
                lambda x: tf.image.random_crop(
                    x,
                    (self._T, len(self._instruments) * self._F, self._n_channels),
                    seed=self._random_seed,
                ),
            ),
        )

    def random_time_stretch(self, sample: Dict) -> Dict:
        """Randomly time stretch the given sample."""
        return dict(
            sample,
            **sync_apply(
                {
                    f"{instrument}_spectrogram": sample[f"{instrument}_spectrogram"]
                    for instrument in self._instruments
                },
                lambda x: random_time_stretch(x, factor_min=0.9, factor_max=1.1),
            ),
        )

    def random_pitch_shift(self, sample: Dict) -> Dict:
        """Randomly pitch shift the given sample."""
        return dict(
            sample,
            **sync_apply(
                {
                    f"{instrument}_spectrogram": sample[f"{instrument}_spectrogram"]
                    for instrument in self._instruments
                },
                lambda x: random_pitch_shift(x, shift_min=-1.0, shift_max=1.0),
                concat_axis=0,
            ),
        )

    def map_features(self, sample: Dict) -> Tuple[Dict, Dict]:
        """Select features and annotation of the given sample."""
        input_ = {
            f"{self._mix_name}_spectrogram": sample[f"{self._mix_name}_spectrogram"]
        }
        output = {
            f"{instrument}_spectrogram": sample[f"{instrument}_spectrogram"]
            for instrument in self._audio_params["instrument_list"]
        }
        return (input_, output)

    def compute_segments(self, dataset: Any, n_chunks_per_song: int) -> Any:
        """
        Computes segments for each song of the dataset.

        Parameters:
            dataset (Any):
                Dataset to compute segments for.
            n_chunks_per_song (int):
                Number of segment per song to compute.

        Returns:
            Any:
                Segmented dataset.
        """
        if n_chunks_per_song <= 0:
            raise ValueError("n_chunks_per_song must be positif")
        datasets = []
        for k in range(n_chunks_per_song):
            if n_chunks_per_song > 1:
                datasets.append(
                    dataset.map(
                        lambda sample: dict(
                            sample,
                            start=tf.maximum(
                                k
                                * (
                                    sample["duration"]
                                    - self._chunk_duration
                                    - 2 * self.MARGIN
                                )
                                / (n_chunks_per_song - 1)
                                + self.MARGIN,
                                0,
                            ),
                        )
                    )
                )
            elif n_chunks_per_song == 1:  # Take central segment.
                datasets.append(
                    dataset.map(
                        lambda sample: dict(
                            sample,
                            start=tf.maximum(
                                sample["duration"] / 2 - self._chunk_duration / 2, 0
                            ),
                        )
                    )
                )
        dataset = datasets[-1]
        for d in datasets[:-1]:
            dataset = dataset.concatenate(d)
        return dataset

    @property
    def instruments(self) -> Any:
        """
        Instrument dataset builder generator.

        Yields:
            Any:
                InstrumentBuilder instance.
        """
        if self._instrument_builders is None:
            self._instrument_builders = []
            for instrument in self._instruments:
                self._instrument_builders.append(
                    InstrumentDatasetBuilder(self, instrument)
                )
        for builder in self._instrument_builders:
            yield builder

    def cache(self, dataset: Any, cache: Optional[str], wait: bool) -> Any:
        """
        Cache the given dataset if cache is enabled. Eventually waits for
        cache to be available (useful if another process is already
        computing cache) if provided wait flag is `True`.

        Parameters:
            dataset (Any):
                Dataset to be cached if cache is required.
            cache (str):
                Path of cache directory to be used, None if no cache.
            wait (bool):
                If caching is enabled, True is cache should be waited.

        Returns:
            Any:
                Cached dataset if needed, original dataset otherwise.
        """
        if cache is not None:
            if wait:
                while not exists(f"{cache}.index"):
                    logger.info(f"Cache not available, wait {self.WAIT_PERIOD}")
                    time.sleep(self.WAIT_PERIOD)
            cache_path = os.path.split(cache)[0]
            os.makedirs(cache_path, exist_ok=True)
            return dataset.cache(cache)
        return dataset

    def build(
        self,
        csv_path: str,
        batch_size: int = 8,
        shuffle: bool = True,
        convert_to_uint: bool = True,
        random_data_augmentation: bool = False,
        random_time_crop: bool = True,
        infinite_generator: bool = True,
        cache_directory: Optional[str] = None,
        wait_for_cache: bool = False,
        num_parallel_calls: int = 4,
        n_chunks_per_song: int = 2,
    ) -> Any:
        dataset = dataset_from_csv(csv_path)
        dataset = self.compute_segments(dataset, n_chunks_per_song)
        # Shuffle data
        if shuffle:
            dataset = dataset.shuffle(
                buffer_size=200000,
                seed=self._random_seed,
                # useless since it is cached :
                reshuffle_each_iteration=True,
            )
        # Expand audio path.
        dataset = dataset.map(self.expand_path)
        # Load waveform, compute spectrogram, and filtering error,
        # K bins frequencies, and waveform.
        N = num_parallel_calls
        for instrument in self.instruments:
            dataset = (
                dataset.map(instrument.load_waveform, num_parallel_calls=N)
                .filter(self.filter_error)
                .map(instrument.compute_spectrogram, num_parallel_calls=N)
                .map(instrument.filter_frequencies)
            )
        dataset = dataset.map(self.filter_waveform)
        # Convert to uint before caching in order to save space.
        if convert_to_uint:
            for instrument in self.instruments:
                dataset = dataset.map(instrument.convert_to_uint)
        dataset = self.cache(dataset, cache_directory, wait_for_cache)
        # Check for INFINITY (should not happen)
        for instrument in self.instruments:
            dataset = dataset.filter(instrument.filter_infinity)
        # Repeat indefinitly
        if infinite_generator:
            dataset = dataset.repeat(count=-1)
        # Ensure same size for vocals and mix spectrograms.
        # NOTE: could be done before caching ?
        dataset = dataset.map(self.harmonize_spectrogram)
        # Filter out too short segment.
        # NOTE: could be done before caching ?
        dataset = dataset.filter(self.filter_short_segments)
        # Random time crop of 11.88s
        if random_time_crop:
            dataset = dataset.map(self.random_time_crop, num_parallel_calls=N)
        else:
            # frame_duration = 11.88/T
            # take central segment (for validation)
            for instrument in self.instruments:
                dataset = dataset.map(instrument.time_crop)
        # Post cache shuffling. Done where the data are the lightest:
        # after croping but before converting back to float.
        if shuffle:
            dataset = dataset.shuffle(
                buffer_size=256, seed=self._random_seed, reshuffle_each_iteration=True
            )
        # Convert back to float32
        if convert_to_uint:
            for instrument in self.instruments:
                dataset = dataset.map(
                    instrument.convert_to_float32, num_parallel_calls=N
                )
        M = 8  # Parallel call post caching.
        # Must be applied with the same factor on mix and vocals.
        if random_data_augmentation:
            dataset = dataset.map(self.random_time_stretch, num_parallel_calls=M).map(
                self.random_pitch_shift, num_parallel_calls=M
            )
        # Filter by shape (remove badly shaped tensors).
        for instrument in self.instruments:
            dataset = dataset.filter(instrument.filter_shape).map(
                instrument.reshape_spectrogram
            )
        # Select features and annotation.
        dataset = dataset.map(self.map_features)
        # Make batch (done after selection to avoid
        # error due to unprocessed instrument spectrogram batching).
        dataset = dataset.batch(batch_size)
        return dataset



File: spleeter/__main__.py

#!/usr/bin/env python
# coding: utf8

"""
Python oneliner script usage.

USAGE: python -m spleeter {train,evaluate,separate} ...

Notes:
    All critical import involving TF, numpy or Pandas are deported to
    command function scope to avoid heavy import on CLI evaluation,
    leading to large bootstraping time.
"""
import json
from functools import partial
from glob import glob
from itertools import product
from os.path import join
from typing import Dict, List, Optional, Tuple

# pyright: reportMissingImports=false
# pylint: disable=import-error
from typer import Exit, Typer

from . import SpleeterError
from .audio import Codec
from .options import (
    AudioAdapterOption,
    AudioBitrateOption,
    AudioCodecOption,
    AudioDurationOption,
    AudioInputArgument,
    AudioInputOption,
    AudioOffsetOption,
    AudioOutputOption,
    FilenameFormatOption,
    ModelParametersOption,
    MUSDBDirectoryOption,
    MWFOption,
    TrainingDataDirectoryOption,
    VerboseOption,
    VersionOption,
)
from .utils.logging import configure_logger, logger

# pylint: enable=import-error

spleeter: Typer = Typer(add_completion=False, no_args_is_help=True, short_help="-h")
""" CLI application. """


@spleeter.callback()
def default(
    version: bool = VersionOption,
) -> None:
    pass


@spleeter.command(no_args_is_help=True)
def train(
    adapter: str = AudioAdapterOption,
    data: str = TrainingDataDirectoryOption,
    params_filename: str = ModelParametersOption,
    verbose: bool = VerboseOption,
) -> None:
    """
    Train a source separation model
    """
    import tensorflow as tf  # type: ignore

    from .audio.adapter import AudioAdapter
    from .dataset import get_training_dataset, get_validation_dataset
    from .model import model_fn
    from .model.provider import ModelProvider
    from .utils.configuration import load_configuration

    configure_logger(verbose)
    audio_adapter = AudioAdapter.get(adapter)
    params = load_configuration(params_filename)
    session_config = tf.compat.v1.ConfigProto()
    session_config.gpu_options.per_process_gpu_memory_fraction = 0.45
    estimator = tf.estimator.Estimator(
        model_fn=model_fn,
        model_dir=params["model_dir"],
        params=params,
        config=tf.estimator.RunConfig(
            save_checkpoints_steps=params["save_checkpoints_steps"],
            tf_random_seed=params["random_seed"],
            save_summary_steps=params["save_summary_steps"],
            session_config=session_config,
            log_step_count_steps=10,
            keep_checkpoint_max=2,
        ),
    )
    input_fn = partial(get_training_dataset, params, audio_adapter, data)
    train_spec = tf.estimator.TrainSpec(
        input_fn=input_fn, max_steps=params["train_max_steps"]
    )
    input_fn = partial(get_validation_dataset, params, audio_adapter, data)
    evaluation_spec = tf.estimator.EvalSpec(
        input_fn=input_fn, steps=None, throttle_secs=params["throttle_secs"]
    )
    logger.info("Start model training")
    tf.estimator.train_and_evaluate(estimator, train_spec, evaluation_spec)
    ModelProvider.writeProbe(params["model_dir"])
    logger.info("Model training done")


@spleeter.command(no_args_is_help=True)
def separate(
    deprecated_files: Optional[str] = AudioInputOption,
    files: List[str] = AudioInputArgument,
    adapter: str = AudioAdapterOption,
    bitrate: str = AudioBitrateOption,
    codec: Codec = AudioCodecOption,
    duration: float = AudioDurationOption,
    offset: float = AudioOffsetOption,
    output_path: str = AudioOutputOption,
    filename_format: str = FilenameFormatOption,
    params_filename: str = ModelParametersOption,
    mwf: bool = MWFOption,
    verbose: bool = VerboseOption,
) -> None:
    """
    Separate audio file(s)
    """
    from .audio.adapter import AudioAdapter
    from .separator import Separator

    configure_logger(verbose)
    if deprecated_files is not None:
        logger.error(
            "⚠️ -i option is not supported anymore, audio files must be supplied "
            "using input argument instead (see spleeter separate --help)"
        )
        raise Exit(20)
    audio_adapter: AudioAdapter = AudioAdapter.get(adapter)
    separator: Separator = Separator(params_filename, MWF=mwf)

    for filename in files:
        separator.separate_to_file(
            filename,
            output_path,
            audio_adapter=audio_adapter,
            offset=offset,
            duration=duration,
            codec=codec,
            bitrate=bitrate,
            filename_format=filename_format,
            synchronous=False,
        )
    separator.join()


EVALUATION_SPLIT: str = "test"
EVALUATION_METRICS_DIRECTORY: str = "metrics"
EVALUATION_INSTRUMENTS: Tuple[str, ...] = ("vocals", "drums", "bass", "other")
EVALUATION_METRICS: Tuple[str, ...] = ("SDR", "SAR", "SIR", "ISR")
EVALUATION_MIXTURE: str = "mixture.wav"
EVALUATION_AUDIO_DIRECTORY: str = "audio"


def _compile_metrics(metrics_output_directory: str) -> Dict:
    """
    Compiles metrics from given directory and returns results as dict.

    Parameters:
        metrics_output_directory (str):
            Directory to get metrics from.

    Returns:
        Dict:
            Compiled metrics as dict.
    """
    import numpy as np
    import pandas as pd  # type: ignore

    songs = glob(join(metrics_output_directory, "test/*.json"))
    index = pd.MultiIndex.from_tuples(
        product(EVALUATION_INSTRUMENTS, EVALUATION_METRICS),
        names=["instrument", "metric"],
    )
    pd.DataFrame([], index=["config1", "config2"], columns=index)
    metrics: Dict = {
        instrument: {k: [] for k in EVALUATION_METRICS}
        for instrument in EVALUATION_INSTRUMENTS
    }
    for song in songs:
        with open(song, "r") as stream:
            data = json.load(stream)
        for target in data["targets"]:
            instrument = target["name"]
            for metric in EVALUATION_METRICS:
                sdr_med = np.median(
                    [
                        frame["metrics"][metric]
                        for frame in target["frames"]
                        if not np.isnan(frame["metrics"][metric])
                    ]
                )
                metrics[instrument][metric].append(sdr_med)
    return metrics


@spleeter.command(no_args_is_help=True)
def evaluate(
    adapter: str = AudioAdapterOption,
    output_path: str = AudioOutputOption,
    params_filename: str = ModelParametersOption,
    mus_dir: str = MUSDBDirectoryOption,
    mwf: bool = MWFOption,
    verbose: bool = VerboseOption,
) -> Dict:
    """
    Evaluate a model on the musDB test dataset
    """
    import numpy as np

    configure_logger(verbose)
    try:
        import musdb  # type: ignore
        import museval  # type: ignore
    except ImportError:
        logger.error("Extra dependencies musdb and museval not found")
        logger.error("Please install musdb and museval first, abort")
        raise Exit(10)
    # Separate musdb sources.
    songs = glob(join(mus_dir, EVALUATION_SPLIT, "*/"))
    mixtures = [join(song, EVALUATION_MIXTURE) for song in songs]
    audio_output_directory = join(output_path, EVALUATION_AUDIO_DIRECTORY)
    separate(
        deprecated_files=None,
        files=mixtures,
        adapter=adapter,
        bitrate="128k",
        codec=Codec.WAV,
        duration=600.0,
        offset=0,
        output_path=join(audio_output_directory, EVALUATION_SPLIT),
        filename_format="{foldername}/{instrument}.{codec}",
        params_filename=params_filename,
        mwf=mwf,
        verbose=verbose,
    )
    # Compute metrics with musdb.
    metrics_output_directory = join(output_path, EVALUATION_METRICS_DIRECTORY)
    logger.info("Starting musdb evaluation (this could be long) ...")
    dataset = musdb.DB(root=mus_dir, is_wav=True, subsets=[EVALUATION_SPLIT])
    museval.eval_mus_dir(
        dataset=dataset,
        estimates_dir=audio_output_directory,
        output_dir=metrics_output_directory,
    )
    logger.info("musdb evaluation done")
    # Compute and pretty print median metrics.
    metrics = _compile_metrics(metrics_output_directory)
    for instrument, metric in metrics.items():
        logger.info(f"{instrument}:")
        for metric, value in metric.items():
            logger.info(f"{metric}: {np.median(value):.3f}")
    return metrics


def entrypoint():
    """Application entrypoint."""
    try:
        spleeter()
    except SpleeterError as e:
        logger.error(e)


if __name__ == "__main__":
    entrypoint()



File: spleeter/resources/__init__.py

#!/usr/bin/env python
# coding: utf8

""" Packages that provides static resources file for the library. """

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"



File: spleeter/utils/configuration.py

#!/usr/bin/env python
# coding: utf8

""" Module that provides configuration loading function. """

import importlib.resources as loader
import json
from os.path import exists
from typing import Dict

from .. import SpleeterError, resources

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"

_EMBEDDED_CONFIGURATION_PREFIX: str = "spleeter:"


def load_configuration(descriptor: str) -> Dict:
    """
    Load configuration from the given descriptor.
    Could be either a `spleeter:` prefixed embedded configuration name
    or a file system path to read configuration from.

    Parameters:
        descriptor (str):
            Configuration descriptor to use for lookup.

    Returns:
        Dict:
            Loaded description as dict.

    Raises:
        ValueError:
            If required embedded configuration does not exists.
        SpleeterError:
            If required configuration file does not exists.
    """
    # Embedded configuration reading.
    if descriptor.startswith(_EMBEDDED_CONFIGURATION_PREFIX):
        name = descriptor[len(_EMBEDDED_CONFIGURATION_PREFIX) :]
        if not loader.is_resource(resources, f"{name}.json"):
            raise SpleeterError(f"No embedded configuration {name} found")
        with loader.open_text(resources, f"{name}.json") as stream:
            return json.load(stream)
    # Standard file reading.
    if not exists(descriptor):
        raise SpleeterError(f"Configuration file {descriptor} not found")
    with open(descriptor, "r") as stream:
        return json.load(stream)



File: spleeter/utils/logging.py

#!/usr/bin/env python
# coding: utf8

"""Centralized logging facilities for Spleeter."""

import logging
import warnings
from os import environ

# pyright: reportMissingImports=false
# pylint: disable=import-error
from typer import echo

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"

environ["TF_CPP_MIN_LOG_LEVEL"] = "3"


class TyperLoggerHandler(logging.Handler):
    """A custom logger handler that use Typer echo."""

    def emit(self, record: logging.LogRecord) -> None:
        echo(self.format(record))


formatter = logging.Formatter("%(levelname)s:%(name)s:%(message)s")
handler = TyperLoggerHandler()
handler.setFormatter(formatter)
logger: logging.Logger = logging.getLogger("spleeter")
logger.addHandler(handler)
logger.setLevel(logging.INFO)


def configure_logger(verbose: bool) -> None:
    """
    Configure application logger.

    Parameters:
        verbose (bool):
            `True` to use verbose logger, `False` otherwise.
    """
    from tensorflow import get_logger  # type: ignore
    from tensorflow.compat.v1 import logging as tf_logging  # type: ignore

    tf_logger = get_logger()
    tf_logger.handlers = [handler]
    if verbose:
        tf_logging.set_verbosity(tf_logging.INFO)
        logger.setLevel(logging.DEBUG)
    else:
        warnings.filterwarnings("ignore")
        tf_logging.set_verbosity(tf_logging.ERROR)



File: spleeter/utils/__init__.py

#!/usr/bin/env python
# coding: utf8

""" This package provides utility function and classes. """

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"



File: spleeter/utils/tensor.py

#!/usr/bin/env python
# coding: utf8

""" Utility function for tensorflow. """

from typing import Any, Callable, Dict

import pandas as pd  # type: ignore

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def sync_apply(
    tensor_dict: Dict[str, tf.Tensor], func: Callable, concat_axis: int = 1
) -> Dict[str, tf.Tensor]:
    """
    Return a function that applies synchronously the provided func on the
    provided dictionnary of tensor. This means that func is applied to the
    concatenation of the tensors in tensor_dict. This is useful for
    performing random operation that needs the same drawn value on multiple
    tensor, such as a random time-crop on both input data and label (the
    same crop should be applied to both input data and label, so random
    crop cannot be applied separately on each of them).

    Note:
        All tensor are assumed to be the same shape.

    Parameters:
        tensor_dict (Dict[str, tf.Tensor]):
            A dictionary of tensor.
        func (Callable):
            Function to be applied to the concatenation of the tensors in
            `tensor_dict`.
        concat_axis (int):
            (Optional) The axis on which to perform the concatenation.

    Returns:
        Dict[str, tf.Tensor]:
            Processed tensors dictionary with the same name (keys) as input
            tensor_dict.
    """
    if concat_axis not in {0, 1}:
        raise NotImplementedError(
            "Function only implemented for concat_axis equal to 0 or 1"
        )
    tensor_list = list(tensor_dict.values())
    concat_tensor = tf.concat(tensor_list, concat_axis)
    processed_concat_tensor = func(concat_tensor)
    tensor_shape = tf.shape(list(tensor_dict.values())[0])
    D = tensor_shape[concat_axis]
    if concat_axis == 0:
        return {
            name: processed_concat_tensor[index * D : (index + 1) * D, :, :]
            for index, name in enumerate(tensor_dict)
        }
    return {
        name: processed_concat_tensor[:, index * D : (index + 1) * D, :]
        for index, name in enumerate(tensor_dict)
    }


def from_float32_to_uint8(
    tensor: tf.Tensor,
    tensor_key: str = "tensor",
    min_key: str = "min",
    max_key: str = "max",
) -> tf.Tensor:
    tensor_min = tf.reduce_min(tensor)
    tensor_max = tf.reduce_max(tensor)
    return {
        tensor_key: tf.cast(
            (tensor - tensor_min) / (tensor_max - tensor_min + 1e-16) * 255.9999,
            dtype=tf.uint8,
        ),
        min_key: tensor_min,
        max_key: tensor_max,
    }


def from_uint8_to_float32(
    tensor: tf.Tensor, tensor_min: tf.Tensor, tensor_max: tf.Tensor
) -> tf.Tensor:
    return (
        tf.cast(tensor, tf.float32) * (tensor_max - tensor_min) / 255.9999 + tensor_min
    )


def pad_and_partition(tensor: tf.Tensor, segment_len: int) -> tf.Tensor:
    """
    Pad and partition a tensor into segment of len `segment_len`
    along the first dimension. The tensor is padded with 0 in order
    to ensure that the first dimension is a multiple of `segment_len`.

    Examples:
    ```python
    >>> tensor = [[1, 2, 3], [4, 5, 6]]
    >>> segment_len = 2
    >>> pad_and_partition(tensor, segment_len)
    [[[1, 2], [4, 5]], [[3, 0], [6, 0]]]
    ````

    Parameters:
        tensor (tf.Tensor):
            Tensor of known fixed rank
        segment_len (int):
            Segment length.

    Returns:
        tf.Tensor:
            Padded and partitioned tensor.
    """
    tensor_size = tf.math.floormod(tf.shape(tensor)[0], segment_len)
    pad_size = tf.math.floormod(segment_len - tensor_size, segment_len)
    padded = tf.pad(tensor, [[0, pad_size]] + [[0, 0]] * (len(tensor.shape) - 1))
    split = (tf.shape(padded)[0] + segment_len - 1) // segment_len
    return tf.reshape(
        padded, tf.concat([[split, segment_len], tf.shape(padded)[1:]], axis=0)
    )


def pad_and_reshape(instr_spec, frame_length, F) -> Any:
    spec_shape = tf.shape(instr_spec)
    extension_row = tf.zeros((spec_shape[0], spec_shape[1], 1, spec_shape[-1]))
    n_extra_row = (frame_length) // 2 + 1 - F
    extension = tf.tile(extension_row, [1, 1, n_extra_row, 1])
    extended_spec = tf.concat([instr_spec, extension], axis=2)
    old_shape = tf.shape(extended_spec)
    new_shape = tf.concat([[old_shape[0] * old_shape[1]], old_shape[2:]], axis=0)
    processed_instr_spec = tf.reshape(extended_spec, new_shape)
    return processed_instr_spec


def dataset_from_csv(csv_path: str, **kwargs) -> Any:
    """
    Load dataset from a CSV file using Pandas.
    kwargs if any are forwarded to the `pandas.read_csv` function.

    Parameters:
        csv_path (str):
            Path of the CSV file to load dataset from.

    Returns:
        Any:
            Loaded dataset.
    """
    df = pd.read_csv(csv_path, **kwargs)
    dataset = tf.data.Dataset.from_tensor_slices({key: df[key].values for key in df})
    return dataset


def check_tensor_shape(tensor_tf: tf.Tensor, target_shape: Any) -> bool:
    """
    Return a Tensorflow boolean graph that indicates whether
    sample[features_key] has the specified target shape.
    Only check not None entries of target_shape.

    Parameters:
        tensor_tf (tensorflow.Tensor):
            Tensor to check shape for.
        target_shape (Any):
            Target shape to compare tensor to.

    Returns:
        bool:
            `True` if shape is valid, `False` otherwise (as TF boolean).
    """
    result = tf.constant(True)
    for i, target_length in enumerate(target_shape):
        if target_length:
            result = tf.logical_and(
                result, tf.equal(tf.constant(target_length), tf.shape(tensor_tf)[i])
            )
    return result


def set_tensor_shape(tensor: tf.Tensor, tensor_shape: Any) -> tf.Tensor:
    """
    Set shape for a tensor (not in place, as opposed to tf.set_shape).

    Parameters:
        tensor (tensorflow.Tensor):
            Tensor to reshape.
        tensor_shape (Any):
            Shape to apply to the tensor.

    Returns:
        tensorflow.Tensor:
            A reshaped tensor.
    """
    tensor.set_shape(tensor_shape)
    return tensor



File: spleeter/audio/adapter.py

#!/usr/bin/env python
# coding: utf8

""" AudioAdapter class defintion. """

from abc import ABC, abstractmethod
from importlib import import_module
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

# pyright: reportMissingImports=false
# pylint: disable=import-error
import numpy as np
import tensorflow as tf  # type: ignore

from spleeter.audio import Codec

from .. import SpleeterError
from ..types import AudioDescriptor, Signal
from ..utils.logging import logger

# pylint: enable=import-error


__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


class AudioAdapter(ABC):
    """An abstract class for manipulating audio signal."""

    _DEFAULT: Optional["AudioAdapter"] = None
    """Default audio adapter singleton instance."""

    @abstractmethod
    def load(
        self,
        audio_descriptor: AudioDescriptor,
        offset: Optional[float] = None,
        duration: Optional[float] = None,
        sample_rate: Optional[float] = None,
        dtype: bytes = b"float32",
    ) -> Signal:
        """
        Loads the audio file denoted by the given audio descriptor and
        returns it data as a waveform. Aims to be implemented by client.

        Parameters:
            audio_descriptor (AudioDescriptor):
                Describe song to load, in case of file based audio adapter,
                such descriptor would be a file path.
            offset (Optional[float]):
                (Optional) Start offset to load from in seconds.
            duration (Optional[float]):
                (Optional) Duration to load in seconds.
            sample_rate (Optional[float]):
                (Optional) Sample rate to load audio with.
            dtype (bytes):
                (Optional) Data type to use, default to `b'float32'`.

        Returns:
            Signal:
                Loaded data as (wf, sample_rate) tuple.
        """
        pass

    def load_waveform(
        self,
        audio_descriptor,
        offset: float = 0.0,
        duration: float = 1800.0,
        sample_rate: int = 44100,
        dtype: bytes = b"float32",
        waveform_name: str = "waveform",
    ) -> Dict[str, Any]:
        """
        Load the audio and convert it to a tensorflow waveform.

        Parameters:
            audio_descriptor (Any):
                Describe song to load, in case of file based audio adapter,
                such descriptor would be a file path.
            offset (float):
                (Optional) Start offset to load from in seconds.
            duration (float):
                (Optional) Duration to load in seconds.
            sample_rate (float):
                (Optional) Sample rate to load audio with.
            dtype (bytes):
                (Optional) Data type to use, default to `b'float32'`.
            waveform_name (str):
                (Optional) Name of the key in output dict, default to
                `'waveform'`.

        Returns:
            Dict[str, Any]:
                TF output dict with waveform as `(T x chan numpy array)`
                and a boolean that tells whether there were an error while
                trying to load the waveform.
        """
        # Cast parameters to TF format.
        offset = tf.cast(offset, tf.float64)
        duration = tf.cast(duration, tf.float64)

        # Defined safe loading function.
        def safe_load(path, offset, duration, sample_rate, dtype):
            logger.info(f"Loading audio {path} from {offset} to {offset + duration}")
            try:
                (data, _) = self.load(
                    path.numpy(),
                    offset.numpy(),
                    duration.numpy(),
                    sample_rate.numpy(),
                    dtype=dtype.numpy(),
                )
                logger.info("Audio data loaded successfully")
                return (data, False)
            except Exception as e:
                logger.exception("An error occurs while loading audio", exc_info=e)
            return (np.float32(-1.0), True)

        # Execute function and format results.
        results = (
            tf.py_function(
                safe_load,
                [audio_descriptor, offset, duration, sample_rate, dtype],
                (tf.float32, tf.bool),
            ),
        )
        waveform, error = results[0]
        return {waveform_name: waveform, f"{waveform_name}_error": error}

    @abstractmethod
    def save(
        self,
        path: Union[Path, str],
        data: np.ndarray,
        sample_rate: float,
        codec: Codec = None,
        bitrate: str = None,
    ) -> None:
        """
        Save the given audio data to the file denoted by the given path.

        Parameters:
            path (Union[Path, str]):
                Path like of the audio file to save data in.
            data (np.ndarray):
                Waveform data to write.
            sample_rate (float):
                Sample rate to write file in.
            codec (Codec):
                (Optional) Writing codec to use, default to `None`.
            bitrate (str):
                (Optional) Bitrate of the written audio file, default to
                `None`.
        """
        pass

    @classmethod
    def default(cls) -> "AudioAdapter":
        """
        Builds and returns a default audio adapter instance.

        Returns:
            AudioAdapter:
                Default adapter instance to use.
        """
        if cls._DEFAULT is None:
            from .ffmpeg import FFMPEGProcessAudioAdapter

            cls._DEFAULT = FFMPEGProcessAudioAdapter()
        return cls._DEFAULT

    @classmethod
    def get(cls, descriptor: str) -> "AudioAdapter":
        """
        Load dynamically an AudioAdapter from given class descriptor.

        Parameters:
            descriptor (str):
                Adapter class descriptor (module.Class)

        Returns:
            AudioAdapter:
                Created adapter instance.
        """
        if not descriptor:
            return cls.default()
        module_desc: List[str] = descriptor.split(".")
        adapter_class_name: str = module_desc[-1]
        module_path: str = ".".join(module_desc[:-1])
        adapter_module = import_module(module_path)
        adapter_class = getattr(adapter_module, adapter_class_name)
        if not issubclass(adapter_class, AudioAdapter):
            raise SpleeterError(
                f"{adapter_class_name} is not a valid AudioAdapter class"
            )
        return adapter_class()



File: spleeter/audio/ffmpeg.py

#!/usr/bin/env python
# coding: utf8

"""
This module provides an AudioAdapter implementation based on FFMPEG
process. Such implementation is POSIXish and depends on nothing except
standard Python libraries. Thus this implementation is the default one
used within this library.
"""

import datetime as dt
import os
import shutil
from pathlib import Path
from typing import Dict, Optional, Union

# pyright: reportMissingImports=false
# pylint: disable=import-error
import ffmpeg  # type: ignore
import numpy as np

from .. import SpleeterError
from ..types import Signal
from ..utils.logging import logger
from . import Codec
from .adapter import AudioAdapter

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


class FFMPEGProcessAudioAdapter(AudioAdapter):
    """
    An AudioAdapter implementation that use FFMPEG binary through
    subprocess in order to perform I/O operation for audio processing.

    When created, FFMPEG binary path will be checked and expended,
    raising exception if not found. Such path could be infered using
    `FFMPEG_PATH` environment variable.
    """

    SUPPORTED_CODECS: Dict[Codec, str] = {
        Codec.M4A: "aac",
        Codec.OGG: "libvorbis",
        Codec.WMA: "wmav2",
    }
    """ FFMPEG codec name mapping. """

    def __init__(_) -> None:
        """
        Default constructor, ensure FFMPEG binaries are available.

        Raises:
            SpleeterError:
                If ffmpeg or ffprobe is not found.
        """
        for binary in ("ffmpeg", "ffprobe"):
            if shutil.which(binary) is None:
                raise SpleeterError("{} binary not found".format(binary))

    def load(
        _,
        path: Union[Path, str],
        offset: Optional[float] = None,
        duration: Optional[float] = None,
        sample_rate: Optional[float] = None,
        dtype: bytes = b"float32",
    ) -> Signal:
        """
        Loads the audio file denoted by the given path
        and returns it data as a waveform.

        Parameters:
            path (Union[Path, str]:
                Path of the audio file to load data from.
            offset (Optional[float]):
                (Optional) Start offset to load from in seconds.
            duration (Optional[float]):
                (Optional) Duration to load in seconds.
            sample_rate (Optional[float]):
                (Optional) Sample rate to load audio with.
            dtype (bytes):
                (Optional) Data type to use, default to `b'float32'`.

        Returns:
            Signal:
                Loaded data a (waveform, sample_rate) tuple.

        Raises:
            SpleeterError:
                If any error occurs while loading audio.
        """
        if isinstance(path, Path):
            path = str(path)
        if not isinstance(path, str):
            path = path.decode()
        try:
            probe = ffmpeg.probe(path)
        except ffmpeg._run.Error as e:
            raise SpleeterError(
                "An error occurs with ffprobe (see ffprobe output below)\n\n{}".format(
                    e.stderr.decode()
                )
            )
        if "streams" not in probe or len(probe["streams"]) == 0:
            raise SpleeterError("No stream was found with ffprobe")
        metadata = next(
            stream for stream in probe["streams"] if stream["codec_type"] == "audio"
        )
        n_channels = metadata["channels"]
        if sample_rate is None:
            sample_rate = metadata["sample_rate"]
        output_kwargs = {"format": "f32le", "ar": sample_rate}
        if duration is not None:
            output_kwargs["t"] = str(dt.timedelta(seconds=duration))
        if offset is not None:
            output_kwargs["ss"] = str(dt.timedelta(seconds=offset))
        process = (
            ffmpeg.input(path)
            .output("pipe:", **output_kwargs)
            .run_async(pipe_stdout=True, pipe_stderr=True)
        )
        buffer, _ = process.communicate()
        waveform = np.frombuffer(buffer, dtype="<f4").reshape(-1, n_channels)
        if not waveform.dtype == np.dtype(dtype):
            waveform = waveform.astype(dtype)
        return (waveform, sample_rate)

    def save(
        self,
        path: Union[Path, str],
        data: np.ndarray,
        sample_rate: float,
        codec: Codec = None,
        bitrate: str = None,
    ) -> None:
        """
        Write waveform data to the file denoted by the given path using
        FFMPEG process.

        Parameters:
            path (Union[Path, str]):
                Path like of the audio file to save data in.
            data (np.ndarray):
                Waveform data to write.
            sample_rate (float):
                Sample rate to write file in.
            codec (Codec):
                (Optional) Writing codec to use, default to `None`.
            bitrate (str):
                (Optional) Bitrate of the written audio file, default to
                `None`.

        Raises:
            IOError:
                If any error occurs while using FFMPEG to write data.
        """
        if isinstance(path, Path):
            path = str(path)
        directory = os.path.dirname(path)
        if not os.path.exists(directory):
            raise SpleeterError(f"output directory does not exists: {directory}")
        logger.debug(f"Writing file {path}")
        input_kwargs = {"ar": sample_rate, "ac": data.shape[1]}
        output_kwargs = {"ar": sample_rate, "strict": "-2"}
        if bitrate:
            output_kwargs["audio_bitrate"] = bitrate
        if codec is not None and codec != "wav":
            output_kwargs["codec"] = self.SUPPORTED_CODECS.get(codec, codec)
        process = (
            ffmpeg.input("pipe:", format="f32le", **input_kwargs)
            .output(path, **output_kwargs)
            .overwrite_output()
            .run_async(pipe_stdin=True, pipe_stderr=True, quiet=True)
        )
        try:
            process.stdin.write(data.astype("<f4").tobytes())
            process.stdin.close()
            process.wait()
        except IOError:
            raise SpleeterError(f"FFMPEG error: {process.stderr.read()}")
        logger.info(f"File {path} written succesfully")



File: spleeter/audio/__init__.py

#!/usr/bin/env python
# coding: utf8

"""
`spleeter.utils.audio` package provides various
tools for manipulating audio content such as :

- Audio adapter class for abstract interaction with audio file.
- FFMPEG implementation for audio adapter.
- Waveform convertion and transforming functions.
"""

from enum import Enum

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


class Codec(str, Enum):
    """Enumeration of supported audio codec."""

    WAV: str = "wav"
    MP3: str = "mp3"
    OGG: str = "ogg"
    M4A: str = "m4a"
    WMA: str = "wma"
    FLAC: str = "flac"



File: spleeter/audio/convertor.py

#!/usr/bin/env python
# coding: utf8

""" This module provides audio data convertion functions. """

# pyright: reportMissingImports=false
# pylint: disable=import-error
import numpy as np
import tensorflow as tf  # type: ignore

from ..utils.tensor import from_float32_to_uint8, from_uint8_to_float32

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def to_n_channels(waveform: tf.Tensor, n_channels: int) -> tf.Tensor:
    """
    Convert a waveform to n_channels by removing or duplicating channels if
    needed (in tensorflow).

    Parameters:
        waveform (tf.Tensor):
            Waveform to transform.
        n_channels (int):
            Number of channel to reshape waveform in.

    Returns:
        tf.Tensor:
            Reshaped waveform.
    """
    return tf.cond(
        tf.shape(waveform)[1] >= n_channels,
        true_fn=lambda: waveform[:, :n_channels],
        false_fn=lambda: tf.tile(waveform, [1, n_channels])[:, :n_channels],
    )


def to_stereo(waveform: np.ndarray) -> np.ndarray:
    """
    Convert a waveform to stereo by duplicating if mono, or truncating
    if too many channels.

    Parameters:
        waveform (np.ndarray):
            a `(N, d)` numpy array.

    Returns:
        np.ndarray:
            A stereo waveform as a `(N, 1)` numpy array.
    """
    if waveform.shape[1] == 1:
        return np.repeat(waveform, 2, axis=-1)
    if waveform.shape[1] > 2:
        return waveform[:, :2]
    return waveform


def gain_to_db(tensor: tf.Tensor, espilon: float = 10e-10) -> tf.Tensor:
    """
    Convert from gain to decibel in tensorflow.

    Parameters:
        tensor (tf.Tensor):
            Tensor to convert
        epsilon (float):
            Operation constant.

    Returns:
        tf.Tensor:
            Converted tensor.
    """
    return 20.0 / np.log(10) * tf.math.log(tf.maximum(tensor, espilon))


def db_to_gain(tensor: tf.Tensor) -> tf.Tensor:
    """
    Convert from decibel to gain in tensorflow.

    Parameters:
        tensor (tf.Tensor):
            Tensor to convert

    Returns:
        tf.Tensor:
            Converted tensor.
    """
    return tf.pow(10.0, (tensor / 20.0))


def spectrogram_to_db_uint(
    spectrogram: tf.Tensor, db_range: float = 100.0, **kwargs
) -> tf.Tensor:
    """
    Encodes given spectrogram into uint8 using decibel scale.

    Parameters:
        spectrogram (tf.Tensor):
            Spectrogram to be encoded as TF float tensor.
        db_range (float):
            Range in decibel for encoding.

    Returns:
        tf.Tensor:
            Encoded decibel spectrogram as `uint8` tensor.
    """
    db_spectrogram: tf.Tensor = gain_to_db(spectrogram)
    max_db_spectrogram: tf.Tensor = tf.reduce_max(db_spectrogram)
    int_db_spectrogram: tf.Tensor = tf.maximum(
        db_spectrogram, max_db_spectrogram - db_range
    )
    return from_float32_to_uint8(int_db_spectrogram, **kwargs)


def db_uint_spectrogram_to_gain(
    db_uint_spectrogram: tf.Tensor, min_db: tf.Tensor, max_db: tf.Tensor
) -> tf.Tensor:
    """
    Decode spectrogram from uint8 decibel scale.

    Paramters:
        db_uint_spectrogram (tf.Tensor):
            Decibel spectrogram to decode.
        min_db (tf.Tensor):
            Lower bound limit for decoding.
        max_db (tf.Tensor):
            Upper bound limit for decoding.

    Returns:
        tf.Tensor:
            Decoded spectrogram as `float32` tensor.
    """
    db_spectrogram: tf.Tensor = from_uint8_to_float32(
        db_uint_spectrogram, min_db, max_db
    )
    return db_to_gain(db_spectrogram)



File: spleeter/audio/spectrogram.py

#!/usr/bin/env python
# coding: utf8

""" Spectrogram specific data augmentation. """

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore
from tensorflow.signal import hann_window, stft  # type: ignore

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def compute_spectrogram_tf(
    waveform: tf.Tensor,
    frame_length: int = 2048,
    frame_step: int = 512,
    spec_exponent: float = 1.0,
    window_exponent: float = 1.0,
) -> tf.Tensor:
    """
    Compute magnitude / power spectrogram from waveform as a
    `n_samples x n_channels` tensor.

    Parameters:
        waveform (tf.Tensor):
            Input waveform as `(times x number of channels)` tensor.
        frame_length (int):
            (Optional) Length of a STFT frame to use.
        frame_step (int):
            (Optional) HOP between successive frames.
        spec_exponent (float):
            (Optional) Exponent of the spectrogram (usually 1 for
            magnitude spectrogram, or 2 for power spectrogram).
        window_exponent (float):
            (Optional) Exponent applied to the Hann windowing function
            (may be useful for making perfect STFT/iSTFT reconstruction).

    Returns:
        tf.Tensor:
            Computed magnitude / power spectrogram as a
            `(T x F x n_channels)` tensor.
    """
    stft_tensor: tf.Tensor = tf.transpose(
        stft(
            tf.transpose(waveform),
            frame_length,
            frame_step,
            window_fn=lambda f, dtype: hann_window(
                f, periodic=True, dtype=waveform.dtype
            )
            ** window_exponent,
        ),
        perm=[1, 2, 0],
    )
    return tf.abs(stft_tensor) ** spec_exponent


def time_stretch(
    spectrogram: tf.Tensor,
    factor: float = 1.0,
    method: tf.image.ResizeMethod = tf.image.ResizeMethod.BILINEAR,
) -> tf.Tensor:
    """
    Time stretch a spectrogram preserving shape in tensorflow. Note that
    this is an approximation in the frequency domain.

    Parameters:
        spectrogram (tf.Tensor):
            Input spectrogram to be time stretched as tensor.
        factor (float):
            (Optional) Time stretch factor, must be > 0, default to `1`.
        method (tf.image.ResizeMethod):
            (Optional) Interpolation method, default to `BILINEAR`.

    Returns:
        tf.Tensor:
            Time stretched spectrogram as tensor with same shape.
    """
    T = tf.shape(spectrogram)[0]
    T_ts = tf.cast(tf.cast(T, tf.float32) * factor, tf.int32)[0]
    F = tf.shape(spectrogram)[1]
    ts_spec = tf.image.resize_images(
        spectrogram, [T_ts, F], method=method, align_corners=True
    )
    return tf.image.resize_image_with_crop_or_pad(ts_spec, T, F)


def random_time_stretch(
    spectrogram: tf.Tensor, factor_min: float = 0.9, factor_max: float = 1.1, **kwargs
) -> tf.Tensor:
    """
    Time stretch a spectrogram preserving shape with random ratio in
    tensorflow. Applies time_stretch to spectrogram with a random ratio
    drawn uniformly in `[factor_min, factor_max]`.

    Parameters:
        spectrogram (tf.Tensor):
            Input spectrogram to be time stretched as tensor.
        factor_min (float):
            (Optional) Min time stretch factor, default to `0.9`.
        factor_max (float):
            (Optional) Max time stretch factor, default to `1.1`.
        ** kwargs:
            Time stretch args.

    Returns:
        tf.Tensor:
            Randomly time stretched spectrogram as tensor with same shape.
    """
    factor = (
        tf.random_uniform(shape=(1,), seed=0) * (factor_max - factor_min) + factor_min
    )
    return time_stretch(spectrogram, factor=factor, **kwargs)


def pitch_shift(
    spectrogram: tf.Tensor,
    semitone_shift: float = 0.0,
    method: tf.image.ResizeMethod = tf.image.ResizeMethod.BILINEAR,
) -> tf.Tensor:
    """
    Pitch shift a spectrogram preserving shape in tensorflow. Note that
    this is an approximation in the frequency domain.

    Parameters:
        spectrogram (tf.Tensor):
            Input spectrogram to be pitch shifted as tensor.
        semitone_shift (float):
            (Optional) Pitch shift in semitone, default to `0.0`.
        method (tf.image.ResizeMethod):
            (Optional) Interpolation method, default to `BILINEAR`.

    Returns:
        tf.Tensor:
            Pitch shifted spectrogram (same shape as spectrogram).
    """
    factor = 2 ** (semitone_shift / 12.0)
    T = tf.shape(spectrogram)[0]
    F = tf.shape(spectrogram)[1]
    F_ps = tf.cast(tf.cast(F, tf.float32) * factor, tf.int32)[0]
    ps_spec = tf.image.resize_images(
        spectrogram, [T, F_ps], method=method, align_corners=True
    )
    paddings = [[0, 0], [0, tf.maximum(0, F - F_ps)], [0, 0]]
    return tf.pad(ps_spec[:, :F, :], paddings, "CONSTANT")


def random_pitch_shift(
    spectrogram: tf.Tensor, shift_min: float = -1.0, shift_max: float = 1.0, **kwargs
) -> tf.Tensor:
    """
    Pitch shift a spectrogram preserving shape with random ratio in
    tensorflow. Applies pitch_shift to spectrogram with a random shift
    amount (expressed in semitones) drawn uniformly in
    `[shift_min, shift_max]`.

    Parameters:
        spectrogram (tf.Tensor):
            Input spectrogram to be pitch shifted as tensor.
        shift_min (float):
            (Optional) Min pitch shift in semitone, default to -1.
        shift_max (float):
            (Optional) Max pitch shift in semitone, default to 1.

    Returns:
        tf.Tensor:
            Randomly pitch shifted spectrogram (same shape as spectrogram).
    """
    semitone_shift = (
        tf.random_uniform(shape=(1,), seed=0) * (shift_max - shift_min) + shift_min
    )
    return pitch_shift(spectrogram, semitone_shift=semitone_shift, **kwargs)



File: spleeter/model/__init__.py

#!/usr/bin/env python
# coding: utf8

""" This package provide an estimator builder as well as model functions. """

import importlib
from typing import Any, Dict, Optional, Tuple

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore
from tensorflow.signal import hann_window, inverse_stft, stft  # type: ignore

from ..utils.tensor import pad_and_partition, pad_and_reshape

# pylint: enable=import-error


__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


placeholder = tf.compat.v1.placeholder


def get_model_function(model_type):
    """
    Get tensorflow function of the model to be applied to the input tensor.
    For instance "unet.softmax_unet" will return the softmax_unet function
    in the "unet.py" submodule of the current module (spleeter.model).

    Parameters:
        model_type (str):
            The relative module path to the model function.

    Returns:
        Function:
            A tensorflow function to be applied to the input tensor to get the
            multitrack output.
    """
    relative_path_to_module = ".".join(model_type.split(".")[:-1])
    model_name = model_type.split(".")[-1]
    main_module = ".".join((__name__, "functions"))
    path_to_module = f"{main_module}.{relative_path_to_module}"
    module = importlib.import_module(path_to_module)
    model_function = getattr(module, model_name)
    return model_function


class InputProvider(object):
    def __init__(self, params):
        self.params = params

    def get_input_dict_placeholders(self):
        raise NotImplementedError()

    @property
    def input_names(self):
        raise NotImplementedError()

    def get_feed_dict(self, features, *args):
        raise NotImplementedError()


class WaveformInputProvider(InputProvider):
    @property
    def input_names(self):
        return ["audio_id", "waveform"]

    def get_input_dict_placeholders(self):
        shape = (None, self.params["n_channels"])
        features = {
            "waveform": placeholder(tf.float32, shape=shape, name="waveform"),
            "audio_id": placeholder(tf.string, name="audio_id"),
        }
        return features

    def get_feed_dict(self, features, waveform, audio_id):
        return {features["audio_id"]: audio_id, features["waveform"]: waveform}


class InputProviderFactory(object):
    @staticmethod
    def get(params):
        return WaveformInputProvider(params)


class EstimatorSpecBuilder(object):
    """
    A builder class that allows to builds a multitrack unet model
    estimator. The built model estimator has a different behaviour when
    used in a train/eval mode and in predict mode.

    * In train/eval mode:
        It takes as input and outputs magnitude spectrogram
    * In predict mode:
        It takes as input and outputs waveform.
        The whole separation process is then done in this function
        for performance reason: it makes it possible to run the whole
        separation process (including STFT and inverse STFT) on GPU.

    Example:
    >>> from spleeter.model import EstimatorSpecBuilder
    >>> builder = EstimatorSpecBuilder()
    >>> builder.build_predict_model()
    >>> builder.build_evaluation_model()
    >>> builder.build_train_model()

    >>> from spleeter.model import model_fn
    >>> estimator = tf.estimator.Estimator(model_fn=model_fn, ...)
    """

    # Supported model functions.
    DEFAULT_MODEL = "unet.unet"

    # Supported loss functions.
    L1_MASK = "L1_mask"
    WEIGHTED_L1_MASK = "weighted_L1_mask"

    # Supported optimizers.
    ADADELTA = "Adadelta"
    SGD = "SGD"

    # Math constants.
    WINDOW_COMPENSATION_FACTOR = 2.0 / 3.0
    EPSILON = 1e-10

    def __init__(self, features: Dict, params: Dict) -> None:
        """
        Default constructor. Depending on built model usage,
        the provided features should be different:

        * In train/eval mode:
            Features is a dictionary with a "mix_spectrogram" key,
            associated to the mix magnitude spectrogram.
        * In predict mode:
            Features is a dictionary with a "waveform" key,
            associated to the waveform of the sound to be separated.

        Parameters:
            features ():
                The input features for the estimator.
            params ():
                Some hyperparameters as a dictionary.
        """

        self._features = features
        self._params = params
        # Get instrument name.
        self._mix_name = params["mix_name"]
        self._instruments = params["instrument_list"]
        # Get STFT/signals parameters
        self._n_channels = params["n_channels"]
        self._T = params["T"]
        self._F = params["F"]
        self._frame_length = params["frame_length"]
        self._frame_step = params["frame_step"]

    def _build_model_outputs(self):
        """
        Created a batch_sizexTxFxn_channels input tensor containing
        mix magnitude spectrogram, then an output dict from it
        according to the selected model in internal parameters.

        Raises:
            ValueError:
                If required model_type is not supported.
        """
        input_tensor = self.spectrogram_feature
        model = self._params.get("model", None)
        if model is not None:
            model_type = model.get("type", self.DEFAULT_MODEL)
        else:
            model_type = self.DEFAULT_MODEL
        try:
            apply_model = get_model_function(model_type)
        except ModuleNotFoundError:
            raise ValueError(f"No model function {model_type} found")
        self._model_outputs = apply_model(
            input_tensor, self._instruments, self._params["model"]["params"]
        )

    def _build_loss(self, labels: Dict) -> Tuple[tf.Tensor, Dict]:
        """
        Construct tensorflow loss and metrics

        Parameters:
            labels (Dict):
                Dictionary of target outputs (key: instrument name,
                value: ground truth spectrogram of the instrument)

        Returns:
            Tuple[tf.Tensor, Dict]:
                Tensorflow (loss, metrics) tuple.
        """
        output_dict = self.model_outputs
        loss_type = self._params.get("loss_type", self.L1_MASK)
        if loss_type == self.L1_MASK:
            losses = {
                name: tf.reduce_mean(tf.abs(output - labels[name]))
                for name, output in output_dict.items()
            }
        elif loss_type == self.WEIGHTED_L1_MASK:
            losses = {
                name: tf.reduce_mean(
                    tf.reduce_mean(labels[name], axis=[1, 2, 3], keep_dims=True)
                    * tf.abs(output - labels[name])
                )
                for name, output in output_dict.items()
            }
        else:
            raise ValueError(f"Unkwnown loss type: {loss_type}")
        loss = tf.reduce_sum(list(losses.values()))
        # Add metrics for monitoring each instrument.
        metrics = {k: tf.compat.v1.metrics.mean(v) for k, v in losses.items()}
        metrics["absolute_difference"] = tf.compat.v1.metrics.mean(loss)
        return loss, metrics

    def _build_optimizer(self) -> tf.Tensor:
        """
        Builds an optimizer instance from internal parameter values.
        Default to AdamOptimizer if not specified.

        Returns:
            tf.Tensor:
                Optimizer instance from internal configuration.
        """
        name = self._params.get("optimizer")
        if name == self.ADADELTA:
            return tf.compat.v1.train.AdadeltaOptimizer()
        rate = self._params["learning_rate"]
        if name == self.SGD:
            return tf.compat.v1.train.GradientDescentOptimizer(rate)
        return tf.compat.v1.train.AdamOptimizer(rate)

    @property
    def instruments(self):
        return self._instruments

    @property
    def stft_name(self):
        return f"{self._mix_name}_stft"

    @property
    def spectrogram_name(self):
        return f"{self._mix_name}_spectrogram"

    def _build_stft_feature(self):
        """
        Compute STFT of waveform and slice the STFT in segment
        with the right length to feed the network.
        """
        stft_name = self.stft_name
        spec_name = self.spectrogram_name

        if stft_name not in self._features:
            # pad input with a frame of zeros
            waveform = tf.concat(
                [
                    tf.zeros((self._frame_length, self._n_channels)),
                    self._features["waveform"],
                ],
                0,
            )
            stft_feature = tf.transpose(
                stft(
                    tf.transpose(waveform),
                    self._frame_length,
                    self._frame_step,
                    window_fn=lambda frame_length, dtype: (
                        hann_window(frame_length, periodic=True, dtype=dtype)
                    ),
                    pad_end=True,
                ),
                perm=[1, 2, 0],
            )
            self._features[f"{self._mix_name}_stft"] = stft_feature
        if spec_name not in self._features:
            self._features[spec_name] = tf.abs(
                pad_and_partition(self._features[stft_name], self._T)
            )[:, :, : self._F, :]

    @property
    def model_outputs(self):
        if not hasattr(self, "_model_outputs"):
            self._build_model_outputs()
        return self._model_outputs

    @property
    def outputs(self):
        if not hasattr(self, "_outputs"):
            self._build_outputs()
        return self._outputs

    @property
    def stft_feature(self):
        if self.stft_name not in self._features:
            self._build_stft_feature()
        return self._features[self.stft_name]

    @property
    def spectrogram_feature(self):
        if self.spectrogram_name not in self._features:
            self._build_stft_feature()
        return self._features[self.spectrogram_name]

    @property
    def masks(self):
        if not hasattr(self, "_masks"):
            self._build_masks()
        return self._masks

    @property
    def masked_stfts(self):
        if not hasattr(self, "_masked_stfts"):
            self._build_masked_stfts()
        return self._masked_stfts

    def _inverse_stft(
        self, stft_t: tf.Tensor, time_crop: Optional[Any] = None
    ) -> tf.Tensor:
        """
        Inverse and reshape the given STFT

        Parameters:
            stft_t (tf.Tensor):
                Input STFT.
            time_crop (Optional[Any]):
                Time cropping.

        Returns:
            tf.Tensor:
                Inverse STFT (waveform).
        """
        inversed = (
            inverse_stft(
                tf.transpose(stft_t, perm=[2, 0, 1]),
                self._frame_length,
                self._frame_step,
                window_fn=lambda frame_length, dtype: (
                    hann_window(frame_length, periodic=True, dtype=dtype)
                ),
            )
            * self.WINDOW_COMPENSATION_FACTOR
        )
        reshaped = tf.transpose(inversed)
        if time_crop is None:
            time_crop = tf.shape(self._features["waveform"])[0]
        return reshaped[self._frame_length : self._frame_length + time_crop, :]

    def _build_mwf_output_waveform(self) -> Dict:
        """
        Perform separation with multichannel Wiener Filtering using Norbert.

        Note: multichannel Wiener Filtering is not coded in Tensorflow
        and thus may be quite slow.

        Returns:
            Dict:
                Dictionary of separated waveforms (key: instrument name,
                value: estimated waveform of the instrument)
        """
        import norbert  # type: ignore # pylint: disable=import-error

        output_dict = self.model_outputs
        x = self.stft_feature
        v = tf.stack(
            [
                pad_and_reshape(
                    output_dict[f"{instrument}_spectrogram"],
                    self._frame_length,
                    self._F,
                )[: tf.shape(x)[0], ...]
                for instrument in self._instruments
            ],
            axis=3,
        )
        input_args = [v, x]
        stft_function = (
            tf.py_function(
                lambda v, x: norbert.wiener(v.numpy(), x.numpy()),
                input_args,
                tf.complex64,
            ),
        )
        return {
            instrument: self._inverse_stft(stft_function[0][:, :, :, k])
            for k, instrument in enumerate(self._instruments)
        }

    def _extend_mask(self, mask: tf.Tensor) -> tf.Tensor:
        """
        Extend mask, from reduced number of frequency bin to
        the number of frequency bin in the STFT.

        Parameters:
            mask (tf.Tensor):
                Restricted mask.

        Returns:
            tf.Tensor:
                Extended mask

        Raises:
            ValueError:
                If invalid mask_extension parameter is set.
        """
        extension = self._params["mask_extension"]
        # Extend with average
        # (dispatch according to energy in the processed band)
        if extension == "average":
            extension_row = tf.reduce_mean(mask, axis=2, keepdims=True)
        # Extend with 0
        # (avoid extension artifacts but not conservative separation)
        elif extension == "zeros":
            mask_shape = tf.shape(mask)
            extension_row = tf.zeros((mask_shape[0], mask_shape[1], 1, mask_shape[-1]))
        else:
            raise ValueError(f"Invalid mask_extension parameter {extension}")
        n_extra_row = self._frame_length // 2 + 1 - self._F
        extension = tf.tile(extension_row, [1, 1, n_extra_row, 1])
        return tf.concat([mask, extension], axis=2)

    def _build_masks(self):
        """
        Compute masks from the output spectrograms of the model.
        """
        output_dict = self.model_outputs
        stft_feature = self.stft_feature
        separation_exponent = self._params["separation_exponent"]
        output_sum = (
            tf.reduce_sum(
                [e ** separation_exponent for e in output_dict.values()], axis=0
            )
            + self.EPSILON
        )
        out = {}
        for instrument in self._instruments:
            output = output_dict[f"{instrument}_spectrogram"]
            # Compute mask with the model.
            instrument_mask = (
                output ** separation_exponent + (self.EPSILON / len(output_dict))
            ) / output_sum
            # Extend mask;
            instrument_mask = self._extend_mask(instrument_mask)
            # Stack back mask.
            old_shape = tf.shape(instrument_mask)
            new_shape = tf.concat(
                [[old_shape[0] * old_shape[1]], old_shape[2:]], axis=0
            )
            instrument_mask = tf.reshape(instrument_mask, new_shape)
            # Remove padded part (for mask having the same size as STFT);

            instrument_mask = instrument_mask[: tf.shape(stft_feature)[0], ...]
            out[instrument] = instrument_mask
        self._masks = out

    def _build_masked_stfts(self):
        input_stft = self.stft_feature
        out = {}
        for instrument, mask in self.masks.items():
            out[instrument] = tf.cast(mask, dtype=tf.complex64) * input_stft
        self._masked_stfts = out

    def _build_manual_output_waveform(self, masked_stft: Dict) -> Dict:
        """
        Perform ratio mask separation

        Parameters:
            masked_stft (Dict):
                Dictionary of estimated spectrogram (key: instrument name,
                value: estimated spectrogram of the instrument).

        Returns:
            Dict:
                Dictionary of separated waveforms (key: instrument name,
                value: estimated waveform of the instrument).
        """
        output_waveform = {}
        for instrument, stft_data in masked_stft.items():
            output_waveform[instrument] = self._inverse_stft(stft_data)
        return output_waveform

    def _build_output_waveform(self, masked_stft: Dict) -> Dict:
        """
        Build output waveform from given output dict in order
        to be used in prediction context. The configuration
        building method will be using MWF.

        masked_stft (Dict):
                Dictionary of estimated spectrogram (key: instrument name,
                value: estimated spectrogram of the instrument).

        Returns:
            Dict:
                Built output waveform.
        """

        if self._params.get("MWF", False):
            output_waveform = self._build_mwf_output_waveform()
        else:
            output_waveform = self._build_manual_output_waveform(masked_stft)
        return output_waveform

    def _build_outputs(self):
        self._outputs = self._build_output_waveform(self.masked_stfts)

        if "audio_id" in self._features:
            self._outputs["audio_id"] = self._features["audio_id"]

    def build_predict_model(self) -> tf.Tensor:
        """
        Builder interface for creating model instance that aims to perform
        prediction / inference over given track. The output of such estimator
        will be a dictionary with a "<instrument>" key per separated instrument,
        associated to the estimated separated waveform of the instrument.

        Returns:
            tf.Tensor:
                An estimator for performing prediction.
        """

        return tf.estimator.EstimatorSpec(
            tf.estimator.ModeKeys.PREDICT, predictions=self.outputs
        )

    def build_evaluation_model(self, labels: Dict) -> tf.Tensor:
        """
        Builder interface for creating model instance that aims
        to perform model evaluation. The output of such estimator
        will be a dictionary with a key "<instrument>_spectrogram"
        per separated instrument, associated to the estimated
        separated instrument magnitude spectrogram.

        Parameters:
            labels (Dict):
                Model labels.

        Returns:
            tf.Tensor:
                An estimator for performing model evaluation.
        """
        loss, metrics = self._build_loss(labels)
        return tf.estimator.EstimatorSpec(
            tf.estimator.ModeKeys.EVAL, loss=loss, eval_metric_ops=metrics
        )

    def build_train_model(self, labels: Dict) -> tf.Tensor:
        """
        Builder interface for creating model instance that aims to perform
        model training. The output of such estimator will be a dictionary
        with a key "<instrument>_spectrogram" per separated instrument,
        associated to the estimated separated instrument magnitude spectrogram.

        Parameters:
            labels (Dict):
                Model labels.

        Returns:
            tf.Tensor:
                An estimator for performing model training.
        """
        loss, metrics = self._build_loss(labels)
        optimizer = self._build_optimizer()
        train_operation = optimizer.minimize(
            loss=loss, global_step=tf.compat.v1.train.get_global_step()
        )
        return tf.estimator.EstimatorSpec(
            mode=tf.estimator.ModeKeys.TRAIN,
            loss=loss,
            train_op=train_operation,
            eval_metric_ops=metrics,
        )


def model_fn(features, labels, mode, params):
    builder = EstimatorSpecBuilder(features, params)
    if mode == tf.estimator.ModeKeys.PREDICT:
        return builder.build_predict_model()
    elif mode == tf.estimator.ModeKeys.EVAL:
        return builder.build_evaluation_model(labels)
    elif mode == tf.estimator.ModeKeys.TRAIN:
        return builder.build_train_model(labels)
    raise ValueError(f"Unknown mode {mode}")



File: spleeter/model/provider/__init__.py

#!/usr/bin/env python
# coding: utf8

"""
This package provides tools for downloading model from network
using remote storage abstraction.

Example:
```python
>>> provider = MyProviderImplementation()
>>> provider.get('/path/to/local/storage', params)
```
"""

from abc import ABC, abstractmethod
from os import environ, makedirs
from os.path import exists, isabs, join, sep

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


class ModelProvider(ABC):
    """
    A ModelProvider manages model files on disk and
    file download is not available.
    """

    DEFAULT_MODEL_PATH: str = environ.get("MODEL_PATH", "pretrained_models")
    MODEL_PROBE_PATH: str = ".probe"

    @abstractmethod
    def download(_, name: str, path: str) -> None:
        """
        Download model denoted by the given name to disk.

        Parameters:
            name (str):
                Name of the model to download.
            path (str):
                Path of the directory to save model into.
        """
        pass

    @staticmethod
    def writeProbe(directory: str) -> None:
        """
        Write a model probe file into the given directory.

        Parameters:
            directory (str):
                Directory to write probe into.
        """
        probe: str = join(directory, ModelProvider.MODEL_PROBE_PATH)
        with open(probe, "w") as stream:
            stream.write("OK")

    def get(self, model_directory: str) -> str:
        """
        Ensures required model is available at given location.

        Parameters:
            model_directory (str):
                Expected model_directory to be available.

        Raises:
            IOError:
                If model can not be retrieved.

        Returns:
            str:
                Available model directory.
        """
        # Expend model directory if needed.
        if not isabs(model_directory):
            model_directory = join(self.DEFAULT_MODEL_PATH, model_directory)
        # Download it if not exists.
        model_probe: str = join(model_directory, self.MODEL_PROBE_PATH)
        if not exists(model_probe):
            if not exists(model_directory):
                makedirs(model_directory)
                self.download(model_directory.split(sep)[-1], model_directory)
                self.writeProbe(model_directory)
        return model_directory

    @classmethod
    def default(_: type) -> "ModelProvider":
        """
        Builds and returns a default model provider.

        Returns:
            ModelProvider:
                A default model provider instance to use.
        """
        from .github import GithubModelProvider

        return GithubModelProvider.from_environ()



File: spleeter/model/provider/github.py

#!/usr/bin/env python
# coding: utf8

"""
A ModelProvider backed by Github Release feature.

Examples:

```python
>>> from spleeter.model.provider import github
>>> provider = github.GithubModelProvider(
        'github.com',
        'Deezer/spleeter',
        'latest')
>>> provider.download('2stems', '/path/to/local/storage')
```
"""

import hashlib
import os
import tarfile
from os import environ
from tempfile import NamedTemporaryFile
from typing import Dict

# pyright: reportMissingImports=false
# pylint: disable=import-error
import httpx

from ...utils.logging import logger
from . import ModelProvider

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def compute_file_checksum(path):
    """
    Computes given path file sha256.

    Parameters:
        path (str):
            Path of the file to compute checksum for.

    Returns:
        str:
            File checksum.
    """
    sha256 = hashlib.sha256()
    with open(path, "rb") as stream:
        for chunk in iter(lambda: stream.read(4096), b""):
            sha256.update(chunk)
    return sha256.hexdigest()


class GithubModelProvider(ModelProvider):
    """A ModelProvider implementation backed on Github for remote storage."""

    DEFAULT_HOST: str = "https://github.com"
    DEFAULT_REPOSITORY: str = "deezer/spleeter"

    CHECKSUM_INDEX: str = "checksum.json"
    LATEST_RELEASE: str = "v1.4.0"
    RELEASE_PATH: str = "releases/download"

    def __init__(self, host: str, repository: str, release: str) -> None:
        """Default constructor.

        Parameters:
            host (str):
                Host to the Github instance to reach.
            repository (str):
                Repository path within target Github.
            release (str):
                Release name to get models from.
        """
        self._host: str = host
        self._repository: str = repository
        self._release: str = release

    @classmethod
    def from_environ(cls) -> "GithubModelProvider":
        """
        Factory method that creates provider from envvars.

        Returns:
            GithubModelProvider:
                Created instance.
        """
        return cls(
            environ.get("GITHUB_HOST", cls.DEFAULT_HOST),
            environ.get("GITHUB_REPOSITORY", cls.DEFAULT_REPOSITORY),
            environ.get("GITHUB_RELEASE", cls.LATEST_RELEASE),
        )

    def checksum(self, name: str) -> str:
        """
        Downloads and returns reference checksum for the given model name.

        Parameters:
            name (str):
                Name of the model to get checksum for.

        Returns:
            str:
                Checksum of the required model.

        Raises:
            ValueError:
                If the given model name is not indexed.
        """
        url: str = "/".join(
            (
                self._host,
                self._repository,
                self.RELEASE_PATH,
                self._release,
                self.CHECKSUM_INDEX,
            )
        )
        response: httpx.Response = httpx.get(url)
        response.raise_for_status()
        index: Dict = response.json()
        if name not in index:
            raise ValueError(f"No checksum for model {name}")
        return index[name]

    def download(self, name: str, path: str) -> None:
        """
        Download model denoted by the given name to disk.

        Parameters:
            name (str):
                Name of the model to download.
            path (str):
                Path of the directory to save model into.
        """
        url: str = "/".join(
            (self._host, self._repository, self.RELEASE_PATH, self._release, name)
        )
        url = f"{url}.tar.gz"
        logger.info(f"Downloading model archive {url}")
        with httpx.Client(http2=True) as client:
            with client.stream("GET", url) as response:
                response.raise_for_status()
                archive = NamedTemporaryFile(delete=False)
                try:
                    with archive as stream:
                        for chunk in response.iter_raw():
                            stream.write(chunk)
                    logger.info("Validating archive checksum")
                    checksum: str = compute_file_checksum(archive.name)
                    if checksum != self.checksum(name):
                        raise IOError("Downloaded file is corrupted, please retry")
                    logger.info(f"Extracting downloaded {name} archive")
                    with tarfile.open(name=archive.name) as tar:
                        tar.extractall(path=path)
                finally:
                    os.unlink(archive.name)
        logger.info(f"{name} model file(s) extracted")



File: spleeter/model/functions/unet.py

#!/usr/bin/env python
# coding: utf8

"""
This module contains building functions for U-net source
separation models in a similar way as in A. Jansson et al. :

"Singing voice separation with deep u-net convolutional networks",
ISMIR 2017

Each instrument is modeled by a single U-net
convolutional / deconvolutional network that take a mix spectrogram
as input and the estimated sound spectrogram as output.
"""

from functools import partial
from typing import Any, Dict, Iterable, Optional

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore
from tensorflow.compat.v1 import logging  # type: ignore
from tensorflow.compat.v1.keras.initializers import he_uniform  # type: ignore
from tensorflow.keras.layers import (  # type: ignore
    ELU,
    BatchNormalization,
    Concatenate,
    Conv2D,
    Conv2DTranspose,
    Dropout,
    LeakyReLU,
    Multiply,
    ReLU,
    Softmax,
)

from . import apply

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def _get_conv_activation_layer(params: Dict) -> Any:
    """
    Parameters:
        params (Dict):
            Model parameters.

    Returns:
        Any:
            Required Activation function.
    """
    conv_activation: str = str(params.get("conv_activation"))
    if conv_activation == "ReLU":
        return ReLU()
    elif conv_activation == "ELU":
        return ELU()
    return LeakyReLU(0.2)


def _get_deconv_activation_layer(params: Dict) -> Any:
    """
    Parameters:
        params (Dict):
            Model parameters.

    Returns:
        Any:
            Required Activation function.
    """
    deconv_activation: str = str(params.get("deconv_activation"))
    if deconv_activation == "LeakyReLU":
        return LeakyReLU(0.2)
    elif deconv_activation == "ELU":
        return ELU()
    return ReLU()


def apply_unet(
    input_tensor: tf.Tensor,
    output_name: str = "output",
    params: Dict = {},
    output_mask_logit: bool = False,
) -> tf.Tensor:
    """
    Apply a convolutionnal U-net to model a single instrument (one U-net
    is used for each instrument).

    Parameters:
        input_tensor (tf.Tensor):
            Input of the model.
        output_name (str):
            (Optional) name of the output, default to 'output'.
        params (Dict):
            (Optional) dict of BLSTM parameters.
        output_mask_logit (bool):
            (Optional) Sigmoid or logit?

    Returns:
        tf.Tensor:
            Output tensor.
    """
    logging.info(f"Apply unet for {output_name}")
    conv_n_filters = params.get("conv_n_filters", [16, 32, 64, 128, 256, 512])
    conv_activation_layer = _get_conv_activation_layer(params)
    deconv_activation_layer = _get_deconv_activation_layer(params)
    kernel_initializer = he_uniform(seed=50)
    conv2d_factory = partial(
        Conv2D, strides=(2, 2), padding="same", kernel_initializer=kernel_initializer
    )
    # First layer.
    conv1 = conv2d_factory(conv_n_filters[0], (5, 5))(input_tensor)
    batch1 = BatchNormalization(axis=-1)(conv1)
    rel1 = conv_activation_layer(batch1)
    # Second layer.
    conv2 = conv2d_factory(conv_n_filters[1], (5, 5))(rel1)
    batch2 = BatchNormalization(axis=-1)(conv2)
    rel2 = conv_activation_layer(batch2)
    # Third layer.
    conv3 = conv2d_factory(conv_n_filters[2], (5, 5))(rel2)
    batch3 = BatchNormalization(axis=-1)(conv3)
    rel3 = conv_activation_layer(batch3)
    # Fourth layer.
    conv4 = conv2d_factory(conv_n_filters[3], (5, 5))(rel3)
    batch4 = BatchNormalization(axis=-1)(conv4)
    rel4 = conv_activation_layer(batch4)
    # Fifth layer.
    conv5 = conv2d_factory(conv_n_filters[4], (5, 5))(rel4)
    batch5 = BatchNormalization(axis=-1)(conv5)
    rel5 = conv_activation_layer(batch5)
    # Sixth layer
    conv6 = conv2d_factory(conv_n_filters[5], (5, 5))(rel5)
    batch6 = BatchNormalization(axis=-1)(conv6)
    _ = conv_activation_layer(batch6)
    #
    #
    conv2d_transpose_factory = partial(
        Conv2DTranspose,
        strides=(2, 2),
        padding="same",
        kernel_initializer=kernel_initializer,
    )
    #
    up1 = conv2d_transpose_factory(conv_n_filters[4], (5, 5))((conv6))
    up1 = deconv_activation_layer(up1)
    batch7 = BatchNormalization(axis=-1)(up1)
    drop1 = Dropout(0.5)(batch7)
    merge1 = Concatenate(axis=-1)([conv5, drop1])
    #
    up2 = conv2d_transpose_factory(conv_n_filters[3], (5, 5))((merge1))
    up2 = deconv_activation_layer(up2)
    batch8 = BatchNormalization(axis=-1)(up2)
    drop2 = Dropout(0.5)(batch8)
    merge2 = Concatenate(axis=-1)([conv4, drop2])
    #
    up3 = conv2d_transpose_factory(conv_n_filters[2], (5, 5))((merge2))
    up3 = deconv_activation_layer(up3)
    batch9 = BatchNormalization(axis=-1)(up3)
    drop3 = Dropout(0.5)(batch9)
    merge3 = Concatenate(axis=-1)([conv3, drop3])
    #
    up4 = conv2d_transpose_factory(conv_n_filters[1], (5, 5))((merge3))
    up4 = deconv_activation_layer(up4)
    batch10 = BatchNormalization(axis=-1)(up4)
    merge4 = Concatenate(axis=-1)([conv2, batch10])
    #
    up5 = conv2d_transpose_factory(conv_n_filters[0], (5, 5))((merge4))
    up5 = deconv_activation_layer(up5)
    batch11 = BatchNormalization(axis=-1)(up5)
    merge5 = Concatenate(axis=-1)([conv1, batch11])
    #
    up6 = conv2d_transpose_factory(1, (5, 5), strides=(2, 2))((merge5))
    up6 = deconv_activation_layer(up6)
    batch12 = BatchNormalization(axis=-1)(up6)
    # Last layer to ensure initial shape reconstruction.
    if not output_mask_logit:
        up7 = Conv2D(
            2,
            (4, 4),
            dilation_rate=(2, 2),
            activation="sigmoid",
            padding="same",
            kernel_initializer=kernel_initializer,
        )((batch12))
        output = Multiply(name=output_name)([up7, input_tensor])
        return output
    return Conv2D(
        2,
        (4, 4),
        dilation_rate=(2, 2),
        padding="same",
        kernel_initializer=kernel_initializer,
    )((batch12))


def unet(
    input_tensor: tf.Tensor, instruments: Iterable[str], params: Optional[Dict] = None
) -> Dict:
    """Model function applier."""
    return apply(apply_unet, input_tensor, instruments, params)


def softmax_unet(
    input_tensor: tf.Tensor, instruments: Iterable[str], params: Dict = {}
) -> Dict:
    """
    Apply softmax to multitrack unet in order to have mask suming to one.

    Parameters:
        input_tensor (tf.Tensor):
            Tensor to apply blstm to.
        instruments (Iterable[str]):
            Iterable that provides a collection of instruments.
        params (Dict):
            (Optional) dict of BLSTM parameters.

    Returns:
        Dict:
            Created output tensor dict.
    """
    logit_mask_list = []
    for instrument in instruments:
        out_name = f"{instrument}_spectrogram"
        logit_mask_list.append(
            apply_unet(
                input_tensor,
                output_name=out_name,
                params=params,
                output_mask_logit=True,
            )
        )
    masks = Softmax(axis=4)(tf.stack(logit_mask_list, axis=4))
    output_dict = {}
    for i, instrument in enumerate(instruments):
        out_name = f"{instrument}_spectrogram"
        output_dict[out_name] = Multiply(name=out_name)([masks[..., i], input_tensor])
    return output_dict



File: spleeter/model/functions/__init__.py

#!/usr/bin/env python
# coding: utf8

""" This package provide model functions. """

from typing import Callable, Dict, Iterable, Optional

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def apply(
    function: Callable,
    input_tensor: tf.Tensor,
    instruments: Iterable[str],
    params: Optional[Dict] = None,
) -> Dict:
    """
    Apply given function to the input tensor.

    Parameters:
        function (Callable):
            Function to be applied to tensor.
        input_tensor (tf.Tensor):
            Tensor to apply blstm to.
        instruments (Iterable[str]):
            Iterable that provides a collection of instruments.
        params (Optional[Dict]):
            (Optional) dict of BLSTM parameters.

    Returns:
        Dict:
            Created output tensor dict.
    """
    output_dict: Dict = {}
    for instrument in instruments:
        out_name = f"{instrument}_spectrogram"
        output_dict[out_name] = function(
            input_tensor, output_name=out_name, params=params or {}
        )
    return output_dict



File: spleeter/model/functions/blstm.py

#!/usr/bin/env python
# coding: utf8

"""
This system (UHL1) uses a bi-directional LSTM network as described in :

`S. Uhlich, M. Porcu, F. Giron, M. Enenkl, T. Kemp, N. Takahashi and
Y. Mitsufuji.

"Improving music source separation based on deep neural networks through
data augmentation and network blending", Proc. ICASSP, 2017.`

It has three BLSTM layers, each having 500 cells.  For each instrument,
a network is trained which predicts the target instrument amplitude from
the mixture amplitude in the STFT domain (frame size: 4096, hop size:
1024). The raw output of each network is then combined by a multichannel
Wiener filter. The network is trained on musdb where we split train into
train_train and train_valid with 86 and 14 songs, respectively. The
validation set is used to perform early stopping and hyperparameter
selection (LSTM layer dropout rate, regularization strength).
"""

from typing import Dict, Optional

# pyright: reportMissingImports=false
# pylint: disable=import-error
import tensorflow as tf  # type: ignore
from tensorflow.compat.v1.keras.initializers import he_uniform  # type: ignore
from tensorflow.compat.v1.keras.layers import CuDNNLSTM  # type: ignore
from tensorflow.keras.layers import (  # type: ignore
    Bidirectional,
    Dense,
    Flatten,
    Reshape,
    TimeDistributed,
)

from . import apply

# pylint: enable=import-error

__email__ = "spleeter@deezer.com"
__author__ = "Deezer Research"
__license__ = "MIT License"


def apply_blstm(
    input_tensor: tf.Tensor, output_name: str = "output", params: Optional[Dict] = None
) -> tf.Tensor:
    """
    Apply BLSTM to the given input_tensor.

    Parameters:
        input_tensor (tf.Tensor):
            Input of the model.
        output_name (str):
            (Optional) name of the output, default to 'output'.
        params (Optional[Dict]):
            (Optional) dict of BLSTM parameters.

    Returns:
        tf.Tensor:
            Output tensor.
    """
    if params is None:
        params = {}
    units: int = params.get("lstm_units", 250)
    kernel_initializer = he_uniform(seed=50)
    flatten_input = TimeDistributed(Flatten())((input_tensor))

    def create_bidirectional():
        return Bidirectional(
            CuDNNLSTM(
                units, kernel_initializer=kernel_initializer, return_sequences=True
            )
        )

    l1 = create_bidirectional()((flatten_input))
    l2 = create_bidirectional()((l1))
    l3 = create_bidirectional()((l2))
    dense = TimeDistributed(
        Dense(
            int(flatten_input.shape[2]),
            activation="relu",
            kernel_initializer=kernel_initializer,
        )
    )((l3))
    output: tf.Tensor = TimeDistributed(
        Reshape(input_tensor.shape[2:]), name=output_name
    )(dense)
    return output


def blstm(
    input_tensor: tf.Tensor, output_name: str = "output", params: Optional[Dict] = None
) -> tf.Tensor:
    """Model function applier."""
    return apply(apply_blstm, input_tensor, output_name, params)



<img src="https://github.com/deezer/spleeter/raw/master/images/spleeter_logo.png" height="80" />

[![Github actions](https://github.com/deezer/spleeter/workflows/pytest/badge.svg)](https://github.com/deezer/spleeter/actions) ![PyPI - Python Version](https://img.shields.io/pypi/pyversions/spleeter) [![PyPI version](https://badge.fury.io/py/spleeter.svg)](https://badge.fury.io/py/spleeter) [![Conda](https://img.shields.io/conda/vn/deezer-research/spleeter)](https://anaconda.org/deezer-research/spleeter) [![Docker Pulls](https://img.shields.io/docker/pulls/deezer/spleeter)](https://hub.docker.com/r/deezer/spleeter) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/deezer/spleeter/blob/master/spleeter.ipynb) [![Gitter chat](https://badges.gitter.im/gitterHQ/gitter.png)](https://gitter.im/spleeter/community) [![status](https://joss.theoj.org/papers/259e5efe669945a343bad6eccb89018b/status.svg)](https://joss.theoj.org/papers/259e5efe669945a343bad6eccb89018b)

> :warning: [Spleeter 2.1.0](https://pypi.org/project/spleeter/) release introduces some breaking changes, including new CLI option naming for input, and the drop
> of dedicated GPU package. Please read [CHANGELOG](CHANGELOG.md) for more details.

## About

**Spleeter** is [Deezer](https://www.deezer.com/) source separation library with pretrained models
written in [Python](https://www.python.org/) and uses [Tensorflow](https://tensorflow.org/). It makes it easy
to train source separation model (assuming you have a dataset of isolated sources), and provides
already trained state of the art model for performing various flavour of separation :

* Vocals (singing voice) / accompaniment separation ([2 stems](https://github.com/deezer/spleeter/wiki/2.-Getting-started#using-2stems-model))
* Vocals / drums / bass / other separation ([4 stems](https://github.com/deezer/spleeter/wiki/2.-Getting-started#using-4stems-model))
* Vocals / drums / bass / piano / other separation ([5 stems](https://github.com/deezer/spleeter/wiki/2.-Getting-started#using-5stems-model))

2 stems and 4 stems models have [high performances](https://github.com/deezer/spleeter/wiki/Separation-Performances) on the [musdb](https://sigsep.github.io/datasets/musdb.html) dataset. **Spleeter** is also very fast as it can perform separation of audio files to 4 stems 100x faster than real-time when run on a GPU.

We designed **Spleeter** so you can use it straight from [command line](https://github.com/deezer/spleeter/wiki/2.-Getting-started#usage)
as well as directly in your own development pipeline as a [Python library](https://github.com/deezer/spleeter/wiki/4.-API-Reference#separator). It can be installed with [pip](https://github.com/deezer/spleeter/wiki/1.-Installation#using-pip) or be used with
[Docker](https://github.com/deezer/spleeter/wiki/2.-Getting-started#using-docker-image).

### Projects and Softwares using **Spleeter**

Since it's been released, there are multiple forks exposing **Spleeter** through either a Guided User Interface (GUI) or a standalone free or paying website. Please note that we do not host, maintain or directly support any of these initiatives.

That being said, many cool projects have been built on top of ours. Notably the porting to the *Ableton Live* ecosystem through the [Spleeter 4 Max](https://github.com/diracdeltas/spleeter4max#spleeter-for-max) project.

**Spleeter** pre-trained models have also been used by professionnal audio softwares. Here's a non-exhaustive list:

* [iZotope](https://www.izotope.com/en/shop/rx-8-standard.html) in its *Music Rebalance* feature within **RX 8**
* [SpectralLayers](https://new.steinberg.net/spectralayers/) in its *Unmix* feature in **SpectralLayers 7**
* [Acon Digital](https://acondigital.com/products/acoustica-audio-editor/) within **Acoustica 7**
* [VirtualDJ](https://www.virtualdj.com/stems/) in their stem isolation feature
* [Algoriddim](https://www.algoriddim.com/apps) in their **NeuralMix** and **djayPRO** app suite

🆕 **Spleeter** is a baseline in the ongoing [Music Demixing Challenge](https://www.aicrowd.com/challenges/music-demixing-challenge-ismir-2021)!

## Spleeter Pro (Commercial version)

Check out our commercial version : [Spleeter Pro](https://www.deezer-techservices.com/solutions/spleeter/). Benefit from our expertise for precise audio separation, faster processing speeds, and dedicated professional support. 

## Quick start

Want to try it out but don't want to install anything ? We have set up a [Google Colab](https://colab.research.google.com/github/deezer/spleeter/blob/master/spleeter.ipynb).

Ready to dig into it ? In a few lines you can install **Spleeter**  and separate the vocal and accompaniment parts from an example audio file.
You need first to install `ffmpeg` and `libsndfile`. It can be done on most platform using [Conda](https://github.com/deezer/spleeter/wiki/1.-Installation#using-conda):

```bash
# install dependencies using conda
conda install -c conda-forge ffmpeg libsndfile
# install spleeter with pip
pip install spleeter
# download an example audio file (if you don't have wget, use another tool for downloading)
wget https://github.com/deezer/spleeter/raw/master/audio_example.mp3
# separate the example audio into two components
spleeter separate -p spleeter:2stems -o output audio_example.mp3
```

> :warning: Note that we no longer recommend using `conda` for installing spleeter.

> ⚠️ There are known issues with Apple M1 chips, mostly due to TensorFlow compatibility. Until these are fixed, you can use [this workaround](https://github.com/deezer/spleeter/issues/607#issuecomment-1021669444).

You should get two separated audio files (`vocals.wav` and `accompaniment.wav`) in the `output/audio_example` folder.

For a detailed documentation, please check the [repository wiki](https://github.com/deezer/spleeter/wiki/1.-Installation)

## Development and Testing

This project is managed using [Poetry](https://python-poetry.org/docs/basic-usage/), to run test suite you
can execute the following set of commands:

```bash
# Clone spleeter repository
git clone https://github.com/Deezer/spleeter && cd spleeter
# Install poetry
pip install poetry
# Install spleeter dependencies
poetry install
# Run unit test suite
poetry run pytest tests/
```

## Reference

* Deezer Research - Source Separation Engine Story - deezer.io blog post:
  * [English version](https://deezer.io/releasing-spleeter-deezer-r-d-source-separation-engine-2b88985e797e)
  * [Japanese version](http://dzr.fm/splitterjp)
* [Music Source Separation tool with pre-trained models / ISMIR2019 extended abstract](http://archives.ismir.net/ismir2019/latebreaking/000036.pdf)

If you use **Spleeter** in your work, please cite:

```BibTeX
@article{spleeter2020,
  doi = {10.21105/joss.02154},
  url = {https://doi.org/10.21105/joss.02154},
  year = {2020},
  publisher = {The Open Journal},
  volume = {5},
  number = {50},
  pages = {2154},
  author = {Romain Hennequin and Anis Khlif and Felix Voituret and Manuel Moussallam},
  title = {Spleeter: a fast and efficient music source separation tool with pre-trained models},
  journal = {Journal of Open Source Software},
  note = {Deezer Research}
}
```

## License

The code of **Spleeter** is [MIT-licensed](LICENSE).

## Disclaimer

If you plan to use **Spleeter** on copyrighted material, make sure you get proper authorization from right owners beforehand.

## Troubleshooting

**Spleeter** is a complex piece of software and although we continously try to improve and test it you may encounter unexpected issues running it. If that's the case please check the [FAQ page](https://github.com/deezer/spleeter/wiki/5.-FAQ) first as well as the list of [currently open issues](https://github.com/deezer/spleeter/issues)

### Windows users

   It appears that sometimes the shortcut command `spleeter` does not work properly on windows. This is a known issue that we will hopefully fix soon. In the meantime replace `spleeter separate` by `python -m spleeter separate` in command line and it should work.

## Contributing

If you would like to participate in the development of **Spleeter** you are more than welcome to do so. Don't hesitate to throw us a pull request and we'll do our best to examine it quickly. Please check out our [guidelines](.github/CONTRIBUTING.md) first.

## Note

This repository include a demo audio file `audio_example.mp3` which is an excerpt
from Slow Motion Dream by Steven M Bryant (c) copyright 2011 Licensed under a Creative
Commons Attribution (3.0) [license](http://dig.ccmixter.org/files/stevieb357/34740)
Ft: CSoul,Alex Beroza & Robert Siekawitch
