import argparse
|
import os
|
|
import numpy as np
|
|
import funasr.modules.eend_ola.utils.feature as feature
|
import funasr.modules.eend_ola.utils.kaldi_data as kaldi_data
|
|
|
def _count_frames(data_len, size, step):
|
return int((data_len - size + step) / step)
|
|
|
def _gen_frame_indices(
|
data_length, size=2000, step=2000,
|
use_last_samples=False,
|
label_delay=0,
|
subsampling=1):
|
i = -1
|
for i in range(_count_frames(data_length, size, step)):
|
yield i * step, i * step + size
|
if use_last_samples and i * step + size < data_length:
|
if data_length - (i + 1) * step - subsampling * label_delay > 0:
|
yield (i + 1) * step, data_length
|
|
|
class KaldiDiarizationDataset():
|
def __init__(
|
self,
|
data_dir,
|
chunk_size=2000,
|
context_size=0,
|
frame_size=1024,
|
frame_shift=256,
|
subsampling=1,
|
rate=16000,
|
input_transform=None,
|
use_last_samples=False,
|
label_delay=0,
|
n_speakers=None,
|
):
|
self.data_dir = data_dir
|
self.chunk_size = chunk_size
|
self.context_size = context_size
|
self.frame_size = frame_size
|
self.frame_shift = frame_shift
|
self.subsampling = subsampling
|
self.input_transform = input_transform
|
self.n_speakers = n_speakers
|
self.chunk_indices = []
|
self.label_delay = label_delay
|
|
self.data = kaldi_data.KaldiData(self.data_dir)
|
|
# make chunk indices: filepath, start_frame, end_frame
|
for rec, path in self.data.wavs.items():
|
data_len = int(self.data.reco2dur[rec] * rate / frame_shift)
|
data_len = int(data_len / self.subsampling)
|
for st, ed in _gen_frame_indices(
|
data_len, chunk_size, chunk_size, use_last_samples,
|
label_delay=self.label_delay,
|
subsampling=self.subsampling):
|
self.chunk_indices.append(
|
(rec, path, st * self.subsampling, ed * self.subsampling))
|
print(len(self.chunk_indices), " chunks")
|
|
|
def convert(args):
|
f = open(out_wav_file, 'w')
|
dataset = KaldiDiarizationDataset(
|
data_dir=args.data_dir,
|
chunk_size=args.num_frames,
|
context_size=args.context_size,
|
input_transform=args.input_transform,
|
frame_size=args.frame_size,
|
frame_shift=args.frame_shift,
|
subsampling=args.subsampling,
|
rate=8000,
|
use_last_samples=True,
|
)
|
length = len(dataset.chunk_indices)
|
for idx, (rec, path, st, ed) in enumerate(dataset.chunk_indices):
|
Y, T = feature.get_labeledSTFT(
|
dataset.data,
|
rec,
|
st,
|
ed,
|
dataset.frame_size,
|
dataset.frame_shift,
|
dataset.n_speakers)
|
Y = feature.transform(Y, dataset.input_transform)
|
Y_spliced = feature.splice(Y, dataset.context_size)
|
Y_ss, T_ss = feature.subsample(Y_spliced, T, dataset.subsampling)
|
st = '{:0>7d}'.format(st)
|
ed = '{:0>7d}'.format(ed)
|
suffix = '_' + st + '_' + ed
|
|
parts = os.readlink('/'.join(path.split('/')[:-1])).split('/')
|
# print('parts: ', parts)
|
parts = parts[:4] + ['numpy_data'] + parts[4:]
|
cur_path = '/'.join(parts)
|
# print('cur path: ', cur_path)
|
out_path = os.path.join(cur_path, path.split('/')[-1].split('.')[0] + suffix + '.npz')
|
# print(out_path)
|
# print(cur_path)
|
if not os.path.exists(cur_path):
|
os.makedirs(cur_path)
|
np.savez(out_path, Y=Y_ss, T=T_ss)
|
if idx == length - 1:
|
f.write(rec + suffix + ' ' + out_path)
|
else:
|
f.write(rec + suffix + ' ' + out_path + '\n')
|
|
|
if __name__ == '__main__':
|
parser = argparse.ArgumentParser()
|
parser.add_argument("data_dir")
|
parser.add_argument("num_frames")
|
parser.add_argument("context_size")
|
parser.add_argument("frame_size")
|
parser.add_argument("frame_shift")
|
parser.add_argument("subsampling")
|
|
|
|
args = parser.parse_args()
|
convert(args)
|
