STFT encoder for speech enhancement and separation
| 6 | |
| 7 | |
| 8 | class STFTEncoder(AbsEncoder): |
| 9 | """STFT encoder for speech enhancement and separation""" |
| 10 | |
| 11 | def __init__( |
| 12 | self, |
| 13 | n_fft: int = 512, |
| 14 | win_length: int = None, |
| 15 | hop_length: int = 128, |
| 16 | window="hann", |
| 17 | center: bool = True, |
| 18 | normalized: bool = False, |
| 19 | onesided: bool = True, |
| 20 | use_builtin_complex: bool = True, |
| 21 | default_fs: int = 16000, |
| 22 | spec_transform_type: str = None, |
| 23 | spec_factor: float = 0.15, |
| 24 | spec_abs_exponent: float = 0.5, |
| 25 | ): |
| 26 | super().__init__() |
| 27 | self.stft = Stft( |
| 28 | n_fft=n_fft, |
| 29 | win_length=win_length, |
| 30 | hop_length=hop_length, |
| 31 | window=window, |
| 32 | center=center, |
| 33 | normalized=normalized, |
| 34 | onesided=onesided, |
| 35 | ) |
| 36 | |
| 37 | self._output_dim = n_fft // 2 + 1 if onesided else n_fft |
| 38 | self.use_builtin_complex = use_builtin_complex |
| 39 | self.win_length = win_length if win_length else n_fft |
| 40 | self.hop_length = hop_length |
| 41 | self.window = window |
| 42 | self.n_fft = n_fft |
| 43 | self.center = center |
| 44 | self.default_fs = default_fs |
| 45 | |
| 46 | # spec transform related. See equation (1) in paper |
| 47 | # 'Speech Enhancement and Dereverberation With Diffusion-Based Generative |
| 48 | # Models'. The default value of 0.15, 0.5 also come from the paper. |
| 49 | # spec_transform_type: "exponent", "log", or "none" |
| 50 | self.spec_transform_type = spec_transform_type |
| 51 | # the output specturm will be scaled with: spec * self.spec_factor |
| 52 | self.spec_factor = spec_factor |
| 53 | # the exponent factor used in the "exponent" transform |
| 54 | self.spec_abs_exponent = spec_abs_exponent |
| 55 | |
| 56 | def spec_transform_func(self, spec): |
| 57 | if self.spec_transform_type == "exponent": |
| 58 | if self.spec_abs_exponent != 1: |
| 59 | # only do this calculation if spec_exponent != 1, |
| 60 | # otherwise it's quite a bit of wasted computation |
| 61 | # and introduced numerical error |
| 62 | e = self.spec_abs_exponent |
| 63 | spec = spec.abs() ** e * torch.exp(1j * spec.angle()) |
| 64 | spec = spec * self.spec_factor |
| 65 | elif self.spec_transform_type == "log": |
no outgoing calls
searching dependent graphs…