import logging
|
import numpy as np
|
import soundfile
|
import kaldiio
|
from funasr.utils.job_runner import MultiProcessRunnerV3
|
from funasr.utils.misc import load_scp_as_list, load_scp_as_dict
|
import os
|
import argparse
|
from collections import OrderedDict
|
import random
|
|
|
class MyRunner(MultiProcessRunnerV3):
|
|
def prepare(self, parser: argparse.ArgumentParser):
|
parser.add_argument("--label_scp", type=str, required=True)
|
parser.add_argument("--wav_scp", type=str, required=True)
|
parser.add_argument("--utt2spk", type=str, required=True)
|
parser.add_argument("--spk2meeting", type=str, required=True)
|
parser.add_argument("--utt2xvec", type=str, required=True)
|
parser.add_argument("--out_dir", type=str, required=True)
|
parser.add_argument("--chunk_size", type=int, default=16)
|
parser.add_argument("--chunk_shift", type=int, default=4)
|
parser.add_argument("--frame_shift", type=float, default=0.01)
|
args = parser.parse_args()
|
|
if not os.path.exists(args.out_dir):
|
os.makedirs(args.out_dir)
|
|
label_list = load_scp_as_list(args.label_scp)
|
wav_scp = load_scp_as_dict(args.wav_scp)
|
utt2spk = load_scp_as_dict(args.utt2spk)
|
utt2xvec = load_scp_as_dict(args.utt2xvec)
|
spk2meeting = load_scp_as_dict(args.spk2meeting)
|
|
meeting2spks = OrderedDict()
|
for spk, meeting in spk2meeting.items():
|
if meeting not in meeting2spks:
|
meeting2spks[meeting] = []
|
meeting2spks[meeting].append(spk)
|
|
spk2utts = OrderedDict()
|
for utt, spk in utt2spk.items():
|
if spk not in spk2utts:
|
spk2utts[spk] = []
|
spk2utts[spk].append(utt)
|
|
return label_list, (wav_scp, utt2xvec, spk2utts, meeting2spks), args
|
|
def post(self, results_list, args):
|
pass
|
|
|
def process(task_args):
|
task_idx, task_list, (wav_scp, utt2xvec, spk2utts, meeting2spks), args = task_args
|
out_path = os.path.join(args.out_dir, "wav_mix.{}".format(task_idx+1))
|
wav_mix_writer = kaldiio.WriteHelper('ark,scp:{}.ark,{}.scp'.format(out_path, out_path))
|
|
out_path = os.path.join(args.out_dir, "wav_sep.{}".format(task_idx + 1))
|
wav_sep_writer = kaldiio.WriteHelper('ark,scp:{}.ark,{}.scp'.format(out_path, out_path))
|
|
out_path = os.path.join(args.out_dir, "label.{}".format(task_idx + 1))
|
label_writer = kaldiio.WriteHelper('ark,scp:{}.ark,{}.scp'.format(out_path, out_path))
|
|
idx = 0
|
for _, label_path in task_list:
|
rand_shift = random.randint(0, int(args.chunk_shift / args.frame_shift))
|
whole_label = kaldiio.load_mat(label_path)
|
whole_label = whole_label[rand_shift:]
|
num_chunk = (whole_label.shape[0] - args.chunk_size) // args.chunk_shift
|
for i in range(num_chunk):
|
utt_id = "part{}_utt{:10d}".format(task_idx + 1, idx + 1)
|
|
|
wav_mix_writer.close()
|
wav_sep_writer.close()
|
label_writer.close()
|
return None
|
|
|
if __name__ == '__main__':
|
my_runner = MyRunner(process)
|
my_runner.run()
|
