107 lines
2.9 KiB
Python
Executable File
107 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
import argparse
|
|
import dataclasses
|
|
|
|
import numpy as np
|
|
from sklearn.cluster import AgglomerativeClustering
|
|
from sklearn.preprocessing import normalize
|
|
from speechbrain.inference.speaker import EncoderClassifier
|
|
import torch
|
|
import torchaudio
|
|
import whisper
|
|
|
|
@dataclasses.dataclass
|
|
class Segment:
|
|
start: np.float32
|
|
end: np.float32
|
|
text: str
|
|
|
|
def srt_timestamp(milis):
|
|
h = milis // (1000 * 60 * 60)
|
|
milis %= 1000 * 60 * 60
|
|
m = milis // (1000 * 60)
|
|
milis %= 1000 * 60
|
|
s = milis // 1000
|
|
milis %= 1000
|
|
return f"{h:02}:{m:02}:{s:02},{milis:03}"
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("input")
|
|
parser.add_argument("-o", "--output", default="output.srt")
|
|
args = parser.parse_args()
|
|
|
|
model = whisper.load_model("turbo")
|
|
result = model.transcribe(args.input, word_timestamps=True, language="en")
|
|
segments = []
|
|
for seg in result["segments"]:
|
|
words = seg["words"]
|
|
start = words[0]["start"]
|
|
end = words[-1]["end"]
|
|
segments.append(Segment(start, end, seg["text"].strip()))
|
|
|
|
print("Transcription done. Starting speaker classification")
|
|
|
|
classifier = EncoderClassifier.from_hparams(
|
|
source="speechbrain/spkrec-ecapa-voxceleb",
|
|
)
|
|
|
|
signal, sample_rate = torchaudio.load(args.input)
|
|
|
|
print(sample_rate)
|
|
# resample to 16khz
|
|
if sample_rate != 16000:
|
|
resampler = torchaudio.transforms.Resample(sample_rate, 16000)
|
|
signal = resampler(signal)
|
|
sample_rate = 16000
|
|
|
|
# convert to mono if needed
|
|
if signal.shape[0] > 1:
|
|
signal = torch.mean(signal, dim=0, keepdim=True)
|
|
|
|
embeddings = []
|
|
for seg in segments:
|
|
start_sample = int(seg.start * sample_rate)
|
|
end_sample = int(seg.end * sample_rate)
|
|
|
|
chunk = signal[:, start_sample:end_sample]
|
|
|
|
# SpeechBrain expects [batch, time]
|
|
emb = classifier.encode_batch(chunk)
|
|
emb = emb.squeeze().detach().cpu().numpy()
|
|
|
|
embeddings.append(emb)
|
|
embeddings = np.vstack(embeddings)
|
|
embeddings = normalize(embeddings)
|
|
|
|
n_speakers = 2
|
|
clustering = AgglomerativeClustering(
|
|
n_clusters=n_speakers,
|
|
metric="cosine",
|
|
linkage="average"
|
|
)
|
|
labels = clustering.fit_predict(embeddings)
|
|
|
|
for seg, label in zip(segments, labels):
|
|
seg.text = f"Speaker {label + 1}: {seg.text}"
|
|
|
|
with open(args.output, "w") as srt_file:
|
|
for i, seg in enumerate(segments, 1):
|
|
start = srt_timestamp(round(seg.start * 1000))
|
|
end = srt_timestamp(round(seg.end * 1000))
|
|
srt_file.write(f"{i}\n{start} --> {end}\n{seg.text}\n\n")
|
|
print(start, end, seg.text)
|
|
|
|
from sklearn.decomposition import PCA
|
|
import matplotlib.pyplot as plt
|
|
|
|
pca = PCA(n_components=2)
|
|
points = pca.fit_transform(embeddings)
|
|
|
|
plt.scatter(points[:,0], points[:,1])
|
|
plt.show()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|