<commit_before>"""Time-frequency Keras layers.

This module has low-level implementations of some popular time-frequency operations such as STFT and inverse STFT.
We're using these layers to compose layers in `kapre.composed` where more high-level and popular layers
such as melspectrogram layer are provided. You should go check it out!

Note:
    **Why time-frequency representation?**

    Every representation (STFT, melspectrogram, etc) has something in common - they're all 2D representations
    (time, frequency-ish) of audio signals. They're helpful because they decompose an audio signal, which is a simultaneous
    mixture of a lot of frequency components into different frequency bins. They have spatial property; the frequency
    bins are *sorted*, so frequency bins nearby has represent only slightly different frequency components. The
    frequency decomposition is also what's happening during human auditory perception through cochlea.

    **Which representation to use as input?**

    For a quick summary, check out my tutorial paper, `A Tutorial on Deep Learning for Music Information Retrieval <https://arxiv.org/abs/1709.04396>`_.

"""
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Layer

from . import backend
from .backend import _CH_FIRST_STR, _CH_LAST_STR, _CH_DEFAULT_STR
from .tflite_compatible_stft import atan2_tflite

__all__ = [
    'STFT',
    'InverseSTFT',
    'Magnitude',
    'Phase',
    'MagnitudeToDecibel',
    'ApplyFilterbank',
    'Delta',
    'ConcatenateFrequencyMap',
]


def _shape_spectrum_output(spectrums, data_format):
    """Shape batch spectrograms into the right format.

    Args:
        spectrums (`Tensor`): result of tf.signal.stft or similar, i.e., (..., time, freq).
        data_format (`str`): 'channels_first' or 'channels_last'

    Returns:
        spectrums (`Tensor`): a transposed version of input `spectrums`

    """
    if data_format == _CH_FIRST_STR:
        pass  # probably it's already (batch, channel, time, freq)
    else:
        spectrums = tf.transpose(spectrums, perm=(0, 2, 3, 1))  # (batch, time, freq, channel)
    return spectrums


class STFT(Layer):
    """
    A Short-time Fourier transform layer.

    It uses `tf.signal.stft` to compute complex STFT. Additionally, it reshapes the output to be a proper 2D batch.

    If `output_data_format == 'channels_last'`, the output shape is (batch, time, freq, channel)
    If `output_data_format == 'channels_first'`, the output shape is (batch, channel, time, freq)

    Args:
        n_fft (int): Number of FFTs. Defaults to `2048`
        win_length (int or None): Window length in sample. Defaults to `n_fft`.
        hop_length (int or None): Hop length in sample between analysis windows. Defaults to `n_fft // 4` following Librosa.
        window_name (str or None): *Name* of `tf.signal` function that returns a 1D tensor window that is used in analysis.
            Defaults to `hann_window` which uses `tf.signal.hann_window`.
            Window availability depends on Tensorflow version. More details are at `kapre.backend.get_window()`.
        pad_begin (bool): Whether to pad with zeros along time axis (length: win_length - hop_length). Defaults to `False`.
        pad_end (bool): Whether to pad with zeros at the finishing end of the signal.
        input_data_format (str): the audio data format of input waveform batch.
            `'channels_last'` if it's `(batch, time, channels)` and
            `'channels_first'` if it's `(batch, channels, time)`.
            Defaults to the setting of your Keras configuration. (`tf.keras.backend.image_data_format()`)
        output_data_format (str): The data format of output STFT.<commit_msg>Fix STFT docs about default value for hop_length (#129)<commit_after>"""Time-frequency Keras layers.

This module has low-level implementations of some popular time-frequency operations such as STFT and inverse STFT.
We're using these layers to compose layers in `kapre.composed` where more high-level and popular layers
such as melspectrogram layer are provided. You should go check it out!

Note:
    **Why time-frequency representation?**

    Every representation (STFT, melspectrogram, etc) has something in common - they're all 2D representations
    (time, frequency-ish) of audio signals. They're helpful because they decompose an audio signal, which is a simultaneous
    mixture of a lot of frequency components into different frequency bins. They have spatial property; the frequency
    bins are *sorted*, so frequency bins nearby has represent only slightly different frequency components. The
    frequency decomposition is also what's happening during human auditory perception through cochlea.

    **Which representation to use as input?**

    For a quick summary, check out my tutorial paper, `A Tutorial on Deep Learning for Music Information Retrieval <https://arxiv.org/abs/1709.04396>`_.

"""
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Layer

from . import backend
from .backend import _CH_FIRST_STR, _CH_LAST_STR, _CH_DEFAULT_STR
from .tflite_compatible_stft import atan2_tflite

__all__ = [
    'STFT',
    'InverseSTFT',
    'Magnitude',
    'Phase',
    'MagnitudeToDecibel',
    'ApplyFilterbank',
    'Delta',
    'ConcatenateFrequencyMap',
]


def _shape_spectrum_output(spectrums, data_format):
    """Shape batch spectrograms into the right format.

    Args:
        spectrums (`Tensor`): result of tf.signal.stft or similar, i.e., (..., time, freq).
        data_format (`str`): 'channels_first' or 'channels_last'

    Returns:
        spectrums (`Tensor`): a transposed version of input `spectrums`

    """
    if data_format == _CH_FIRST_STR:
        pass  # probably it's already (batch, channel, time, freq)
    else:
        spectrums = tf.transpose(spectrums, perm=(0, 2, 3, 1))  # (batch, time, freq, channel)
    return spectrums


class STFT(Layer):
    """
    A Short-time Fourier transform layer.

    It uses `tf.signal.stft` to compute complex STFT. Additionally, it reshapes the output to be a proper 2D batch.

    If `output_data_format == 'channels_last'`, the output shape is (batch, time, freq, channel)
    If `output_data_format == 'channels_first'`, the output shape is (batch, channel, time, freq)

    Args:
        n_fft (int): Number of FFTs. Defaults to `2048`
        win_length (int or None): Window length in sample. Defaults to `n_fft`.
        hop_length (int or None): Hop length in sample between analysis windows. Defaults to `win_length // 4` following Librosa.
        window_name (str or None): *Name* of `tf.signal` function that returns a 1D tensor window that is used in analysis.
            Defaults to `hann_window` which uses `tf.signal.hann_window`.
            Window availability depends on Tensorflow version. More details are at `kapre.backend.get_window()`.
        pad_begin (bool): Whether to pad with zeros along time axis (length: win_length - hop_length). Defaults to `False`.
        pad_end (bool): Whether to pad with zeros at the finishing end of the signal.
        input_data_format (str): the audio data format of input waveform batch.
            `'channels_last'` if it's `(batch, time, channels)` and
            `'channels_first'` if it's `(batch, channels, time)`.
            Defaults to the setting of your Keras configuration. (`tf.keras.backend.image_data_format()`)
        output_data_format (str): The data format of output STFT.
{'fraction_non_alphanumeric': 0.08539653117966371, 'fraction_numerical': 0.006222692969680922, 'mean_word_length': 3.5754088431253788, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2723480', 'n_tokens_mistral': 2226, 'n_tokens_neox': 2082, 'n_words': 875}