File size: 5,230 Bytes
cfdc687
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright (c) 2021 Kazuhiro KOBAYASHI <root.4mac@gmail.com>
#
# Distributed under terms of the MIT license.

"""

"""

import librosa
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class LogMelSpectrogram(nn.Module):
    def __init__(
        self,
        fs=24000,
        hop_size=128,
        fft_size=1024,
        win_length=None,
        window="hann",
        center=True,
        pad_mode="reflect",
        n_mels=80,
        fmin=None,
        fmax=None,
        scaler_file=None,
    ):
        super().__init__()
        self.mag_layer = Magnitude(fs, hop_size, fft_size, win_length, window)
        self.mel_layer = Magnitude2LogMelSpectrogram(fs, fft_size, n_mels, fmin, fmax)
        if scaler_file is not None:
            self.melspc_scaler = LogMelSpectrogramScaler(scaler_file)
        else:
            self.melspc_scaler = None

    def forward(self, x):
        mag = self.mag_layer(x)
        log_melspc = self.mel_layer(mag)
        if self.melspc_scaler is not None:
            log_melspc = self.scaler_layer(log_melspc)
        return log_melspc


class Magnitude(torch.nn.Module):
    def __init__(
        self,
        fs=24000,
        hop_size=128,
        fft_size=1024,
        win_length=None,
        window="hann",
        center=True,
        pad_mode="reflect",
        return_complex=True,
    ):
        super().__init__()
        self.hop_size = hop_size
        self.fft_size = fft_size
        self.win_length = fft_size if win_length is None else win_length
        self.window = window
        self.center = center
        self.pad_mode = pad_mode
        self.return_complex = return_complex

    def forward(self, x):
        """
        x: (B, 1, T)
        ret: (B, T, fft_size // 2 + 1)
        """
        f = getattr(torch, f"{self.window}_window")
        window = f(self.win_length, dtype=x.dtype, device=x.device)
        y = torch.stft(
            x,
            n_fft=self.fft_size,
            win_length=self.win_length,
            hop_length=self.hop_size,
            window=window,
            center=self.center,
            pad_mode=self.pad_mode,
            return_complex=self.return_complex,
        )
        return y.abs().transpose(1, 2)


class Magnitude2LogMelSpectrogram(torch.nn.Module):
    def __init__(
        self, fs=24000, fft_size=1024, n_mels=80, fmin=None, fmax=None, eps=1.0e-10
    ):
        super().__init__()
        self.eps = eps
        fmin = 0 if fmin is None else fmin
        fmax = fs / 2 if fmax is None else fmax
        mel_basis = librosa.filters.mel(
            fs, fft_size, n_mels=n_mels, fmin=fmin, fmax=fmax
        )
        self.register_buffer("mel_basis", torch.from_numpy(mel_basis.T).float())

    def forward(self, x):
        melspc = torch.matmul(x, self.mel_basis)
        log_melspc = torch.clamp(melspc, min=self.eps).log10()
        return log_melspc


class LogMelSpectrogram2LogMagnitude(nn.Module):
    def __init__(
        self,
        fs,
        fft_size,
        n_mels=80,
        fmin=None,
        fmax=None,
        eps=1.0e-10,
        roll_size=24,
        melspc_scaler_fn=None,
    ):
        super().__init__()
        self.eps = eps
        self.roll_size = roll_size
        self.melspc_scaler_fn = melspc_scaler_fn
        fmin = 0 if fmin is None else fmin
        fmax = fs / 2 if fmax is None else fmax

        mel_basis = librosa.filters.mel(
            fs, fft_size, n_mels=n_mels, fmin=fmin, fmax=fmax
        )
        inv_mel_basis = np.linalg.pinv(mel_basis)
        self.register_buffer("inv_mel_basis", torch.from_numpy(inv_mel_basis.T).float())

    def forward(self, x):
        if self.melspc_scaler_fn is not None:
            # denorm mlfb
            x = self.melspc_scaler_fn.inverse_transform(x)

        x = torch.pow(10.0, x)
        spc = torch.matmul(x, self.inv_mel_basis)
        log_spc = 10 * torch.clamp(spc, min=self.eps).log10()
        z = F.pad(log_spc, (self.roll_size // 2 - 1, self.roll_size // 2))
        z = z.unfold(-1, self.roll_size, step=1)
        log_spc = torch.median(z, dim=-1)[0]
        return log_spc


class CepstrumLiftering(nn.Module):
    def __init__(self, lifter_size=None):
        super().__init__()
        if lifter_size <= 0:
            raise ValueError("lifter_size must be > 0.")
        else:
            self.lifter_size = lifter_size

    def forward(self, x):
        cep = torch.fft.ifft(x, dim=-1)
        cep[..., self.lifter_size : -self.lifter_size] = 0
        x = torch.fft.fft(cep, dim=-1)
        return x


class LogMelSpectrogramScaler(nn.Module):
    def __init__(self, scaler):
        super().__init__()
        self.register_parameter(
            "mean",
            nn.Parameter(torch.from_numpy(scaler.mean_).float(), requires_grad=False),
        )
        self.register_parameter(
            "scale",
            nn.Parameter(
                torch.from_numpy(scaler.var_).float().sqrt(), requires_grad=False
            ),
        )

    def forward(self, x):
        return (x - self.mean) / self.scale

    def inverse_transform(self, x):
        return x * self.scale + self.mean