SpeechSynthesisDataset returns speaker_ids

lhotse/dataset/speech_synthesis.py

@@ -1,5 +1,6 @@
 from typing import Callable, Dict, List, Sequence, Union
 
+import numpy as np
 import torch
 
 from lhotse import validate
@@ -23,6 +24,7 @@
             'audio_lens': (B, ) int tensor
             'features_lens': (B, ) int tensor
             'tokens_lens': (B, ) int tensor
+            'speakers': (B) long tensor (optional)
         }
     """
 
@@ -32,6 +34,7 @@
         cut_transforms: List[Callable[[CutSet], CutSet]] = None,
         feature_input_strategy: BatchIO = PrecomputedFeatures(),
         feature_transforms: Union[Sequence[Callable], Callable] = None,
+        speaker_id_mapping: Dict[str, int] = None,
         add_eos: bool = True,
         add_bos: bool = True,
     ) -> None:
@@ -41,6 +44,7 @@
         self.token_collater = TokenCollater(cuts, add_eos=add_eos, add_bos=add_bos)
         self.cut_transforms = ifnone(cut_transforms, [])
         self.feature_input_strategy = feature_input_strategy
+        self.speaker_id_mapping = speaker_id_mapping
 
         if feature_transforms is None:
             feature_transforms = []
@@ -65,15 +69,21 @@
             features = transform(features)
 
         tokens, tokens_lens = self.token_collater(cuts)
-
-        return {
+        batch = {
             "audio": audio,
             "features": features,
             "tokens": tokens,
             "audio_lens": audio_lens,
             "features_lens": features_lens,
             "tokens_lens": tokens_lens,
         }
+        if self.speaker_id_mapping is not None:
+            batch["speakers"] = torch.tensor(
+                [self.speaker_id_mapping[cut.supervisions[0].speaker] for cut in cuts],
+                dtype=torch.long,
+            )
+        else:
+            return batch
 
 
 def validate_for_tts(cuts: CutSet) -> None:
@@ -82,3 +92,10 @@
         assert (
             len(cut.supervisions) == 1
         ), "Only the Cuts with single supervision are supported."
+
+
+def get_sid_to_index_map(sid_list) -> Dict[str, np.ndarray]:
+    sid_to_onehot_map = {}
+    for index, sid in enumerate(sid_list):
+        sid_to_onehot_map[sid] = index
+    return sid_to_onehot_map