mvp

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from main_pipeline import main_pipeline
 from scipy.io.wavfile import write
 
-title = "audio_denoise and speakser diarization"
+title = "audio_denoise and speakser diarization. Fast inference https://t.me/diarizarion_bot"
 
 
 def app_pipeline(audio):
@@ -14,9 +14,15 @@ def app_pipeline(audio):
     return result_diarization + [None] * (10 - len(result_diarization))
 
 
+example_list = [
+    ["dialog.mp3"]
+]
 gr.Interface(
     app_pipeline,
     gr.Audio(type="numpy", label="Input"),
     [gr.Audio(visible=True) for i in range(10)],
     title=title,
+    examples=example_list,
+    cache_examples=False
+
 ).launch(enable_queue=True)

main_pipeline.py

@@ -21,7 +21,6 @@ def save_speaker_audios(segments, denoised_audio_path, out_folder='out', out_f=4
         for _, r in temp_df.iterrows():
             start = int(r["start"] * out_f)
             end = int(r["end"] * out_f)
-            # output_signal[start:end] = signal[start:end]
             output_signal.append(signal[start:end])
 
         out_wav_path = f'{out_folder}/{label}.wav'
@@ -42,4 +41,4 @@ def main_pipeline(audio_path):
 
 
 if __name__ == '__main__':
-    main_pipeline('out.wav')
+    main_pipeline('dialog.mp3')

utils/denoise_pipeline.py

@@ -10,9 +10,12 @@ from demucs.pretrained import get_model
 demucs_model = get_model('htdemucs')
 
 
-def denoise(filename, device):
+def denoise(filename: str, device: str) -> str:
+
     wav_ref, sr = librosa.load(filename, mono=False, sr=44100)
+
     wav = torch.tensor(wav_ref)
+    wav = torch.cat([wav.unsqueeze(0), wav.unsqueeze(0)]) if len(wav.shape) == 1 else wav
     ref = wav.mean(0)
     wav = (wav - ref.mean()) / wav.std()
     sources = apply_model(
@@ -32,4 +35,4 @@ def denoise(filename, device):
 
 if __name__ == '__main__':
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    denoise(filename='../out.wav', device=device)
+    denoise(filename='../oxx.wav', device=device)

utils/diarization_pipeline.py

@@ -1,9 +1,5 @@
 from simple_diarizer.diarizer import Diarizer
 
-import librosa
-
-import soundfile as sf
-
 
 class DiarizationPipeline:
     def __init__(self, ):
@@ -16,10 +12,7 @@ class DiarizationPipeline:
         )
 
     def __call__(self, wav_file):
-        y_16k, sr = librosa.load(wav_file, sr=16000)
-        sf.write("converted.wav", y_16k, 16000, 'PCM_24')
-
-        segments = self.diar.diarize("converted.wav",
+        segments = self.diar.diarize(wav_file,
                                      num_speakers=None,
                                      threshold=9e-1, )