114 lines
3.6 KiB
Python
Executable File
114 lines
3.6 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.dataio import audio_io
|
|
from speechbrain.dataio.preprocess import AudioNormalizer
|
|
from speechbrain.inference.speaker import EncoderClassifier
|
|
import torch
|
|
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(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
|
parser.add_argument("input")
|
|
parser.add_argument("-o", "--output", default="output.srt",
|
|
help="Output subtitle (SRT) file path")
|
|
parser.add_argument("-s", "--speakers", type=int, default=2,
|
|
help="Number of speakers in the audio")
|
|
parser.add_argument("-m", "--model", default='turbo', choices=whisper.available_models(),
|
|
help="Whisper model to use")
|
|
parser.add_argument("--gpu", action="store_true", help="run on GPU if available")
|
|
args = parser.parse_args()
|
|
|
|
|
|
has_gpu = torch.cuda.is_available()
|
|
if args.gpu and not has_gpu:
|
|
print("GPU not available. Falling back to CPU")
|
|
dev = torch.device("cuda" if has_gpu and args.gpu else "cpu")
|
|
|
|
model = whisper.load_model("turbo").to(dev)
|
|
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",
|
|
run_opts={"device": str(dev)},
|
|
)
|
|
|
|
signal, sample_rate = audio_io.load(args.input, channels_first=False)
|
|
audio_normalizer = AudioNormalizer()
|
|
signal = audio_normalizer(signal, sample_rate)
|
|
sample_rate = audio_normalizer.sample_rate
|
|
|
|
embedding_segments = []
|
|
embeddings = []
|
|
for seg in segments:
|
|
if seg.end - seg.start < 0.1:
|
|
# Segment is too short to bother trying to classify
|
|
continue
|
|
|
|
start_sample = int(seg.start * sample_rate)
|
|
end_sample = int(seg.end * sample_rate)
|
|
|
|
chunk = signal[start_sample:end_sample]
|
|
|
|
# normalize
|
|
chunk /= chunk.abs().max()
|
|
|
|
# SpeechBrain expects [batch, time]
|
|
emb = classifier.encode_batch(chunk.unsqueeze(0).to(dev))
|
|
emb = emb.squeeze().detach().cpu().numpy()
|
|
|
|
embeddings.append(emb)
|
|
embedding_segments.append(seg)
|
|
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(embedding_segments, labels):
|
|
seg.text = f"Speaker {label + 1}: {seg.text}"
|
|
|
|
print("Final transcription:")
|
|
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)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|