wehub-resource-sync
120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
#!/usr/bin/env python3 -u
|
|
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import argparse
|
|
import os
|
|
import os.path as osp
|
|
import tqdm
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from shutil import copyfile
|
|
|
|
from npy_append_array import NpyAppendArray
|
|
|
|
import fairseq
|
|
import soundfile as sf
|
|
|
|
|
|
def get_parser():
|
|
parser = argparse.ArgumentParser(
|
|
description="compute kmeans codebook from kaldi-computed feats"
|
|
)
|
|
# fmt: off
|
|
parser.add_argument('data', help='location of tsv files')
|
|
parser.add_argument('--split', help='which split to read', required=True)
|
|
parser.add_argument('--save-dir', help='where to save the output', required=True)
|
|
parser.add_argument('--checkpoint', type=str, help='checkpoint for wav2vec ctc model', required=True)
|
|
parser.add_argument('--layer', type=int, default=14, help='which layer to use')
|
|
# fmt: on
|
|
|
|
return parser
|
|
|
|
|
|
class Wav2VecFeatureReader(object):
|
|
def __init__(self, cp_file, layer):
|
|
model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task(
|
|
[cp_file]
|
|
)
|
|
model = model[0]
|
|
model.eval()
|
|
model.cuda()
|
|
self.model = model
|
|
self.task = task
|
|
self.layer = layer
|
|
|
|
def read_audio(self, fname):
|
|
"""Load an audio file and return PCM along with the sample rate"""
|
|
wav, sr = sf.read(fname)
|
|
assert sr == 16e3
|
|
|
|
return wav
|
|
|
|
def get_feats(self, loc):
|
|
x = self.read_audio(loc)
|
|
with torch.no_grad():
|
|
source = torch.from_numpy(x).float().cuda()
|
|
if self.task.cfg.normalize:
|
|
assert source.dim() == 1, source.dim()
|
|
with torch.no_grad():
|
|
source = F.layer_norm(source, source.shape)
|
|
source = source.view(1, -1)
|
|
|
|
m_res = self.model(source=source, mask=False, features_only=True, layer=self.layer)
|
|
return m_res["x"].squeeze(0).cpu()
|
|
|
|
|
|
def get_iterator(args):
|
|
with open(osp.join(args.data, args.split) + ".tsv", "r") as fp:
|
|
lines = fp.read().split("\n")
|
|
root = lines.pop(0).strip()
|
|
files = [osp.join(root, line.split("\t")[0]) for line in lines if len(line) > 0]
|
|
|
|
num = len(files)
|
|
reader = Wav2VecFeatureReader(args.checkpoint, args.layer)
|
|
|
|
def iterate():
|
|
for fname in files:
|
|
w2v_feats = reader.get_feats(fname)
|
|
yield w2v_feats
|
|
|
|
return iterate, num
|
|
|
|
|
|
def main():
|
|
parser = get_parser()
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.save_dir, exist_ok=True)
|
|
|
|
def create_files(dest):
|
|
copyfile(osp.join(args.data, args.split) + ".tsv", dest + ".tsv")
|
|
if osp.exists(osp.join(args.data, args.split) + ".wrd"):
|
|
copyfile(osp.join(args.data, args.split) + ".wrd", dest + ".wrd")
|
|
if osp.exists(osp.join(args.data, args.split) + ".phn"):
|
|
copyfile(osp.join(args.data, args.split) + ".phn", dest + ".phn")
|
|
|
|
if osp.exists(dest + ".npy"):
|
|
os.remove(dest + ".npy")
|
|
npaa = NpyAppendArray(dest + ".npy")
|
|
return npaa
|
|
|
|
save_path = osp.join(args.save_dir, args.split)
|
|
npaa = create_files(save_path)
|
|
|
|
generator, num = get_iterator(args)
|
|
iterator = generator()
|
|
|
|
with open(save_path + ".lengths", "w") as l_f:
|
|
for w2v_feats in tqdm.tqdm(iterator, total=num):
|
|
print(len(w2v_feats), file=l_f)
|
|
|
|
if len(w2v_feats) > 0:
|
|
npaa.append(w2v_feats.numpy())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|