| import argparse | |
| from utils.denoise_pipeline import denoise | |
| from utils.diarization_pipeline import DiarizationPipeline | |
| class CleaningPipeline: | |
| def __init__(self, device): | |
| """ | |
| Cleaning audio pipeline. Contains: | |
| - denoising | |
| - diarization | |
| """ | |
| self.device = device | |
| self.denoiser = denoise | |
| self.diarization = DiarizationPipeline(device) | |
| def __call__(self, input_audio_path: str): | |
| denoised_audio_path = self.denoiser(input_audio_path, self.device) | |
| result_diarization = self.diarization(denoised_audio_path) | |
| if result_diarization != {}: | |
| output_diar_audio_paths = result_diarization['output_diar_audio_paths'] | |
| count_speakers = result_diarization['count_speakers'] | |
| return [denoised_audio_path] + output_diar_audio_paths + [None] * (19 - count_speakers) | |
| else: | |
| return [] | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--audio-path', default='dialog.mp3', help='Path to audio') | |
| parser.add_argument('--device', default='cpu', help='Path to audio') | |
| parser.add_argument('--out-folder-path', default='out', help='Path to result folder') | |
| opt = parser.parse_args() | |
| cleaning_pipeline = CleaningPipeline('cuda:0') | |
| cleaning_pipeline(input_audio_path=opt.audio_path) | |