B fgÃã@sddlmZGdd„dƒZdS)é)Únetwork_wrapperc@s*eZdZdZdd„Zd dd„Zdd „ZdS) Ú ClearVoicez˜ The main class inferface to the end users for performing speech processing this class provides the desired model to perform the given task cCs<tƒ|_g|_x(|D] }| ||¡}|j|g7_qWdS)aÈ Load the desired models for the specified task. Perform all the given models and return all results. Parameters: ---------- task: str the task matching any of the provided tasks: 'speech_enhancement' 'speech_separation' 'target_speaker_extraction' model_names: str or list of str the model names matching any of the provided models: 'FRCRN_SE_16K' 'MossFormer2_SE_48K' 'MossFormerGAN_SE_16K' 'MossFormer2_SS_16K' 'AV_MossFormer2_TSE_16K' Returns: -------- A ModelsList object, that can be run to get the desired results N)rÚmodels)ÚselfÚtaskÚ model_namesZ model_nameÚmodel©r úV/mnt/nas/mit_sg/shengkui.zhao/speech_codec/clear_speech_local/clearvoice/clearvoice.pyÚ__init__s   zClearVoice.__init__FNcCsRi}x*|jD] }| |||¡}|s |||j<q W|sNt|ƒdkrJ||jS|SdS)Né)rZprocessÚnameÚlen)rÚ input_pathÚ online_writeÚ output_pathÚresultsrÚresultr r r Ú__call__#s   zClearVoice.__call__cCsŠd}d}t|jƒdkrd}xL|jD]B}t|tƒrb|j|krRt||jƒdkrbd}nt|ƒdkrbd}Pq"Wx|jD]}| |||¡qpWdS)NFr T)rrÚ isinstanceÚdictr Úwrite)rrrZ add_subdirZuse_keyrr r r r0s     zClearVoice.write)FN)Ú__name__Ú __module__Ú __qualname__Ú__doc__r rrr r r r rs rN)rrr r r r Ús