import argparse
|
import json
|
import os
|
|
import numpy as np
|
import torchaudio
|
import torchaudio.compliance.kaldi as kaldi
|
import yaml
|
from funasr.models.frontend.default import DefaultFrontend
|
import torch
|
|
def get_parser():
|
parser = argparse.ArgumentParser(
|
description="computer global cmvn",
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
)
|
parser.add_argument(
|
"--dim",
|
default=80,
|
type=int,
|
help="feature dimension",
|
)
|
parser.add_argument(
|
"--wav_path",
|
default=False,
|
required=True,
|
type=str,
|
help="the path of wav scps",
|
)
|
parser.add_argument(
|
"--config_file",
|
type=str,
|
help="the config file for computing cmvn",
|
)
|
parser.add_argument(
|
"--idx",
|
default=1,
|
required=True,
|
type=int,
|
help="index",
|
)
|
return parser
|
|
|
def compute_fbank(wav_file,
|
num_mel_bins=80,
|
frame_length=25,
|
frame_shift=10,
|
dither=0.0,
|
resample_rate=16000,
|
speed=1.0,
|
window_type="hamming"):
|
waveform, sample_rate = torchaudio.load(wav_file)
|
if resample_rate != sample_rate:
|
waveform = torchaudio.transforms.Resample(orig_freq=sample_rate,
|
new_freq=resample_rate)(waveform)
|
if speed != 1.0:
|
waveform, _ = torchaudio.sox_effects.apply_effects_tensor(
|
waveform, resample_rate,
|
[['speed', str(speed)], ['rate', str(resample_rate)]]
|
)
|
|
waveform = waveform * (1 << 15)
|
mat = kaldi.fbank(waveform,
|
num_mel_bins=num_mel_bins,
|
frame_length=frame_length,
|
frame_shift=frame_shift,
|
dither=dither,
|
energy_floor=0.0,
|
window_type=window_type,
|
sample_frequency=resample_rate)
|
|
return mat.numpy()
|
|
|
def main():
|
parser = get_parser()
|
args = parser.parse_args()
|
|
wav_scp_file = os.path.join(args.wav_path, "wav.{}.scp".format(args.idx))
|
cmvn_file = os.path.join(args.wav_path, "cmvn.{}.json".format(args.idx))
|
|
mean_stats = np.zeros(args.dim)
|
var_stats = np.zeros(args.dim)
|
total_frames = 0
|
|
# with ReadHelper('ark:{}'.format(ark_file)) as ark_reader:
|
# for key, mat in ark_reader:
|
# mean_stats += np.sum(mat, axis=0)
|
# var_stats += np.sum(np.square(mat), axis=0)
|
# total_frames += mat.shape[0]
|
|
with open(args.config_file) as f:
|
configs = yaml.safe_load(f)
|
frontend_configs = configs.get("frontend_conf", {})
|
num_mel_bins = frontend_configs.get("n_mels", 80)
|
frame_length = frontend_configs.get("frame_length", 25)
|
frame_shift = frontend_configs.get("frame_shift", 10)
|
window_type = frontend_configs.get("window", "hamming")
|
resample_rate = frontend_configs.get("fs", 16000)
|
n_fft = frontend_configs.get("n_fft", "400")
|
use_channel = frontend_configs.get("use_channel", None)
|
assert num_mel_bins == args.dim
|
frontend = DefaultFrontend(
|
fs=resample_rate,
|
n_fft=n_fft,
|
win_length=frame_length * 16,
|
hop_length=frame_shift * 16,
|
window=window_type,
|
n_mels=num_mel_bins,
|
use_channel=use_channel,
|
)
|
with open(wav_scp_file) as f:
|
lines = f.readlines()
|
for line in lines:
|
_, wav_file = line.strip().split()
|
wavform, _ = torchaudio.load(wav_file)
|
fbank, _ = frontend(wavform.transpose(0, 1).unsqueeze(0), torch.tensor([wavform.shape[1]]))
|
fbank = fbank.squeeze(0).numpy()
|
mean_stats += np.sum(fbank, axis=0)
|
var_stats += np.sum(np.square(fbank), axis=0)
|
total_frames += fbank.shape[0]
|
|
cmvn_info = {
|
'mean_stats': list(mean_stats.tolist()),
|
'var_stats': list(var_stats.tolist()),
|
'total_frames': total_frames
|
}
|
with open(cmvn_file, 'w') as fout:
|
fout.write(json.dumps(cmvn_info))
|
|
|
if __name__ == '__main__':
|
main()
|
