Scale correction for signal reconstruction

[narrietal]

Hi,

I recently found a small bug on the code. When a signal is reconstructed via overlap and add, the signal has twice the amplitude than the original one. This is caused due to the amount of frame overlap used, anything above 50% overlap leads to an scale on the amplitude. In this case, there is a 75% overlap, as to form the STFT you take frames of 32ms with 8ms shifts, that is 75% overlap. In order to know more about this phenomenon I would recommend you to take a look at this post: https://dsp.stackexchange.com/questions/53367/overlap-add-time-domain-audio-frames-how-does-normalization-scaling-work-with-o

This problem was previously addressed by the community, so TensorFlow released a function to correct the amplitude on the iSTFT: https://www.tensorflow.org/api_docs/python/tf/signal/inverse_stft_window_fn. However, this solution is only valid if you use the "inverse_stft" TF function (https://www.tensorflow.org/api_docs/python/tf/signal/inverse_stft), which you do not. In order to solve this you can always create a custom window function with all ones, however I believe it is cleaner to just multply each frame by the right scaling factor, in this case frame_step/frame_len

Below you may find a toy example I did with a simple sinusoide signal:

Side note: I would personally recommend using a Hanning window function to obtain a signal with less artefacts. See the same example as above after using a Hanning window:

Please, find below the toy code I used to obtain the previous plots:

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

def stft_layer(signal, frame_length, frame_step, use_hanning_window=False):
    # creating frames from the continuous waveform
    frames = tf.signal.frame(signal, frame_length, frame_step)

    if use_hanning_window:
        # Create the Hanning window
        window = tf.signal.hann_window(frame_length, periodic=False)
        window = tf.cast(window, frames.dtype)

        # Apply the Hanning window to the STFT
        frames *= window

    # calculating the fft over the time frames. rfft returns NFFT/2+1 bins.
    stft = tf.signal.rfft(frames)
    return stft

def istft_layer(stft, frame_length, frame_step, amplitude_correction=True, use_hanning_window=False):
    frames = tf.signal.irfft(stft)

    if amplitude_correction:
        if use_hanning_window:
            if frame_step/frame_length < 1/2: frames *= 2*(frame_step/frame_length)
        else:
            if frame_step/frame_length < 1/2: frames *= (frame_step/frame_length)


    signal = tf.signal.overlap_and_add(frames, frame_step)

    return signal

# Example usage
# Assuming you have a 1-dimensional audio signal 'audio_signal' and the desired parameters
frame_length = 512
frame_step = 128

# Creating a sample audio signal (Replace this with your actual audio signal)
size = 2048
audio_signal = np.sin(np.arange(size) * 1 / 100)

stft = stft_layer(audio_signal, frame_length, frame_step)

default_reconstruction = istft_layer(stft, frame_length, frame_step, amplitude_correction=False)

corrected_reconstruction = istft_layer(stft, frame_length, frame_step, amplitude_correction=True)

# Plot the original signal and the reconstructed signal
plt.figure(figsize=(20, 8))

# Original Signal
plt.subplot(2, 1, 1)
# Original Signal
plt.plot(audio_signal, label='Original Signal', color='blue')
# Reconstructed Signal
plt.plot(default_reconstruction, label='Reconstructed Signal', color='red', linestyle='dashed')
plt.title('Default Signal Reconstruction')
plt.xlabel('Time')
plt.ylabel('Amplitude')

# Reconstructed Signal
plt.subplot(2, 1, 2)
# Original Signal
plt.plot(audio_signal, label='Original Signal', color='blue')
# Reconstructed Signal
plt.plot(corrected_reconstruction, label='Reconstructed Signal', color='red', linestyle='dashed')
plt.title('Corrected Signal Reconstruction')
plt.xlabel('Time')
plt.ylabel('Amplitude')

plt.tight_layout()
plt.show()

plt.savefig('default_vs_corrected_signal_reconstruction_without_TF_solution.png')