Skip to content

spk_feat_extractor

Author: Heli Qi Affiliation: NAIST Date: 2022.12

main(src_file, result_path=None, spk_emb_model='ecapa', batch_size=10, ncpu=8, ngpu=0)

Main function to extract speaker embeddings from audio waveforms.

Parameters:

Name Type Description Default
src_file str

Path to the idx2wav file containing the source waveforms for speaker embedding extraction.

 required
result_path str

Path to save the extracted speaker embeddings.

 None
spk_emb_model str

Speaker recognition model to use for extracting speaker embeddings.

 'ecapa'
ncpu int

Number of CPU processes for extracting speaker embeddings.

 8
ngpu int

Number of GPUs for extracting speaker embeddings.

 0
Source code in speechain/datasets/pyscripts/spk_feat_extractor.py 
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def main(
    src_file: str,
    result_path: str = None,
    spk_emb_model: str = "ecapa",
    batch_size: int = 10,
    ncpu: int = 8,
    ngpu: int = 0,
):
    """
    Main function to extract speaker embeddings from audio waveforms.

    Args:
        src_file (str):
            Path to the idx2wav file containing the source waveforms for speaker embedding extraction.
        result_path (str):
            Path to save the extracted speaker embeddings.
        spk_emb_model (str):
            Speaker recognition model to use for extracting speaker embeddings.
        ncpu (int):
            Number of CPU processes for extracting speaker embeddings.
        ngpu (int):
            Number of GPUs for extracting speaker embeddings.
    """
    spk_emb_model = spk_emb_model.lower()
    assert spk_emb_model in [
        "ecapa",
        "xvector",
    ], f"Your input spk_emb_model should be one of ['ecapa', 'xvector'], but got {spk_emb_model}!"

    # initialize the result path
    src_idx2wav_path = parse_path_args(src_file)
    if result_path is None:
        result_path = os.path.dirname(src_idx2wav_path)
    else:
        result_path = parse_path_args(result_path)

    # --- 1. Speaker Embedding Feature Extraction for Individual Utterances --- #
    idx2spk_feat_path = os.path.join(result_path, f"idx2{spk_emb_model}_spk_feat")
    spk2aver_spk_feat_path = os.path.join(
        result_path, f"spk2aver_{spk_emb_model}_spk_feat"
    )
    # skip if idx2spk_feat and spk2aver_spk_feat has already existed
    if not os.path.exists(idx2spk_feat_path) or not os.path.exists(
        spk2aver_spk_feat_path
    ):
        # read the source idx2wav file into a Dict, str -> Dict[str, str]
        src_idx2wav = load_idx2data_file(src_idx2wav_path)
        # read the idx2spk file in the directory of idx2wav into a Dict, str -> Dict[str, str]
        idx2spk = load_idx2data_file(
            os.path.join(os.path.dirname(src_idx2wav_path), "idx2spk")
        )

        # gather waveforms for each speaker
        spk2wav_dict = {}
        for idx, spk_id in idx2spk.items():
            if spk_id not in spk2wav_dict.keys():
                spk2wav_dict[spk_id] = {}
            spk2wav_dict[spk_id][idx] = src_idx2wav[idx]

        # initialize the arguments for the execution function
        extract_spk_feat_func = partial(
            extract_spk_feat,
            spk_emb_model=spk_emb_model,
            batch_size=batch_size,
            save_path=os.path.join(result_path, spk_emb_model),
        )

        device_list = (
            get_idle_gpu(ngpu, id_only=True) if ngpu > 0 else [-1 for _ in range(ncpu)]
        )
        n_proc = len(device_list) if ngpu > 0 else ncpu
        spk2wav_list = list(spk2wav_dict.items())
        func_args = [
            [dict(spk2wav_list[i::n_proc]), device_list[i]] for i in range(n_proc)
        ]

        # start the executing jobs
        with Pool(n_proc) as executor:
            extraction_results = executor.starmap(extract_spk_feat_func, func_args)

        # gather the results from all the processes
        idx2spk_feat, spk2aver_spk_feat = {}, {}
        for _idx2spk_feat, _spk2aver_spk_feat in extraction_results:
            idx2spk_feat.update(_idx2spk_feat)
            spk2aver_spk_feat.update(_spk2aver_spk_feat)

        # record the address of the .npy file of each vector
        np.savetxt(
            idx2spk_feat_path,
            sorted(idx2spk_feat.items(), key=lambda x: x[0]),
            fmt="%s",
        )
        # record the address of the .npy file of each vector
        np.savetxt(
            spk2aver_spk_feat_path,
            sorted(spk2aver_spk_feat.items(), key=lambda x: x[0]),
            fmt="%s",
        )