class_docstr: DiarizeSpeech class
  
  Examples:
      >>> import soundfile
      >>> diarization = DiarizeSpeech("diar_config.yaml", "diar.pth")
      >>> audio, rate = soundfile.read("speech.wav")
      >>> diarization(audio)
      [(spk_id, start, end), (spk_id2, start2, end2)]
class_docstr_tok: ["DiarizeSpeech","class","Examples",":",">",">",">","import","soundfile",">",">",">","diarization","=","DiarizeSpeech","(","``","diar_config.yaml","''",",","``","diar.pth","''",")",">",">",">","audio",",","rate","=","soundfile.read","(","``","speech.wav","''",")",">",">",">","diarization","(","audio",")","[","(","spk_id",",","start",",","end",")",",","(","spk_id2",",","start2",",","end2",")","]"]
docstr: Inference
  
  Args:
      speech: Input speech data (Batch, Nsamples [, Channels])
      fs: sample rate
  Returns:
      [speaker_info1, speaker_info2, ...]
docstr_tok: ["Inference","Args",":","speech",":","Input","speech","data","(","Batch",",","Nsamples","[",",","Channels","]",")","fs",":","sample","rate","Returns",":","[","speaker_info1",",","speaker_info2",",","...","]"]
code: def __call__(
      self, speech: Union[torch.Tensor, np.ndarray], fs: int = 8000
  ) -> List[torch.Tensor]:
      """Inference
  
      Args:
          speech: Input speech data (Batch, Nsamples [, Channels])
          fs: sample rate
      Returns:
          [speaker_info1, speaker_info2, ...]
  
      """
      assert check_argument_types()
  
      # Input as audio signal
      if isinstance(speech, np.ndarray):
          speech = torch.as_tensor(speech)
  
      assert speech.dim() > 1, speech.size()
      batch_size = speech.size(0)
      speech = speech.to(getattr(torch, self.dtype))
      # lengths: (B,)
      lengths = speech.new_full(
          [batch_size], dtype=torch.long, fill_value=speech.size(1)
      )
  
      # a. To device
      speech = to_device(speech, device=self.device)
      lengths = to_device(lengths, device=self.device)
  
      if self.segmenting_diar and lengths[0] > self.segment_size * fs:
          # Segment-wise speaker diarization
          # Note that the segments are processed independently for now
          # i.e., no speaker tracing is performed
          num_segments = int(
              np.ceil(speech.size(1) / (self.segment_size * fs))
          )
          t = T = int(self.segment_size * fs)
          pad_shape = speech[:, :T].shape
          diarized_wavs = []
          range_ = trange if self.show_progressbar else range
          for i in range_(num_segments):
              st = int(i * self.segment_size * fs)
              en = st + T
              if en >= lengths[0]:
                  # en - st < T (last segment)
                  en = lengths[0]
                  speech_seg = speech.new_zeros(pad_shape)
                  t = en - st
                  speech_seg[:, :t] = speech[:, st:en]
              else:
                  t = T
                  speech_seg = speech[:, st:en]  # B x T [x C]
  
              lengths_seg = speech.new_full(
                  [batch_size], dtype=torch.long, fill_value=T
              )
              # b. Diarization Forward
              encoder_out, encoder_out_lens = self.encode(
                  speech_seg,
                  lengths_seg,
              )
              spk_prediction, _ = self.decode(
                  encoder_out, encoder_out_lens
              )
              # List[torch.Tensor(B, T, num_spks)]
              diarized_wavs.append(spk_prediction)
          # Determine maximum estimated number of speakers among the segments
          max_len = max([x.size(2) for x in diarized_wavs])
          # pad tensors in diarized_wavs with "float('-inf')" to have same size
          diarized_wavs = [
              torch.nn.functional.pad(
                  x, (0, max_len - x.size(2)), "constant", float("-inf")
              )
              for x in diarized_wavs
          ]
          spk_prediction = torch.cat(diarized_wavs, dim=1)
          waves = None
      else:
          # b. Diarization Forward
          encoder_out, encoder_out_lens = self.encode(speech, lengths)
          spk_prediction, num_spk = self.decode(
              encoder_out, encoder_out_lens
          )
          if self.enh_s2t_task:
              # Segment-wise speech separation
              # Note that this is done after diarization using the whole sequence
              if (
                  self.segmenting_enh_diar
                  and lengths[0] > self.segment_size * fs
              ):
                  overlap_length = int(
                      np.round(fs * (self.segment_size - self.hop_size))
                  )
                  num_segments = int(
                      np.ceil(
                          (speech.size(1) - overlap_length)
                          / (self.hop_size * fs)
                      )
                  )
                  t = T = int(self.segment_size * fs)
                  pad_shape = speech[:, :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.new_zeros(pad_shape)
                          t = en - st
                          speech_seg[:, :t] = speech[:, st:en]
                      else:
                          t = T
                          speech_seg = speech[:, st:en]  # B x T [x C]
  
                      lengths_seg = speech.new_full(
                          [batch_size], dtype=torch.long, fill_value=T
                      )
                      # Separation Forward
                      _, _, processed_wav = self.diar_model.encode_diar(
                          speech_seg, lengths_seg, num_spk
                      )
                      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.size(1), (
                      waves.shape,
                      speech.shape,
                  )
                  waves = torch.unbind(waves, dim=0)
              else:
                  # Separation Forward using the whole signal
                  _, _, waves = self.diar_model.encode_diar(
                      speech, lengths, num_spk
                  )
              # multiply diarization result and separation result
              # by calculating the correlation
              if self.multiply_diar_result:
                  (
                      spk_prediction,
                      interp_prediction,
                      _,
                  ) = self.permute_diar(waves, spk_prediction)
                  waves = [
                      waves[i] * interp_prediction[:, :, i]
                      for i in range(num_spk)
                  ]
              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]
          else:
              waves = None
  
      if self.num_spk is not None:
          assert spk_prediction.size(2) == self.num_spk, (
              spk_prediction.size(2),
              self.num_spk,
          )
      assert spk_prediction.size(0) == batch_size, (
          spk_prediction.size(0),
          batch_size,
      )
      spk_prediction = spk_prediction.cpu().numpy()
      spk_prediction = 1 / (1 + np.exp(-spk_prediction))
  
      return (
          waves,
          spk_prediction if self.enh_s2t_task else spk_prediction,
      )
  
imports: import argparse
  import logging
  import sys
  from itertools import permutations
  from pathlib import Path
  from typing import Any, List, Optional, Sequence, Tuple, Union
  import numpy
  import torch
  import torch.nn.functional
  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.npy_scp import NpyScpWriter
  from espnet2.fileio.sound_scp import SoundScpWriter
  from espnet2.tasks.diar import DiarizationTask
  from espnet2.tasks.enh_s2t import EnhS2TTask
  from espnet2.torch_utils.device_funcs import to_device
  from espnet2.torch_utils.set_all_random_seed import set_all_random_seed
  from espnet2.utils import config_argparse
  from espnet2.utils.types import humanfriendly_parse_size_or_none, int_or_none, 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 DiarizeSpeech obj's underlying member method code with example usage: obj.__call__(speech, fs) and returns: waves,  unknown