mirror of
https://github.com/snakers4/silero-vad.git
synced 2026-02-05 01:49:22 +08:00
300 lines
9.6 KiB
Python
300 lines
9.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import Dataset
|
|
import torchaudio
|
|
import numpy as np
|
|
import random
|
|
import gc
|
|
from sklearn.metrics import roc_auc_score
|
|
from tqdm import tqdm
|
|
import pandas as pd
|
|
import warnings
|
|
warnings.filterwarnings('ignore')
|
|
|
|
|
|
def read_audio(path: str,
|
|
sampling_rate: int = 16000,
|
|
normalize=False):
|
|
|
|
wav, sr = torchaudio.load(path)
|
|
|
|
if wav.size(0) > 1:
|
|
wav = wav.mean(dim=0, keepdim=True)
|
|
|
|
if sampling_rate:
|
|
if sr != sampling_rate:
|
|
transform = torchaudio.transforms.Resample(orig_freq=sr,
|
|
new_freq=sampling_rate)
|
|
wav = transform(wav)
|
|
sr = sampling_rate
|
|
|
|
if normalize and wav.abs().max() != 0:
|
|
wav = wav / wav.abs().max()
|
|
|
|
return wav.squeeze(0)
|
|
|
|
|
|
def build_audiomentations_augs(p):
|
|
from audiomentations import SomeOf, AirAbsorption, BandPassFilter, BandStopFilter, ClippingDistortion, HighPassFilter, HighShelfFilter, \
|
|
LowPassFilter, LowShelfFilter, Mp3Compression, PeakingFilter, PitchShift, RoomSimulator, SevenBandParametricEQ, \
|
|
Aliasing, AddGaussianNoise
|
|
transforms = [Aliasing(p=1),
|
|
AddGaussianNoise(p=1),
|
|
AirAbsorption(p=1),
|
|
BandPassFilter(p=1),
|
|
BandStopFilter(p=1),
|
|
ClippingDistortion(p=1),
|
|
HighPassFilter(p=1),
|
|
HighShelfFilter(p=1),
|
|
LowPassFilter(p=1),
|
|
LowShelfFilter(p=1),
|
|
Mp3Compression(p=1),
|
|
PeakingFilter(p=1),
|
|
PitchShift(p=1),
|
|
RoomSimulator(p=1, leave_length_unchanged=True),
|
|
SevenBandParametricEQ(p=1)]
|
|
tr = SomeOf((1, 3), transforms=transforms, p=p)
|
|
return tr
|
|
|
|
|
|
class SileroVadDataset(Dataset):
|
|
def __init__(self,
|
|
config,
|
|
mode='train'):
|
|
|
|
self.num_samples = 512 # constant, do not change
|
|
self.sr = 16000 # constant, do not change
|
|
|
|
self.resample_to_8k = config.tune_8k
|
|
self.noise_loss = config.noise_loss
|
|
self.max_train_length_sec = config.max_train_length_sec
|
|
self.max_train_length_samples = config.max_train_length_sec * self.sr
|
|
|
|
assert self.max_train_length_samples % self.num_samples == 0
|
|
assert mode in ['train', 'val']
|
|
|
|
dataset_path = config.train_dataset_path if mode == 'train' else config.val_dataset_path
|
|
self.dataframe = pd.read_feather(dataset_path).reset_index(drop=True)
|
|
self.index_dict = self.dataframe.to_dict('index')
|
|
self.mode = mode
|
|
print(f'DATASET SIZE : {len(self.dataframe)}')
|
|
|
|
if mode == 'train':
|
|
self.augs = build_audiomentations_augs(p=config.aug_prob)
|
|
else:
|
|
self.augs = None
|
|
|
|
def __getitem__(self, idx):
|
|
idx = None if self.mode == 'train' else idx
|
|
wav, gt, mask = self.load_speech_sample(idx)
|
|
|
|
if self.mode == 'train':
|
|
wav = self.add_augs(wav)
|
|
if len(wav) > self.max_train_length_samples:
|
|
wav = wav[:self.max_train_length_samples]
|
|
gt = gt[:int(self.max_train_length_samples / self.num_samples)]
|
|
mask = mask[:int(self.max_train_length_samples / self.num_samples)]
|
|
|
|
wav = torch.FloatTensor(wav)
|
|
if self.resample_to_8k:
|
|
transform = torchaudio.transforms.Resample(orig_freq=self.sr,
|
|
new_freq=8000)
|
|
wav = transform(wav)
|
|
return wav, torch.FloatTensor(gt), torch.from_numpy(mask)
|
|
|
|
def __len__(self):
|
|
return len(self.index_dict)
|
|
|
|
def load_speech_sample(self, idx=None):
|
|
if idx is None:
|
|
idx = random.randint(0, len(self.index_dict) - 1)
|
|
wav = read_audio(self.index_dict[idx]['audio_path'], self.sr).numpy()
|
|
|
|
if len(wav) % self.num_samples != 0:
|
|
pad_num = self.num_samples - (len(wav) % (self.num_samples))
|
|
wav = np.pad(wav, (0, pad_num), 'constant', constant_values=0)
|
|
|
|
gt, mask = self.get_ground_truth_annotated(self.index_dict[idx]['speech_ts'], len(wav))
|
|
|
|
assert len(gt) == len(wav) / self.num_samples
|
|
|
|
mask[gt == 0]
|
|
|
|
return wav, gt, mask
|
|
|
|
def get_ground_truth_annotated(self, annotation, audio_length_samples):
|
|
gt = np.zeros(audio_length_samples)
|
|
|
|
for i in annotation:
|
|
gt[int(i['start'] * self.sr): int(i['end'] * self.sr)] = 1
|
|
|
|
squeezed_predicts = np.average(gt.reshape(-1, self.num_samples), axis=1)
|
|
squeezed_predicts = (squeezed_predicts > 0.5).astype(int)
|
|
mask = np.ones(len(squeezed_predicts))
|
|
mask[squeezed_predicts == 0] = self.noise_loss
|
|
return squeezed_predicts, mask
|
|
|
|
def add_augs(self, wav):
|
|
while True:
|
|
try:
|
|
wav_aug = self.augs(wav, self.sr)
|
|
if np.isnan(wav_aug.max()) or np.isnan(wav_aug.min()):
|
|
return wav
|
|
return wav_aug
|
|
except Exception as e:
|
|
continue
|
|
|
|
|
|
def SileroVadPadder(batch):
|
|
wavs = [batch[i][0] for i in range(len(batch))]
|
|
labels = [batch[i][1] for i in range(len(batch))]
|
|
masks = [batch[i][2] for i in range(len(batch))]
|
|
|
|
wavs = torch.nn.utils.rnn.pad_sequence(
|
|
wavs, batch_first=True, padding_value=0)
|
|
|
|
labels = torch.nn.utils.rnn.pad_sequence(
|
|
labels, batch_first=True, padding_value=0)
|
|
|
|
masks = torch.nn.utils.rnn.pad_sequence(
|
|
masks, batch_first=True, padding_value=0)
|
|
|
|
return wavs, labels, masks
|
|
|
|
|
|
class VADDecoderRNNJIT(nn.Module):
|
|
|
|
def __init__(self):
|
|
super(VADDecoderRNNJIT, self).__init__()
|
|
|
|
self.rnn = nn.LSTMCell(128, 128)
|
|
self.decoder = nn.Sequential(nn.Dropout(0.1),
|
|
nn.ReLU(),
|
|
nn.Conv1d(128, 1, kernel_size=1),
|
|
nn.Sigmoid())
|
|
|
|
def forward(self, x, state=torch.zeros(0)):
|
|
x = x.squeeze(-1)
|
|
if len(state):
|
|
h, c = self.rnn(x, (state[0], state[1]))
|
|
else:
|
|
h, c = self.rnn(x)
|
|
|
|
x = h.unsqueeze(-1).float()
|
|
state = torch.stack([h, c])
|
|
x = self.decoder(x)
|
|
return x, state
|
|
|
|
|
|
class AverageMeter(object):
|
|
"""Computes and stores the average and current value"""
|
|
|
|
def __init__(self):
|
|
self.reset()
|
|
|
|
def reset(self):
|
|
self.val = 0
|
|
self.avg = 0
|
|
self.sum = 0
|
|
self.count = 0
|
|
|
|
def update(self, val, n=1):
|
|
self.val = val
|
|
self.sum += val * n
|
|
self.count += n
|
|
self.avg = self.sum / self.count
|
|
|
|
|
|
def train(config,
|
|
loader,
|
|
jit_model,
|
|
decoder,
|
|
criterion,
|
|
optimizer,
|
|
device):
|
|
|
|
losses = AverageMeter()
|
|
decoder.train()
|
|
|
|
context_size = 32 if config.tune_8k else 64
|
|
num_samples = 256 if config.tune_8k else 512
|
|
stft_layer = jit_model._model_8k.stft if config.tune_8k else jit_model._model.stft
|
|
encoder_layer = jit_model._model_8k.encoder if config.tune_8k else jit_model._model.encoder
|
|
|
|
with torch.enable_grad():
|
|
for _, (x, targets, masks) in tqdm(enumerate(loader), total=len(loader)):
|
|
targets = targets.to(device)
|
|
x = x.to(device)
|
|
masks = masks.to(device)
|
|
x = torch.nn.functional.pad(x, (context_size, 0))
|
|
|
|
outs = []
|
|
state = torch.zeros(0)
|
|
for i in range(context_size, x.shape[1], num_samples):
|
|
input_ = x[:, i-context_size:i+num_samples]
|
|
out = stft_layer(input_)
|
|
out = encoder_layer(out)
|
|
out, state = decoder(out, state)
|
|
outs.append(out)
|
|
stacked = torch.cat(outs, dim=2).squeeze(1)
|
|
|
|
loss = criterion(stacked, targets)
|
|
loss = (loss * masks).mean()
|
|
loss.backward()
|
|
optimizer.step()
|
|
losses.update(loss.item(), masks.numel())
|
|
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return losses.avg
|
|
|
|
|
|
def validate(config,
|
|
loader,
|
|
jit_model,
|
|
decoder,
|
|
criterion,
|
|
device):
|
|
|
|
losses = AverageMeter()
|
|
decoder.eval()
|
|
|
|
predicts = []
|
|
gts = []
|
|
|
|
context_size = 32 if config.tune_8k else 64
|
|
num_samples = 256 if config.tune_8k else 512
|
|
stft_layer = jit_model._model_8k.stft if config.tune_8k else jit_model._model.stft
|
|
encoder_layer = jit_model._model_8k.encoder if config.tune_8k else jit_model._model.encoder
|
|
|
|
with torch.no_grad():
|
|
for _, (x, targets, masks) in tqdm(enumerate(loader), total=len(loader)):
|
|
targets = targets.to(device)
|
|
x = x.to(device)
|
|
masks = masks.to(device)
|
|
x = torch.nn.functional.pad(x, (context_size, 0))
|
|
|
|
outs = []
|
|
state = torch.zeros(0)
|
|
for i in range(context_size, x.shape[1], num_samples):
|
|
input_ = x[:, i-context_size:i+num_samples]
|
|
out = stft_layer(input_)
|
|
out = encoder_layer(out)
|
|
out, state = decoder(out, state)
|
|
outs.append(out)
|
|
stacked = torch.cat(outs, dim=2).squeeze(1)
|
|
|
|
predicts.extend(stacked[masks != 0].tolist())
|
|
gts.extend(targets[masks != 0].tolist())
|
|
|
|
loss = criterion(stacked, targets)
|
|
loss = (loss * masks).mean()
|
|
losses.update(loss.item(), masks.numel())
|
|
score = roc_auc_score(gts, predicts)
|
|
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
return losses.avg, round(score, 3)
|