pytorch_wpe_multichannel_stft_dereverberation_quickstart.py

1import torch
2from torch_wpe import wpe
3
4# Configuration parameters
5taps = 10
6delay = 3
7iterations = 3
8
9# Generate a dummy input signal (STFT domain)
10# Shape: (Batch, Frequency, Channel, Time)
11# For example: 1 batch, 513 frequency bins, 4 channels, 100 time frames
12observed_stft = torch.randn(1, 513, 4, 100, dtype=torch.complex64)
13
14# Apply WPE dereverberation
15# The function returns the estimated clean (dereverberated) signal in the STFT domain
16dereverberated_stft = wpe(
17    observed_stft,
18    taps=taps,
19    delay=delay,
20    iterations=iterations
21)
22
23print(f"Input shape: {observed_stft.shape}")
24print(f"Output shape: {dereverberated_stft.shape}")