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()