|
import torch
|
|
from speechbrain.inference.interfaces import Pretrained
|
|
import torchaudio
|
|
|
|
|
|
def merge_overlapping_segments(segments):
|
|
"""
|
|
Merges segments that overlap or are contiguous, ensuring each speaker segment is represented once.
|
|
|
|
Args:
|
|
segments (list of tuples): List of tuples representing (start, end, label) of segments.
|
|
|
|
Returns:
|
|
list of tuples: Merged list of segments.
|
|
"""
|
|
if not segments:
|
|
return []
|
|
merged = [segments[0]]
|
|
for current in segments[1:]:
|
|
prev = merged[-1]
|
|
if current[0] <= prev[1]:
|
|
if current[2] == prev[2]:
|
|
merged[-1] = (prev[0], max(prev[1], current[1]), prev[2])
|
|
else:
|
|
merged.append(current)
|
|
else:
|
|
merged.append(current)
|
|
return merged
|
|
|
|
|
|
def refine_transitions(aggregated_predictions):
|
|
"""
|
|
Refines transitions between speaker segments to enhance accuracy.
|
|
|
|
Args:
|
|
aggregated_predictions (list of tuples): The aggregated predictions with potential overlaps.
|
|
|
|
Returns:
|
|
list of tuples: Predictions with adjusted transitions.
|
|
"""
|
|
refined_predictions = []
|
|
for i in range(len(aggregated_predictions)):
|
|
if i == 0:
|
|
refined_predictions.append(aggregated_predictions[i])
|
|
continue
|
|
|
|
current_start, current_end, current_label = aggregated_predictions[i]
|
|
prev_start, prev_end, prev_label = aggregated_predictions[i - 1]
|
|
|
|
if current_start - prev_end <= 1.0:
|
|
new_start = prev_end
|
|
else:
|
|
new_start = current_start
|
|
|
|
refined_predictions.append((new_start, current_end, current_label))
|
|
|
|
return refined_predictions
|
|
|
|
|
|
def refine_transitions_with_confidence(aggregated_predictions, segment_confidences):
|
|
"""
|
|
Refines transitions between segments based on confidence levels.
|
|
|
|
Args:
|
|
aggregated_predictions (list of tuples): Initial aggregated predictions.
|
|
segment_confidences (list of float): Confidence scores corresponding to each segment.
|
|
|
|
Returns:
|
|
list of tuples: Refined segment predictions.
|
|
"""
|
|
refined_predictions = []
|
|
for i in range(len(aggregated_predictions)):
|
|
if i == 0:
|
|
refined_predictions.append(aggregated_predictions[i])
|
|
continue
|
|
|
|
current_start, current_end, current_label = aggregated_predictions[i]
|
|
prev_start, prev_end, prev_label, prev_confidence = refined_predictions[-1] + (segment_confidences[i - 1],)
|
|
|
|
current_confidence = segment_confidences[i]
|
|
|
|
if current_label != prev_label:
|
|
if prev_confidence < current_confidence:
|
|
transition_point = current_start
|
|
else:
|
|
transition_point = prev_end
|
|
refined_predictions[-1] = (prev_start, transition_point, prev_label)
|
|
refined_predictions.append((transition_point, current_end, current_label))
|
|
else:
|
|
if prev_confidence < current_confidence:
|
|
refined_predictions[-1] = (prev_start, current_end, current_label)
|
|
else:
|
|
refined_predictions.append((current_start, current_end, current_label))
|
|
|
|
return refined_predictions
|
|
|
|
|
|
def aggregate_segments_with_overlap(segment_predictions):
|
|
"""
|
|
Aggregates overlapping segments into single segments based on speaker labels.
|
|
|
|
Args:
|
|
segment_predictions (list of tuples): List of tuples representing (start, end, label) of segments.
|
|
|
|
Returns:
|
|
list of tuples: Aggregated segments.
|
|
"""
|
|
aggregated_predictions = []
|
|
last_start, last_end, last_label = segment_predictions[0]
|
|
|
|
for start, end, label in segment_predictions[1:]:
|
|
if label == last_label and start <= last_end:
|
|
last_end = max(last_end, end)
|
|
else:
|
|
aggregated_predictions.append((last_start, last_end, last_label))
|
|
last_start, last_end, last_label = start, end, label
|
|
|
|
aggregated_predictions.append((last_start, last_end, last_label))
|
|
|
|
merged = merge_overlapping_segments(aggregated_predictions)
|
|
return merged
|
|
|
|
|
|
class SpeakerCounter(Pretrained):
|
|
"""
|
|
A class for counting speakers in an audio file, built upon the SpeechBrain Pretrained class.
|
|
This class integrates several preprocessing and prediction modules to handle speaker diarization tasks.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
"""
|
|
Initialize the SpeakerCounter with standard and custom parameters.
|
|
Args:
|
|
*args: Variable length argument list.
|
|
**kwargs: Arbitrary keyword arguments.
|
|
"""
|
|
super().__init__(*args, **kwargs)
|
|
self.sample_rate = self.hparams.sample_rate
|
|
|
|
MODULES_NEEDED = [
|
|
"compute_features",
|
|
"mean_var_norm",
|
|
"embedding_model",
|
|
"classifier",
|
|
]
|
|
|
|
def resample_waveform(self, waveform, orig_sample_rate):
|
|
"""
|
|
Resamples the input waveform to the target sample rate specified in the object.
|
|
|
|
Args:
|
|
waveform (Tensor): The input waveform tensor.
|
|
orig_sample_rate (int): The original sample rate of the waveform.
|
|
|
|
Returns:
|
|
Tensor: The resampled waveform.
|
|
"""
|
|
if orig_sample_rate != self.sample_rate:
|
|
resample_transform = torchaudio.transforms.Resample(orig_freq=orig_sample_rate, new_freq=self.sample_rate)
|
|
waveform = resample_transform(waveform)
|
|
return waveform
|
|
|
|
def encode_batch(self, wavs, wav_lens=None):
|
|
"""
|
|
Encodes a batch of waveforms into embeddings using the loaded models.
|
|
|
|
Args:
|
|
wavs (Tensor): Batch of waveforms.
|
|
wav_lens (Tensor, optional): Lengths of the waveforms for normalization.
|
|
|
|
Returns:
|
|
Tensor: Batch of embeddings.
|
|
"""
|
|
if len(wavs.shape) == 1:
|
|
wavs = wavs.unsqueeze(0)
|
|
|
|
if wav_lens is None:
|
|
wav_lens = torch.ones(wavs.shape[0], device=self.device)
|
|
|
|
wavs, wav_lens = wavs.to(self.device), wav_lens.to(self.device)
|
|
wavs = wavs.float()
|
|
|
|
|
|
feats = self.mods.compute_features(wavs)
|
|
feats = self.mods.mean_var_norm(feats, wav_lens)
|
|
embeddings = self.mods.embedding_model(feats, wav_lens)
|
|
return embeddings
|
|
|
|
def create_segments(self, waveform, segment_length, overlap):
|
|
"""
|
|
Creates segments from a single waveform for batch processing.
|
|
|
|
Args:
|
|
waveform (Tensor): Input waveform tensor.
|
|
segment_length (float): Length of each segment in seconds.
|
|
overlap (float): Overlap between segments in seconds.
|
|
|
|
Returns:
|
|
tuple: (segments, segment_times) where segments is a list of tensors, and segment_times
|
|
is a list of (start, end) times.
|
|
"""
|
|
num_samples = waveform.shape[1]
|
|
segment_samples = int(segment_length * self.sample_rate)
|
|
overlap_samples = int(overlap * self.sample_rate)
|
|
step_samples = segment_samples - overlap_samples
|
|
segments = []
|
|
segment_times = []
|
|
|
|
for start in range(0, num_samples - segment_samples + 1, step_samples):
|
|
end = start + segment_samples
|
|
segments.append(waveform[:, start:end])
|
|
start_time = start / self.sample_rate
|
|
end_time = end / self.sample_rate
|
|
segment_times.append((start_time, end_time))
|
|
|
|
return segments, segment_times
|
|
|
|
def classify_file(self, path, segment_length=2.0, overlap=1.47):
|
|
"""
|
|
Processes an audio file to classify and count speakers within segments.
|
|
Utilizes multiple stages of processing to handle overlapping speech and transitions.
|
|
|
|
Args:
|
|
path (str): Path to the audio file.
|
|
segment_length (float): Length of each segment in seconds.
|
|
overlap (float): Overlap between segments in seconds.
|
|
|
|
Outputs:
|
|
Writes the number of speakers in each segment to a text file.
|
|
"""
|
|
waveform, osr = torchaudio.load(path)
|
|
waveform = self.resample_waveform(waveform, osr)
|
|
|
|
segments, segment_times = self.create_segments(waveform, segment_length, overlap)
|
|
segment_predictions = []
|
|
|
|
for segment, (start_time, end_time) in zip(segments, segment_times):
|
|
rel_length = torch.tensor([1.0])
|
|
emb = self.encode_batch(segment, rel_length)
|
|
out_prob = self.mods.classifier(emb).squeeze(1)
|
|
score, index = torch.max(out_prob, dim=-1)
|
|
text_lab = index.item()
|
|
segment_predictions.append((start_time, end_time, text_lab))
|
|
|
|
aggregated_predictions = aggregate_segments_with_overlap(segment_predictions)
|
|
refined_predictions = refine_transitions(aggregated_predictions)
|
|
preds = refine_transitions_with_confidence(aggregated_predictions, refined_predictions)
|
|
|
|
with open("sample_segment_predictions.txt", "w") as file:
|
|
for start_time, end_time, prediction in preds:
|
|
speaker_text = "no speech" if str(prediction) == "0" else (
|
|
"1 speaker" if str(prediction) == "1" else f"{prediction} speakers")
|
|
print(f"{start_time:.2f}-{end_time:.2f} has {speaker_text}")
|
|
file.write(f"{start_time:.2f}-{end_time:.2f} has {speaker_text}\n")
|
|
|
|
def forward(self, wavs, wav_lens=None):
|
|
"""
|
|
Forward pass for classifying audio using preloaded modules.
|
|
|
|
Args:
|
|
wavs (Tensor): Input waveforms.
|
|
wav_lens (Tensor, optional): Lengths of the input waveforms.
|
|
|
|
Returns:
|
|
Output from classify_file method.
|
|
"""
|
|
return self.classify_file(wavs, wav_lens)
|
|
|
