Some additional options for supervision-related methods (#1115)

* add transform attribute for MixedCut

* add mix_first option in normalize_loudness

* handle the case when mix is called on MixedCut with existing transforms

* add test for mixing with transformed MixedCut

* enhancements and bug fixes

* small changes in some cutset methods

* small fix in error message

* return word alignments from ami recipe

* add word alignments for ICSI

* remove unwanted whitespace

* fix IHM preparation

* remove words with zero or negative duration

* ensure word alignments respect segment boundary

* add save-to-wav option for icsi

* add test for mixing cut with recording

* style fix

* add data prep for voxpopuli

* small changes in recipes

* changes for max segment duration

* remove extra code

* made suggested changes

* apply change to multi custom merge func

* remove old code

* fix failing tests

* add tests for trim to alignments with max segment duration

* add tests for merge supervisions

Author: desh2608

lhotse/bin/modes/__init__.py

@@ -6,5 +6,6 @@
 from .manipulation import *
 from .recipes import *
 from .shar import *
+from .supervision import *
 from .validate import *
 from .workflows import *

lhotse/bin/modes/supervision.py

@@ -0,0 +1,73 @@
+import click
+from tqdm import tqdm
+
+from lhotse.bin.modes.cli_base import cli
+from lhotse.serialization import load_manifest_lazy_or_eager
+from lhotse.supervision import SupervisionSet
+from lhotse.utils import Pathlike
+
+
+@cli.group()
+def supervision():
+    """Commands related to manipulating supervision manifests."""
+    pass
+
+
+@supervision.command()
+@click.argument("in_supervision_manifest", type=click.Path(allow_dash=True))
+@click.argument("out_supervision_manifest", type=click.Path(allow_dash=True))
+@click.option(
+    "--ctm-file",
+    type=click.Path(exists=True, dir_okay=False),
+    help="CTM file containing alignments to add.",
+)
+@click.option(
+    "--alignment-type",
+    type=str,
+    default="word",
+    help="Type of alignment to add (default = `word`).",
+)
+@click.option(
+    "--match-channel/--no-match-channel",
+    default=False,
+    help="Whether to match channel between CTM and SupervisionSegment (default = False).",
+)
+@click.option(
+    "--verbose",
+    "-v",
+    is_flag=True,
+    default=False,
+    help="Whether to print verbose output.",
+)
+def with_alignment_from_ctm(
+    in_supervision_manifest: Pathlike,
+    out_supervision_manifest: Pathlike,
+    ctm_file: Pathlike,
+    alignment_type: str,
+    match_channel: bool,
+    verbose: bool,
+):
+    """
+    Add alignments from CTM file to the supervision set.
+
+    :param in_supervision_manifest: Path to input supervision manifest.
+    :param out_supervision_manifest: Path to output supervision manifest.
+    :param ctm_file: Path to CTM file.
+    :param alignment_type: Alignment type (optional, default = `word`).
+    :param match_channel: if True, also match channel between CTM and SupervisionSegment
+    :param verbose: Whether to print verbose output.
+    :return: A new SupervisionSet with AlignmentItem objects added to the segments.
+    """
+    supervisions = load_manifest_lazy_or_eager(in_supervision_manifest, SupervisionSet)
+    supervisions = supervisions.with_alignment_from_ctm(
+        ctm_file=ctm_file,
+        type=alignment_type,
+        match_channel=match_channel,
+        verbose=verbose,
+    )
+    with SupervisionSet.open_writer(out_supervision_manifest, overwrite=True) as writer:
+        supervisions = (
+            tqdm(supervisions, desc="Writing supervisions") if verbose else supervisions
+        )
+        for s in supervisions:
+            writer.write(s)

lhotse/cut/base.py

@@ -510,14 +510,16 @@ def trim_to_alignments(
         self,
         type: str,
         max_pause: Optional[Seconds] = None,
+        max_segment_duration: Optional[Seconds] = None,
         delimiter: str = " ",
         keep_all_channels: bool = False,
     ) -> "CutSet":  # noqa: F821
         """
         Splits the current :class:`.Cut` into its constituent alignment items (:class:`.AlignmentItem`).
         These cuts have identical start times and durations as the alignment item. Additionally,
         the `max_pause` option can be used to merge alignment items that are separated by a pause
-        shorter than `max_pause`.
+        shorter than `max_pause`. If `max_segment_duration` is specified, we will keep merging
+        consecutive segments until the duration of the merged segment exceeds `max_segment_duration`.
 
         For the case of a multi-channel cut with multiple alignments, we can either trim
         while respecting the supervision channels (in which case output cut has the same channels
@@ -531,6 +533,21 @@ def trim_to_alignments(
         .. hint:: If a MultiCut is trimmed and the resulting trimmed cut contains a single channel,
             we convert it to a MonoCut.
 
+        .. hint:: If you have a Cut with multiple supervision segments and you want to trim it to
+            the word-level alignment, you can use the :meth:`.Cut.merge_supervisions` method
+            first to merge the supervisions into a single one, followed by the
+            :meth:`.Cut.trim_to_alignments` method. For example::
+
+                >>> cut = cut.merge_supervisions(type='word', delimiter=' ')
+                >>> cut = cut.trim_to_alignments(type='word', max_pause=1.0)
+
+        .. hint:: The above technique can also be used to segment long cuts into roughly equal
+            duration segments, while respecting alignment boundaries. For example, to split a
+            Cut into 10s segments, you can do::
+
+                >>> cut = cut.merge_supervisions(type='word', delimiter=' ')
+                >>> cut = cut.trim_to_alignments(type='word', max_pause=10.0, max_segment_duration=10.0)
+
         :param type: The type of the alignment to trim to (e.g. "word").
         :param max_pause: The maximum pause allowed between the alignments to merge them. If ``None``,
             no merging will be performed. [default: None]
@@ -546,6 +563,10 @@ def trim_to_alignments(
             # Set to a negative value so that no merging is performed.
             max_pause = -1.0
 
+        if max_segment_duration is None:
+            # Set to the cut duration so that resulting segments are always smaller.
+            max_segment_duration = self.duration
+
         # For the implementation, we first create new supervisions for the cut, and then
         # use the `trim_to_supervisions` method to do the actual trimming.
         new_supervisions = []
@@ -561,14 +582,20 @@ def trim_to_alignments(
             # Merge the alignments if needed. We also keep track of the indices of the
             # merged alignments in the original list. This is needed to create the
             # `alignment` field in the new supervisions.
+            # NOTE: We use the `AlignmentItem` class here for convenience --- the merged
+            # alignments are not actual alignment items, but rather just a way to keep
+            # track of merged segments.
             merged_alignments = [(alignments[0], [0])]
             for i, item in enumerate(alignments[1:]):
                 # If alignment item is blank, skip it. Sometimes, blank alignment items
                 # are used to denote pauses in the utterance.
                 if item.symbol.strip() == "":
                     continue
                 prev_item, prev_indices = merged_alignments[-1]
-                if item.start - prev_item.end <= max_pause:
+                if (
+                    item.start - prev_item.end <= max_pause
+                    and item.end - prev_item.start <= max_segment_duration
+                ):
                     new_item = AlignmentItem(
                         symbol=delimiter.join([prev_item.symbol, item.symbol]),
                         start=prev_item.start,

lhotse/cut/data.py

@@ -1050,6 +1050,7 @@ def filter_supervisions(
     @abstractmethod
     def merge_supervisions(
         self,
+        merge_policy: str = "delimiter",
         custom_merge_fn: Optional[Callable[[str, Iterable[Any]], Any]] = None,
         **kwargs,
     ) -> "DataCut":

lhotse/cut/mixed.py

@@ -1341,37 +1341,46 @@ def map_supervisions(
         return new_mixed_cut
 
     def merge_supervisions(
-        self, custom_merge_fn: Optional[Callable[[str, Iterable[Any]], Any]] = None
+        self,
+        merge_policy: str = "delimiter",
+        custom_merge_fn: Optional[Callable[[str, Iterable[Any]], Any]] = None,
     ) -> "MixedCut":
         """
         Return a copy of the cut that has all of its supervisions merged into
         a single segment.
 
         The new start is the start of the earliest superivion, and the new duration
-        is a minimum spanning duration for all the supervisions.
-
-        The text fields are concatenated with a whitespace, and all other string fields
-        (including IDs) are prefixed with "cat#" and concatenated with a hash symbol "#".
-        This is also applied to ``custom`` fields. Fields with a ``None`` value are omitted.
+        is a minimum spanning duration for all the supervisions. The text fields are
+        concatenated with a whitespace.
 
         .. note:: If you're using individual tracks of a mixed cut, note that this transform
              drops all the supervisions in individual tracks and assigns the merged supervision
              in the first :class:`.DataCut` found in ``self.tracks``.
 
+        :param merge_policy: one of "keep_first" or "delimiter". If "keep_first", we
+            keep only the first segment's field value, otherwise all string fields
+            (including IDs) are prefixed with "cat#" and concatenated with a hash symbol "#".
+            This is also applied to ``custom`` fields. Fields with a ``None`` value are omitted.
         :param custom_merge_fn: a function that will be called to merge custom fields values.
             We expect ``custom_merge_fn`` to handle all possible custom keys.
             When not provided, we will treat all custom values as strings.
             It will be called roughly like:
             ``custom_merge_fn(custom_key, [s.custom[custom_key] for s in sups])``
         """
+        merge_func_ = partial(
+            merge_items_with_delimiter,
+            delimiter="#",
+            return_first=(merge_policy == "keep_first"),
+        )
+
         # "m" stands for merged in variable names below
 
         if custom_merge_fn is not None:
             # Merge custom fields with the user-provided function.
             merge_custom = custom_merge_fn
         else:
             # Merge the string representations of custom fields.
-            merge_custom = lambda k, vs: merge_items_with_delimiter(map(str, vs))
+            merge_custom = lambda k, vs: merge_func_(map(str, vs))
 
         sups = sorted(self.supervisions, key=lambda s: s.start)
 
@@ -1400,18 +1409,18 @@ def merge_supervisions(
             )
 
         msup = SupervisionSegment(
-            id=merge_items_with_delimiter(s.id for s in sups),
+            id=merge_func_(s.id for s in sups),
             # Make merged recording_id is a mix of recording_ids.
-            recording_id=merge_items_with_delimiter(s.recording_id for s in sups),
+            recording_id=merge_func_(s.recording_id for s in sups),
             start=mstart,
             duration=mduration,
             # Hardcode -1 to indicate no specific channel, as the supervisions might have
             # come from different channels in their original recordings.
             channel=-1,
             text=" ".join(s.text for s in sups if s.text),
-            speaker=merge_items_with_delimiter(s.speaker for s in sups if s.speaker),
-            language=merge_items_with_delimiter(s.language for s in sups if s.language),
-            gender=merge_items_with_delimiter(s.gender for s in sups if s.gender),
+            speaker=merge_func_(s.speaker for s in sups if s.speaker),
+            language=merge_func_(s.language for s in sups if s.language),
+            gender=merge_func_(s.gender for s in sups if s.gender),
             custom={
                 k: merge_custom(
                     k,

lhotse/cut/mono.py

@@ -2,7 +2,7 @@
 import math
 import warnings
 from dataclasses import dataclass
-from functools import reduce
+from functools import partial, reduce
 from operator import add
 from typing import Any, Callable, Iterable, List, Optional
 
@@ -194,33 +194,42 @@ def reverb_rir(
             )
 
     def merge_supervisions(
-        self, custom_merge_fn: Optional[Callable[[str, Iterable[Any]], Any]] = None
+        self,
+        merge_policy: str = "delimiter",
+        custom_merge_fn: Optional[Callable[[str, Iterable[Any]], Any]] = None,
     ) -> "MonoCut":
         """
         Return a copy of the cut that has all of its supervisions merged into
         a single segment.
 
         The new start is the start of the earliest superivion, and the new duration
-        is a minimum spanning duration for all the supervisions.
-
-        The text fields are concatenated with a whitespace, and all other string fields
-        (including IDs) are prefixed with "cat#" and concatenated with a hash symbol "#".
-        This is also applied to ``custom`` fields. Fields with a ``None`` value are omitted.
+        is a minimum spanning duration for all the supervisions. The text fields of
+        all segments are concatenated with a whitespace.
 
+        :param merge_policy: one of "keep_first" or "delimiter". If "keep_first", we
+            keep only the first segment's field value, otherwise all string fields
+            (including IDs) are prefixed with "cat#" and concatenated with a hash symbol "#".
+            This is also applied to ``custom`` fields. Fields with a ``None`` value are omitted.
         :param custom_merge_fn: a function that will be called to merge custom fields values.
             We expect ``custom_merge_fn`` to handle all possible custom keys.
             When not provided, we will treat all custom values as strings.
             It will be called roughly like:
             ``custom_merge_fn(custom_key, [s.custom[custom_key] for s in sups])``
         """
+        merge_func_ = partial(
+            merge_items_with_delimiter,
+            delimiter="#",
+            return_first=(merge_policy == "keep_first"),
+        )
+
         # "m" stands for merged in variable names below
 
         if custom_merge_fn is not None:
             # Merge custom fields with the user-provided function.
             merge_custom = custom_merge_fn
         else:
             # Merge the string representations of custom fields.
-            merge_custom = lambda k, vs: merge_items_with_delimiter(map(str, vs))
+            merge_custom = lambda k, vs: merge_func_(map(str, vs))
 
         sups = sorted(self.supervisions, key=lambda s: s.start)
 
@@ -249,15 +258,15 @@ def merge_supervisions(
             )
 
         msup = SupervisionSegment(
-            id=merge_items_with_delimiter(s.id for s in sups),
+            id=merge_func_(s.id for s in sups),
             recording_id=sups[0].recording_id,
             start=mstart,
             duration=mduration,
             channel=sups[0].channel,
             text=" ".join(s.text for s in sups if s.text),
-            speaker=merge_items_with_delimiter(s.speaker for s in sups if s.speaker),
-            language=merge_items_with_delimiter(s.language for s in sups if s.language),
-            gender=merge_items_with_delimiter(s.gender for s in sups if s.gender),
+            speaker=merge_func_(s.speaker for s in sups if s.speaker),
+            language=merge_func_(s.language for s in sups if s.language),
+            gender=merge_func_(s.gender for s in sups if s.gender),
             custom={
                 k: merge_custom(
                     k,