FunASR/egs/aishell2/transformer/utils/apply_cmvn.py
2023-05-12 17:25:54 +08:00

80 lines
1.8 KiB
Python
Executable File

from kaldiio import ReadHelper
from kaldiio import WriteHelper
import argparse
import json
import math
import numpy as np
def get_parser():
parser = argparse.ArgumentParser(
description="apply cmvn",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--ark-file",
"-a",
default=False,
required=True,
type=str,
help="fbank ark file",
)
parser.add_argument(
"--cmvn-file",
"-c",
default=False,
required=True,
type=str,
help="cmvn file",
)
parser.add_argument(
"--ark-index",
"-i",
default=1,
required=True,
type=int,
help="ark index",
)
parser.add_argument(
"--output-dir",
"-o",
default=False,
required=True,
type=str,
help="output dir",
)
return parser
def main():
parser = get_parser()
args = parser.parse_args()
ark_file = args.output_dir + "/feats." + str(args.ark_index) + ".ark"
scp_file = args.output_dir + "/feats." + str(args.ark_index) + ".scp"
ark_writer = WriteHelper('ark,scp:{},{}'.format(ark_file, scp_file))
with open(args.cmvn_file) as f:
cmvn_stats = json.load(f)
means = cmvn_stats['mean_stats']
vars = cmvn_stats['var_stats']
total_frames = cmvn_stats['total_frames']
for i in range(len(means)):
means[i] /= total_frames
vars[i] = vars[i] / total_frames - means[i] * means[i]
if vars[i] < 1.0e-20:
vars[i] = 1.0e-20
vars[i] = 1.0 / math.sqrt(vars[i])
with ReadHelper('ark:{}'.format(args.ark_file)) as ark_reader:
for key, mat in ark_reader:
mat = (mat - means) * vars
ark_writer(key, mat)
if __name__ == '__main__':
main()