stft.py

import numpy as np
import math


def stft(data, window_size, step_size):
    """
    :type data: np.ndarray
    :param data: data.shape = (T,) mono
    """
    length = data.shape[0]
    window_num = math.ceil(length / step_size)
    spec_list = None
    window = np.hanning(window_size)
    for window_no in range(window_num):
        idx_begin = window_no*step_size
        idx_end = idx_begin + window_size
        
        windowed = data[idx_begin:idx_end].flatten()
        windowed.resize(window_size)
        windowed *=  window
        
        spec = np.fft.fft(windowed)
        
        if spec_list is None:
            spec_list = spec
        else:
            spec_list = np.vstack([spec_list, spec])
    return spec_list


def istft(spec_list, length, window_size, step_size):
    res = np.zeros(length+window_size)
    duplicate = np.zeros(length+window_size)
    window = np.hanning(window_size)
    for window_no, spec in enumerate(list(spec_list)):
        idx_begin = window_no*step_size
        idx_end = idx_begin + window_size
        res[idx_begin:idx_end] += np.fft.ifft(spec).real
        duplicate[idx_begin:idx_end]+=1
    return res[:length]/duplicate[:length]*2 # why 2?