def __init__(
self,
enh_train_config: Union[Path, str],
enh_model_file: Union[Path, str] = None,
segment_size: Optional[float] = None,
hop_size: Optional[float] = None,
normalize_segment_scale: bool = False,
show_progressbar: bool = False,
ref_channel: Optional[int] = None,
normalize_output_wav: bool = False,
device: str = "cpu",
dtype: str = "float32",
):
assert check_argument_types()
# 1. Build Enh model
enh_model, enh_train_args = EnhancementTask.build_model_from_file(
enh_train_config, enh_model_file, device
)
enh_model.to(dtype=getattr(torch, dtype)).eval()
self.device = device
self.dtype = dtype
self.enh_train_args = enh_train_args
self.enh_model = enh_model
# only used when processing long speech, i.e.
# segment_size is not None and hop_size is not None
self.segment_size = segment_size
self.hop_size = hop_size
self.normalize_segment_scale = normalize_segment_scale
self.normalize_output_wav = normalize_output_wav
self.show_progressbar = show_progressbar
self.num_spk = enh_model.num_spk
task = "enhancement" if self.num_spk == 1 else "separation"
# reference channel for processing multi-channel speech
if ref_channel is not None:
logging.info(
"Overwrite enh_model.separator.ref_channel with {}".format(ref_channel)
)
enh_model.separator.ref_channel = ref_channel
self.ref_channel = ref_channel
else:
self.ref_channel = enh_model.ref_channel
self.segmenting = segment_size is not None and hop_size is not None
if self.segmenting:
logging.info("Perform segment-wise speech %s" % task)
logging.info(
"Segment length = {} sec, hop length = {} sec".format(
segment_size, hop_size
)
)
else:
logging.info("Perform direct speech %s on the input" % task)
def inference(
output_dir: str,
batch_size: int,
dtype: str,
fs: int,
ngpu: int,
seed: int,
num_workers: int,
log_level: Union[int, str],
data_path_and_name_and_type: Sequence[Tuple[str, str, str]],
key_file: Optional[str],
enh_train_config: str,
enh_model_file: str,
allow_variable_data_keys: bool,
normalize_output_wav: bool,
):
assert check_argument_types()
if batch_size > 1:
raise NotImplementedError("batch decoding is not implemented")
if ngpu > 1:
raise NotImplementedError("only single GPU decoding is supported")
logging.basicConfig(
level=log_level,
format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
)
if ngpu >= 1:
device = "cuda"
else:
device = "cpu"
# 1. Set random-seed
set_all_random_seed(seed)
# 2. Build Enh model
enh_model, enh_train_args = EnhancementTask.build_model_from_file(
enh_train_config, enh_model_file, device)
enh_model.eval()
num_spk = enh_model.num_spk
# 3. Build data-iterator
loader = EnhancementTask.build_streaming_iterator(
data_path_and_name_and_type,
dtype=dtype,
batch_size=batch_size,
key_file=key_file,
num_workers=num_workers,
preprocess_fn=EnhancementTask.build_preprocess_fn(
enh_train_args, False),
collate_fn=EnhancementTask.build_collate_fn(enh_train_args),
allow_variable_data_keys=allow_variable_data_keys,
inference=True,
)
writers = []
for i in range(num_spk):
writers.append(
SoundScpWriter(f"{output_dir}/wavs/{i + 1}",
f"{output_dir}/spk{i + 1}.scp"))
for keys, batch in loader:
assert isinstance(batch, dict), type(batch)
assert all(isinstance(s, str) for s in keys), keys
_bs = len(next(iter(batch.values())))
assert len(keys) == _bs, f"{len(keys)} != {_bs}"
with torch.no_grad():
# a. To device
batch = to_device(batch, device)
# b. Forward Enhancement Frontend
waves, _, _ = enh_model.enh_model.forward_rawwav(
batch["speech_mix"], batch["speech_mix_lengths"])
assert len(waves[0]) == batch_size, len(waves[0])
# FIXME(Chenda): will be incorrect when
# batch size is not 1 or multi-channel case
if normalize_output_wav:
waves = [
(w / abs(w).max(dim=1, keepdim=True)[0] * 0.9).T.cpu().numpy()
for w in waves
] # list[(sample,batch)]
else:
waves = [w.T.cpu().numpy() for w in waves]
for (i, w) in enumerate(waves):
writers[i][keys[0]] = fs, w
for writer in writers:
writer.close()