Spaces:
Runtime error
Runtime error
# Source: https://github.com/snakers4/silero-vad | |
# | |
# Copyright (c) 2024 snakers4 | |
# | |
# This code is from a MIT-licensed repository. The full license text is available at the root of the source repository. | |
# | |
# Note: This code has been modified to fit the context of this repository. | |
import librosa | |
import torch | |
import numpy as np | |
VAD_THRESHOLD = 20 | |
SAMPLING_RATE = 16000 | |
class SileroVAD: | |
""" | |
Voice Activity Detection (VAD) using Silero-VAD. | |
""" | |
def __init__(self, local=False, model="silero_vad", device=torch.device("cpu")): | |
""" | |
Initialize the VAD object. | |
Args: | |
local (bool, optional): Whether to load the model locally. Defaults to False. | |
model (str, optional): The VAD model name to load. Defaults to "silero_vad". | |
device (torch.device, optional): The device to run the model on. Defaults to 'cpu'. | |
Returns: | |
None | |
Raises: | |
RuntimeError: If loading the model fails. | |
""" | |
try: | |
vad_model, utils = torch.hub.load( | |
repo_or_dir="snakers4/silero-vad" if not local else "vad/silero-vad", | |
model=model, | |
force_reload=False, | |
onnx=True, | |
source="github" if not local else "local", | |
) | |
self.vad_model = vad_model | |
(get_speech_timestamps, _, _, _, _) = utils | |
self.get_speech_timestamps = get_speech_timestamps | |
except Exception as e: | |
raise RuntimeError(f"Failed to load VAD model: {e}") | |
def segment_speech(self, audio_segment, start_time, end_time, sampling_rate): | |
""" | |
Segment speech from an audio segment and return a list of timestamps. | |
Args: | |
audio_segment (np.ndarray): The audio segment to be segmented. | |
start_time (int): The start time of the audio segment in frames. | |
end_time (int): The end time of the audio segment in frames. | |
sampling_rate (int): The sampling rate of the audio segment. | |
Returns: | |
list: A list of timestamps, each containing the start and end times of speech segments in frames. | |
Raises: | |
ValueError: If the audio segment is invalid. | |
""" | |
if audio_segment is None or not isinstance(audio_segment, (np.ndarray, list)): | |
raise ValueError("Invalid audio segment") | |
speech_timestamps = self.get_speech_timestamps( | |
audio_segment, self.vad_model, sampling_rate=sampling_rate | |
) | |
adjusted_timestamps = [ | |
(ts["start"] + start_time, ts["end"] + start_time) | |
for ts in speech_timestamps | |
] | |
if not adjusted_timestamps: | |
return [] | |
intervals = [ | |
end[0] - start[1] | |
for start, end in zip(adjusted_timestamps[:-1], adjusted_timestamps[1:]) | |
] | |
segments = [] | |
def split_timestamps(start_index, end_index): | |
if ( | |
start_index == end_index | |
or adjusted_timestamps[end_index][1] | |
- adjusted_timestamps[start_index][0] | |
< 20 * sampling_rate | |
): | |
segments.append([start_index, end_index]) | |
else: | |
if not intervals[start_index:end_index]: | |
return | |
max_interval_index = intervals[start_index:end_index].index( | |
max(intervals[start_index:end_index]) | |
) | |
split_index = start_index + max_interval_index | |
split_timestamps(start_index, split_index) | |
split_timestamps(split_index + 1, end_index) | |
split_timestamps(0, len(adjusted_timestamps) - 1) | |
merged_timestamps = [ | |
[adjusted_timestamps[start][0], adjusted_timestamps[end][1]] | |
for start, end in segments | |
] | |
return merged_timestamps | |
def vad(self, speakerdia, audio): | |
""" | |
Process the audio based on the given speaker diarization dataframe. | |
Args: | |
speakerdia (pd.DataFrame): The diarization dataframe containing start, end, and speaker info. | |
audio (dict): A dictionary containing the audio waveform and sample rate. | |
Returns: | |
list: A list of dictionaries containing processed audio segments with start, end, and speaker. | |
""" | |
sampling_rate = audio["sample_rate"] | |
audio_data = audio["waveform"] | |
out = [] | |
last_end = 0 | |
speakers_seen = set() | |
count_id = 0 | |
for index, row in speakerdia.iterrows(): | |
start = float(row["start"]) | |
end = float(row["end"]) | |
if end <= last_end: | |
continue | |
last_end = end | |
start_frame = int(start * sampling_rate) | |
end_frame = int(end * sampling_rate) | |
if row["speaker"] not in speakers_seen: | |
speakers_seen.add(row["speaker"]) | |
if end - start <= VAD_THRESHOLD: | |
out.append( | |
{ | |
"index": str(count_id).zfill(5), | |
"start": start, # in seconds | |
"end": end, | |
"speaker": row["speaker"], # same for all | |
} | |
) | |
count_id += 1 | |
continue | |
temp_audio = audio_data[start_frame:end_frame] | |
# resample from 24k to 16k | |
temp_audio_resampled = librosa.resample( | |
temp_audio, orig_sr=sampling_rate, target_sr=SAMPLING_RATE | |
) | |
for start_frame_sub, end_frame_sub in self.segment_speech( | |
temp_audio_resampled, | |
int(start * SAMPLING_RATE), | |
int(end * SAMPLING_RATE), | |
SAMPLING_RATE, | |
): | |
out.append( | |
{ | |
"index": str(count_id).zfill(5), | |
"start": start_frame_sub / SAMPLING_RATE, # in seconds | |
"end": end_frame_sub / SAMPLING_RATE, | |
"speaker": row["speaker"], # same for all | |
} | |
) | |
count_id += 1 | |
return out | |