Spaces:
Running on Zero
Running on Zero
File size: 1,668 Bytes
fa1da78 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from distutils.version import LooseVersion
import torch
from torch_complex.tensor import ComplexTensor
from espnet2.enh.encoder.abs_encoder import AbsEncoder
from espnet2.layers.stft import Stft
is_torch_1_9_plus = LooseVersion(torch.__version__) >= LooseVersion("1.9.0")
class STFTEncoder(AbsEncoder):
"""STFT encoder for speech enhancement and separation"""
def __init__(
self,
n_fft: int = 512,
win_length: int = None,
hop_length: int = 128,
window="hann",
center: bool = True,
normalized: bool = False,
onesided: bool = True,
use_builtin_complex: bool = True,
):
super().__init__()
self.stft = Stft(
n_fft=n_fft,
win_length=win_length,
hop_length=hop_length,
window=window,
center=center,
normalized=normalized,
onesided=onesided,
)
self._output_dim = n_fft // 2 + 1 if onesided else n_fft
self.use_builtin_complex = use_builtin_complex
@property
def output_dim(self) -> int:
return self._output_dim
def forward(self, input: torch.Tensor, ilens: torch.Tensor):
"""Forward.
Args:
input (torch.Tensor): mixed speech [Batch, sample]
ilens (torch.Tensor): input lengths [Batch]
"""
spectrum, flens = self.stft(input, ilens)
if is_torch_1_9_plus and self.use_builtin_complex:
spectrum = torch.complex(spectrum[..., 0], spectrum[..., 1])
else:
spectrum = ComplexTensor(spectrum[..., 0], spectrum[..., 1])
return spectrum, flens
|