File size: 6,356 Bytes
5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 2e9a543 5e9bb10 |
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 |
# coding=utf-8
# Copyright 2023 Xiaomi Corporation and The HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Feature extractor class for CED.
"""
from typing import List, Optional, Union
import numpy as np
import torch
import torchaudio.transforms as audio_transforms
from transformers.feature_extraction_sequence_utils import SequenceFeatureExtractor
from transformers.feature_extraction_utils import BatchFeature
from transformers.utils import logging
logger = logging.get_logger(__name__)
class CedFeatureExtractor(SequenceFeatureExtractor):
r"""
CedFeatureExtractor extracts Mel spectrogram features from audio signals.
Args:
f_min (int, *optional*, defaults to 0): Minimum frequency for the Mel filterbank.
sampling_rate (int, *optional*, defaults to 16000):
Sampling rate of the input audio signal.
win_size (int, *optional*, defaults to 512): Window size for the STFT.
center (bool, *optional*, defaults to `True`):
Whether to pad the signal on both sides to center it.
n_fft (int, *optional*, defaults to 512): Number of FFT points for the STFT.
f_max (int, optional, *optional*): Maximum frequency for the Mel filterbank.
hop_size (int, *optional*, defaults to 160): Hop size for the STFT.
feature_size (int, *optional*, defaults to 64): Number of Mel bands to generate.
padding_value (float, *optional*, defaults to 0.0): Value for padding.
Returns:
BatchFeature: A BatchFeature object containing the extracted features.
"""
def __init__(
self,
f_min: int = 0,
sampling_rate: int = 16000,
win_size: int = 512,
center: bool = True,
n_fft: int = 512,
f_max: Optional[int] = None,
hop_size: int = 160,
feature_size: int = 64,
padding_value: float = 0.0,
**kwargs,
):
super().__init__(
feature_size=feature_size,
sampling_rate=sampling_rate,
padding_value=padding_value,
**kwargs,
)
self.f_min = f_min
self.win_size = win_size
self.center = center
self.n_fft = n_fft
self.f_max = f_max
self.hop_size = hop_size
self.model_input_names = ["input_values"]
def __call__(
self,
x: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
sampling_rate: Optional[int] = None,
max_length: Optional[int] = 16000,
truncation: bool = True,
return_tensors="pt",
) -> BatchFeature:
r"""
Extracts Mel spectrogram features from an audio signal tensor.
Args:
x: Input audio signal tensor.
sampling_rate (int, *optional*, defaults to `None`):
Sampling rate of the input audio signal.
max_length (int, *optional*, defaults to 16000):
Maximum length of the input audio signal.
truncation (bool, *optional*, defaults to `True`):
Whether to truncate the input signal to max_length.
return_tensors (str, *optional*, defaults to "pt"):
If set to "pt", the return type will be a PyTorch tensor.
Returns:
BatchFeature: A dictionary containing the extracted features.
"""
if sampling_rate is None:
sampling_rate = self.sampling_rate
if return_tensors != "pt":
raise NotImplementedError("Only return_tensors='pt' is currently supported.")
mel_spectrogram = audio_transforms.MelSpectrogram(
f_min=self.f_min,
sample_rate=sampling_rate,
win_length=self.win_size,
center=self.center,
n_fft=self.n_fft,
f_max=self.f_max,
hop_length=self.hop_size,
n_mels=self.feature_size,
)
amplitude_to_db = audio_transforms.AmplitudeToDB(top_db=120)
if isinstance(x, np.ndarray):
if x.ndim == 1:
x = x[np.newaxis, :]
if x.ndim != 2:
raise ValueError("np.ndarray input must be a 1D or 2D.")
x = torch.from_numpy(x)
elif isinstance(x, torch.Tensor):
if x.dim() == 1:
x = x.unsqueeze(0)
if x.dim() != 2:
raise ValueError("torch.Tensor input must be a 1D or 2D.")
elif isinstance(x, (list, tuple)):
longest_length = max(x_.shape[0] for x_ in x)
if not truncation and max_length < longest_length:
max_length = longest_length
if all(isinstance(x_, np.ndarray) for x_ in x):
if not all(x_.ndim == 1 for x_ in x):
raise ValueError("All np.ndarray in a list must be 1D.")
x_trim = [x_[:max_length] for x_ in x]
x_pad = [np.pad(x_, (0, max_length - x_.shape[0]), mode="constant", constant_values=0) for x_ in x_trim]
x = torch.stack([torch.from_numpy(x_) for x_ in x_pad])
elif all(isinstance(x_, torch.Tensor) for x_ in x):
if not all(x_.dim() == 1 for x_ in x):
raise ValueError("All torch.Tensor in a list must be 1D.")
x_pad = [torch.nn.functional.pad(x_, (0, max_length - x_.shape[0]), value=0) for x_ in x]
x = torch.stack(x_pad)
else:
raise ValueError("Input list must be numpy arrays or PyTorch tensors.")
else:
raise ValueError(
"Input must be a numpy array, a list of numpy arrays, a PyTorch tensor, or a list of PyTorch tensor."
)
x = x.float()
x = mel_spectrogram(x)
x = amplitude_to_db(x)
return BatchFeature({"input_values": x})
|