class_docstr: SeparateSpeech class
  
  Examples:
      >>> import soundfile
      >>> separate_speech = SeparateSpeech("enh_config.yml", "enh.pth")
      >>> audio, rate = soundfile.read("speech.wav")
      >>> separate_speech(audio)
      [separated_audio1, separated_audio2, ...]
class_docstr_tok: ["SeparateSpeech","class","Examples",":",">",">",">","import","soundfile",">",">",">","separate_speech","=","SeparateSpeech","(","``","enh_config.yml","''",",","``","enh.pth","''",")",">",">",">","audio",",","rate","=","soundfile.read","(","``","speech.wav","''",")",">",">",">","separate_speech","(","audio",")","[","separated_audio1",",","separated_audio2",",","...","]"]
docstr: Inference
  
  Args:
      speech_mix: Input speech data (Batch, Nsamples [, Channels])
      fs: sample rate
      enroll_ref1: enrollment for speaker 1
      enroll_ref2: enrollment for speaker 2
      ...
  Returns:
      [separated_audio1, separated_audio2, ...]
docstr_tok: ["Inference","Args",":","speech_mix",":","Input","speech","data","(","Batch",",","Nsamples","[",",","Channels","]",")","fs",":","sample","rate","enroll_ref1",":","enrollment","for","speaker","1","enroll_ref2",":","enrollment","for","speaker","2","...","Returns",":","[","separated_audio1",",","separated_audio2",",","...","]"]
code: def __call__(
      self,
      speech_mix: Union[torch.Tensor, np.ndarray],
      fs: int = 8000,
      **kwargs,
  ) -> List[torch.Tensor]:
      """Inference
  
      Args:
          speech_mix: Input speech data (Batch, Nsamples [, Channels])
          fs: sample rate
          enroll_ref1: enrollment for speaker 1
          enroll_ref2: enrollment for speaker 2
          ...
      Returns:
          [separated_audio1, separated_audio2, ...]
  
      """
      assert check_argument_types()
  
      enroll_ref = [
          # (Batch, samples_aux)
          torch.as_tensor(kwargs["enroll_ref{}".format(spk + 1)])
          for spk in range(self.num_spk)
          if "enroll_ref{}".format(spk + 1) in kwargs
      ]
  
      # Input as audio signal
      if isinstance(speech_mix, np.ndarray):
          speech_mix = torch.as_tensor(speech_mix)
  
      assert speech_mix.dim() > 1, speech_mix.size()
      batch_size = speech_mix.size(0)
      speech_mix = speech_mix.to(getattr(torch, self.dtype))
      # lengths: (B,)
      lengths = speech_mix.new_full(
          [batch_size], dtype=torch.long, fill_value=speech_mix.size(1)
      )
      aux_lengths = [
          aux.new_full(
              [batch_size], dtype=torch.long, fill_value=aux.size(1)
          )
          for aux in enroll_ref
      ]
  
      # a. To device
      speech_mix = to_device(speech_mix, device=self.device)
      enroll_ref = to_device(enroll_ref, device=self.device)
      if self.enh_model.share_encoder:
          feats_aux, flens_aux = zip(
              *[
                  self.enh_model.encoder(
                      enroll_ref[spk], aux_lengths[spk]
                  )
                  for spk in range(len(enroll_ref))
              ]
          )
      else:
          feats_aux = enroll_ref
          flens_aux = aux_lengths
  
      if self.segmenting and lengths[0] > self.segment_size * fs:
          # Segment-wise speech enhancement/separation
          overlap_length = int(
              np.round(fs * (self.segment_size - self.hop_size))
          )
          num_segments = int(
              np.ceil(
                  (speech_mix.size(1) - overlap_length)
                  / (self.hop_size * fs)
              )
          )
          t = T = int(self.segment_size * fs)
          pad_shape = speech_mix[:, :T].shape
          enh_waves = []
          range_ = trange if self.show_progressbar else range
          for i in range_(num_segments):
              st = int(i * self.hop_size * fs)
              en = st + T
              if en >= lengths[0]:
                  # en - st < T (last segment)
                  en = lengths[0]
                  speech_seg = speech_mix.new_zeros(pad_shape)
                  t = en - st
                  speech_seg[:, :t] = speech_mix[:, st:en]
              else:
                  t = T
                  speech_seg = speech_mix[:, st:en]  # B x T [x C]
  
              lengths_seg = speech_mix.new_full(
                  [batch_size], dtype=torch.long, fill_value=T
              )
              # b. Enhancement/Separation Forward
              feats, f_lens = self.enh_model.encoder(
                  speech_seg, lengths_seg
              )
              feature_pre, _, others = zip(
                  *[
                      self.enh_model.extractor(
                          feats,
                          f_lens,
                          feats_aux[spk],
                          flens_aux[spk],
                          suffix_tag=f"_spk{spk + 1}",
                      )
                      for spk in range(len(enroll_ref))
                  ]
              )
              processed_wav = [
                  self.enh_model.decoder(f, lengths_seg)[0]
                  for f in feature_pre
              ]
              if speech_seg.dim() > 2:
                  # multi-channel speech
                  speech_seg_ = speech_seg[:, self.ref_channel]
              else:
                  speech_seg_ = speech_seg
  
              if self.normalize_segment_scale:
                  # normalize the scale to match the input mixture scale
                  mix_energy = torch.sqrt(
                      torch.mean(
                          speech_seg_[:, :t].pow(2), dim=1, keepdim=True
                      )
                  )
                  enh_energy = torch.sqrt(
                      torch.mean(
                          sum(processed_wav)[:, :t].pow(2),
                          dim=1,
                          keepdim=True,
                      )
                  )
                  processed_wav = [
                      w * (mix_energy / enh_energy)
                      for w in processed_wav
                  ]
              # List[torch.Tensor(num_spk, B, T)]
              enh_waves.append(torch.stack(processed_wav, dim=0))
  
          # c. Stitch the enhanced segments together
          waves = enh_waves[0]
          for i in range(1, num_segments):
              # permutation between separated streams in last and current segments
              perm = self.cal_permumation(
                  waves[:, :, -overlap_length:],
                  enh_waves[i][:, :, :overlap_length],
                  criterion="si_snr",
              )
              # repermute separated streams in current segment
              for batch in range(batch_size):
                  enh_waves[i][:, batch] = enh_waves[i][
                      perm[batch], batch
                  ]
  
              if i == num_segments - 1:
                  enh_waves[i][:, :, t:] = 0
                  enh_waves_res_i = enh_waves[i][:, :, overlap_length:t]
              else:
                  enh_waves_res_i = enh_waves[i][:, :, overlap_length:]
  
              # overlap-and-add (average over the overlapped part)
              waves[:, :, -overlap_length:] = (
                  waves[:, :, -overlap_length:]
                  + enh_waves[i][:, :, :overlap_length]
              ) / 2
              # concatenate the residual parts of the later segment
              waves = torch.cat([waves, enh_waves_res_i], dim=2)
          # ensure the stitched length is same as input
          assert waves.size(2) == speech_mix.size(1), (
              waves.shape,
              speech_mix.shape,
          )
          waves = torch.unbind(waves, dim=0)
      else:
          # b. Enhancement/Separation Forward
          feats, f_lens = self.enh_model.encoder(speech_mix, lengths)
  
          feature_pre, _, others = zip(
              *[
                  self.enh_model.extractor(
                      feats,
                      f_lens,
                      feats_aux[spk],
                      flens_aux[spk],
                      suffix_tag=f"_spk{spk + 1}",
                  )
                  for spk in range(len(enroll_ref))
              ]
          )
          others = {k: v for dic in others for k, v in dic.items()}
  
          waves = [
              self.enh_model.decoder(f, lengths)[0] for f in feature_pre
          ]
  
      assert len(waves[0]) == batch_size, (len(waves[0]), batch_size)
      if self.normalize_output_wav:
          waves = [
              (w / abs(w).max(dim=1, keepdim=True)[0] * 0.9)
              .cpu()
              .numpy()
              for w in waves
          ]  # list[(batch, sample)]
      else:
          waves = [w.cpu().numpy() for w in waves]
  
      return waves
  
imports: import argparse
  import logging
  import sys
  from itertools import chain
  from pathlib import Path
  from typing import Any, List, Optional, Sequence, Tuple, Union
  import humanfriendly
  import numpy
  import torch
  import yaml
  from tqdm import trange
  from typeguard import check_argument_types
  from espnet2.enh.loss.criterions.tf_domain import FrequencyDomainMSE
  from espnet2.enh.loss.criterions.time_domain import SISNRLoss
  from espnet2.enh.loss.wrappers.pit_solver import PITSolver
  from espnet2.fileio.sound_scp import SoundScpWriter
  from espnet2.tasks.enh_tse import TargetSpeakerExtractionTask
  from espnet2.torch_utils.device_funcs import to_device
  from espnet2.torch_utils.set_all_random_seed import set_all_random_seed
  from espnet2.train.abs_espnet_model import AbsESPnetModel
  from espnet2.utils import config_argparse
  from espnet2.utils.types import str2bool, str2triple_str, str_or_none
  from espnet.utils.cli_utils import get_commandline_args
label_desc: Use image node_id 2 for calling the SeparateSpeech obj's underlying member method code with example usage: obj.__call__(speech_mix, fs) and returns: waves