Function compute_stft

utils/strand_util.py:372–405

(input, fft_size, hop_size, win_length, window_type="hann_window")

#assumes the input is (Batch,Time,dim) or (Batch,Time)
#return (B,Freq,T,dim) containing real and imaginary values
def compute_stft(input, fft_size, hop_size, win_length, window_type="hann_window"):
    window = get_window(window_type, win_length).cuda()

    if len(input.shape)==2:
        #we have a (B,T) input so we can just run stft
         x_stft = torch.stft(
            input,
            fft_size,
            hop_size,
            win_length,
            window,
            return_complex=True,
        )
    elif len(input.shape)==3:
        #do stft for each dimension
        dim=input.shape[-1]
        stft_total=[]
        for i in range(dim):
            input_axis = input[:,:,i]
            x_stft = torch.stft(
                input_axis,
                fft_size,
                hop_size,
                win_length,
                window,
                return_complex=True,
            ).unsqueeze(-1)
            stft_total.append(x_stft)
        x_stft = torch.cat(stft_total,-1)
    else:
        return None
    

    return x_stft

#inverse of compute_stft
def compute_istft(input, fft_size, hop_size, win_length, window_type="hann_window", spatial_size=256):