Spaces:
Running
Running
| #importing all the necessary packages | |
| import torch | |
| import transformers | |
| import gradio as gr | |
| from torchaudio.sox_effects import apply_effects_file | |
| from termcolor import colored | |
| from transformers import Wav2Vec2FeatureExtractor, UniSpeechSatForAudioFrameClassification | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Defines the effects to apply to the audio file | |
| EFFECTS = [ | |
| ['remix', '-'], # merge all the channels | |
| ["channels", "1"], #channel-->mono | |
| ["rate", "16000"], # resample to 16000 Hz | |
| ["gain", "-1.0"], #Attenuation -1 dB | |
| ["silence", "1", "0.1", "0.1%", "-1", "0.1", "0.1%"], | |
| #['pad', '0', '1.5'], # add 1.5 seconds silence at the end | |
| ['trim', '0', '10'], # get the first 10 seconds | |
| ] | |
| THRESHOLD = 0.85 #depends on dataset | |
| model_name = "microsoft/unispeech-sat-base-sd" | |
| feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name) | |
| model = UniSpeechSatForAudioFrameClassification.from_pretrained(model_name).to(device) | |
| def fn(path): | |
| #Applying the effects to the audio input file | |
| wav, _ = apply_effects_file(path, EFFECTS) | |
| #Extracting features | |
| input = feature_extractor(wav.squeeze(0), return_tensors="pt", sampling_rate=16000).input_values.to(device) | |
| with torch.no_grad(): | |
| logits = model(input).logits | |
| logits = logits.to(device) | |
| probabilities = torch.sigmoid(logits[0]) | |
| # labels is a one-hot array of shape (num_frames, num_speakers) | |
| labels = (probabilities > 0.5).long() | |
| return labels | |
| inputs = [ | |
| gr.inputs.Audio(source="microphone", type="filepath", optional=True, label="Speaker #1"), | |
| ] | |
| output = gr.outputs.Textbox(label="Output Text") | |
| gr.Interface( | |
| fn=fn, | |
| inputs=inputs, | |
| outputs=output, | |
| theme = "grass", | |
| title="Speaker diarization using UniSpeech-SAT and X-Vectors").launch(enable_queue=True) | |