260 lines
7.7 KiB
Python
260 lines
7.7 KiB
Python
|
from typing import List
|
||
|
|
from typing import Tuple
|
||
|
|
from typing import Union
|
||
|
|
|
||
|
|
import librosa
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torch_complex.tensor import ComplexTensor
|
||
|
|
|
||
|
|
from funasr.models.transformer.utils.nets_utils import make_pad_mask
|
||
|
|
|
||
|
|
|
||
|
|
class FeatureTransform(torch.nn.Module):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
# Mel options,
|
||
|
|
fs: int = 16000,
|
||
|
|
n_fft: int = 512,
|
||
|
|
n_mels: int = 80,
|
||
|
|
fmin: float = 0.0,
|
||
|
|
fmax: float = None,
|
||
|
|
# Normalization
|
||
|
|
stats_file: str = None,
|
||
|
|
apply_uttmvn: bool = True,
|
||
|
|
uttmvn_norm_means: bool = True,
|
||
|
|
uttmvn_norm_vars: bool = False,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.apply_uttmvn = apply_uttmvn
|
||
|
|
|
||
|
|
self.logmel = LogMel(fs=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
|
||
|
|
self.stats_file = stats_file
|
||
|
|
if stats_file is not None:
|
||
|
|
self.global_mvn = GlobalMVN(stats_file)
|
||
|
|
else:
|
||
|
|
self.global_mvn = None
|
||
|
|
|
||
|
|
if self.apply_uttmvn is not None:
|
||
|
|
self.uttmvn = UtteranceMVN(norm_means=uttmvn_norm_means, norm_vars=uttmvn_norm_vars)
|
||
|
|
else:
|
||
|
|
self.uttmvn = None
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self, x: ComplexTensor, ilens: Union[torch.LongTensor, np.ndarray, List[int]]
|
||
|
|
) -> Tuple[torch.Tensor, torch.LongTensor]:
|
||
|
|
# (B, T, F) or (B, T, C, F)
|
||
|
|
if x.dim() not in (3, 4):
|
||
|
|
raise ValueError(f"Input dim must be 3 or 4: {x.dim()}")
|
||
|
|
if not torch.is_tensor(ilens):
|
||
|
|
ilens = torch.from_numpy(np.asarray(ilens)).to(x.device)
|
||
|
|
|
||
|
|
if x.dim() == 4:
|
||
|
|
# h: (B, T, C, F) -> h: (B, T, F)
|
||
|
|
if self.training:
|
||
|
|
# Select 1ch randomly
|
||
|
|
ch = np.random.randint(x.size(2))
|
||
|
|
h = x[:, :, ch, :]
|
||
|
|
else:
|
||
|
|
# Use the first channel
|
||
|
|
h = x[:, :, 0, :]
|
||
|
|
else:
|
||
|
|
h = x
|
||
|
|
|
||
|
|
# h: ComplexTensor(B, T, F) -> torch.Tensor(B, T, F)
|
||
|
|
h = h.real**2 + h.imag**2
|
||
|
|
|
||
|
|
h, _ = self.logmel(h, ilens)
|
||
|
|
if self.stats_file is not None:
|
||
|
|
h, _ = self.global_mvn(h, ilens)
|
||
|
|
if self.apply_uttmvn:
|
||
|
|
h, _ = self.uttmvn(h, ilens)
|
||
|
|
|
||
|
|
return h, ilens
|
||
|
|
|
||
|
|
|
||
|
|
class LogMel(torch.nn.Module):
|
||
|
|
"""Convert STFT to fbank feats
|
||
|
|
|
||
|
|
The arguments is same as librosa.filters.mel
|
||
|
|
|
||
|
|
Args:
|
||
|
|
fs: number > 0 [scalar] sampling rate of the incoming signal
|
||
|
|
n_fft: int > 0 [scalar] number of FFT components
|
||
|
|
n_mels: int > 0 [scalar] number of Mel bands to generate
|
||
|
|
fmin: float >= 0 [scalar] lowest frequency (in Hz)
|
||
|
|
fmax: float >= 0 [scalar] highest frequency (in Hz).
|
||
|
|
If `None`, use `fmax = fs / 2.0`
|
||
|
|
htk: use HTK formula instead of Slaney
|
||
|
|
norm: {None, 1, np.inf} [scalar]
|
||
|
|
if 1, divide the triangular mel weights by the width of the mel band
|
||
|
|
(area normalization). Otherwise, leave all the triangles aiming for
|
||
|
|
a peak value of 1.0
|
||
|
|
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
fs: int = 16000,
|
||
|
|
n_fft: int = 512,
|
||
|
|
n_mels: int = 80,
|
||
|
|
fmin: float = 0.0,
|
||
|
|
fmax: float = None,
|
||
|
|
htk: bool = False,
|
||
|
|
norm=1,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
_mel_options = dict(
|
||
|
|
sr=fs, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax, htk=htk, norm=norm
|
||
|
|
)
|
||
|
|
self.mel_options = _mel_options
|
||
|
|
|
||
|
|
# Note(kamo): The mel matrix of librosa is different from kaldi.
|
||
|
|
melmat = librosa.filters.mel(**_mel_options)
|
||
|
|
# melmat: (D2, D1) -> (D1, D2)
|
||
|
|
self.register_buffer("melmat", torch.from_numpy(melmat.T).float())
|
||
|
|
|
||
|
|
def extra_repr(self):
|
||
|
|
return ", ".join(f"{k}={v}" for k, v in self.mel_options.items())
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self, feat: torch.Tensor, ilens: torch.LongTensor
|
||
|
|
) -> Tuple[torch.Tensor, torch.LongTensor]:
|
||
|
|
# feat: (B, T, D1) x melmat: (D1, D2) -> mel_feat: (B, T, D2)
|
||
|
|
mel_feat = torch.matmul(feat, self.melmat)
|
||
|
|
|
||
|
|
logmel_feat = (mel_feat + 1e-20).log()
|
||
|
|
# Zero padding
|
||
|
|
logmel_feat = logmel_feat.masked_fill(make_pad_mask(ilens, logmel_feat, 1), 0.0)
|
||
|
|
return logmel_feat, ilens
|
||
|
|
|
||
|
|
|
||
|
|
class GlobalMVN(torch.nn.Module):
|
||
|
|
"""Apply global mean and variance normalization
|
||
|
|
|
||
|
|
Args:
|
||
|
|
stats_file(str): npy file of 1-dim array or text file.
|
||
|
|
From the _first element to
|
||
|
|
the {(len(array) - 1) / 2}th element are treated as
|
||
|
|
the sum of features,
|
||
|
|
and the rest excluding the last elements are
|
||
|
|
treated as the sum of the square value of features,
|
||
|
|
and the last elements eqauls to the number of samples.
|
||
|
|
std_floor(float):
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
stats_file: str,
|
||
|
|
norm_means: bool = True,
|
||
|
|
norm_vars: bool = True,
|
||
|
|
eps: float = 1.0e-20,
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.norm_means = norm_means
|
||
|
|
self.norm_vars = norm_vars
|
||
|
|
|
||
|
|
self.stats_file = stats_file
|
||
|
|
stats = np.load(stats_file)
|
||
|
|
|
||
|
|
stats = stats.astype(float)
|
||
|
|
assert (len(stats) - 1) % 2 == 0, stats.shape
|
||
|
|
|
||
|
|
count = stats.flatten()[-1]
|
||
|
|
mean = stats[: (len(stats) - 1) // 2] / count
|
||
|
|
var = stats[(len(stats) - 1) // 2 : -1] / count - mean * mean
|
||
|
|
std = np.maximum(np.sqrt(var), eps)
|
||
|
|
|
||
|
|
self.register_buffer("bias", torch.from_numpy(-mean.astype(np.float32)))
|
||
|
|
self.register_buffer("scale", torch.from_numpy(1 / std.astype(np.float32)))
|
||
|
|
|
||
|
|
def extra_repr(self):
|
||
|
|
return (
|
||
|
|
f"stats_file={self.stats_file}, "
|
||
|
|
f"norm_means={self.norm_means}, norm_vars={self.norm_vars}"
|
||
|
|
)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self, x: torch.Tensor, ilens: torch.LongTensor
|
||
|
|
) -> Tuple[torch.Tensor, torch.LongTensor]:
|
||
|
|
# feat: (B, T, D)
|
||
|
|
if self.norm_means:
|
||
|
|
x += self.bias.type_as(x)
|
||
|
|
x.masked_fill(make_pad_mask(ilens, x, 1), 0.0)
|
||
|
|
|
||
|
|
if self.norm_vars:
|
||
|
|
x *= self.scale.type_as(x)
|
||
|
|
return x, ilens
|
||
|
|
|
||
|
|
|
||
|
|
class UtteranceMVN(torch.nn.Module):
|
||
|
|
def __init__(self, norm_means: bool = True, norm_vars: bool = False, eps: float = 1.0e-20):
|
||
|
|
super().__init__()
|
||
|
|
self.norm_means = norm_means
|
||
|
|
self.norm_vars = norm_vars
|
||
|
|
self.eps = eps
|
||
|
|
|
||
|
|
def extra_repr(self):
|
||
|
|
return f"norm_means={self.norm_means}, norm_vars={self.norm_vars}"
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self, x: torch.Tensor, ilens: torch.LongTensor
|
||
|
|
) -> Tuple[torch.Tensor, torch.LongTensor]:
|
||
|
|
return utterance_mvn(
|
||
|
|
x, ilens, norm_means=self.norm_means, norm_vars=self.norm_vars, eps=self.eps
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def utterance_mvn(
|
||
|
|
x: torch.Tensor,
|
||
|
|
ilens: torch.LongTensor,
|
||
|
|
norm_means: bool = True,
|
||
|
|
norm_vars: bool = False,
|
||
|
|
eps: float = 1.0e-20,
|
||
|
|
) -> Tuple[torch.Tensor, torch.LongTensor]:
|
||
|
|
"""Apply utterance mean and variance normalization
|
||
|
|
|
||
|
|
Args:
|
||
|
|
x: (B, T, D), assumed zero padded
|
||
|
|
ilens: (B, T, D)
|
||
|
|
norm_means:
|
||
|
|
norm_vars:
|
||
|
|
eps:
|
||
|
|
|
||
|
|
"""
|
||
|
|
ilens_ = ilens.type_as(x)
|
||
|
|
# mean: (B, D)
|
||
|
|
mean = x.sum(dim=1) / ilens_[:, None]
|
||
|
|
|
||
|
|
if norm_means:
|
||
|
|
x -= mean[:, None, :]
|
||
|
|
x_ = x
|
||
|
|
else:
|
||
|
|
x_ = x - mean[:, None, :]
|
||
|
|
|
||
|
|
# Zero padding
|
||
|
|
x_.masked_fill(make_pad_mask(ilens, x_, 1), 0.0)
|
||
|
|
if norm_vars:
|
||
|
|
var = x_.pow(2).sum(dim=1) / ilens_[:, None]
|
||
|
|
var = torch.clamp(var, min=eps)
|
||
|
|
x /= var.sqrt()[:, None, :]
|
||
|
|
x_ = x
|
||
|
|
return x_, ilens
|
||
|
|
|
||
|
|
|
||
|
|
def feature_transform_for(args, n_fft):
|
||
|
|
return FeatureTransform(
|
||
|
|
# Mel options,
|
||
|
|
fs=args.fbank_fs,
|
||
|
|
n_fft=n_fft,
|
||
|
|
n_mels=args.n_mels,
|
||
|
|
fmin=args.fbank_fmin,
|
||
|
|
fmax=args.fbank_fmax,
|
||
|
|
# Normalization
|
||
|
|
stats_file=args.stats_file,
|
||
|
|
apply_uttmvn=args.apply_uttmvn,
|
||
|
|
uttmvn_norm_means=args.uttmvn_norm_means,
|
||
|
|
uttmvn_norm_vars=args.uttmvn_norm_vars,
|
||
|
|
)
|