corpus transform: add SplitRecordingsIntoPiecesJob

corpus/transform.py

@@ -6,6 +6,7 @@
     "MergeStrategy",
     "MergeCorpusSegmentsAndAudioJob",
     "ShiftCorpusSegmentStartJob",
+    "SplitRecordingsIntoPiecesJob",
     "ApplyLexiconToCorpusJob",
     "MapRecordingsJob",
     "IntegrateContextIntoOrthJob",
@@ -14,6 +15,7 @@
 import bisect
 import collections
 import enum
+import json
 import logging
 import math
 import os
@@ -566,6 +568,60 @@ def run(self):
             segments_outfile.writelines(segment_file_names)
 
 
+class SplitRecordingsIntoPiecesJob(Job):
+    """
+    Replaces the segments of each recording by pieces of at most `max_piece_seconds`, cut at the
+    pauses of the diarization (see `split_pieces`) or at `max_piece_seconds` without diarization.
+
+    `corpus.xml.gz` holds the pieces as segments `1`, `2`, ... without orth, `summary.json` the
+    recordings, pieces and hours.
+    """
+
+    def __init__(self, *, corpus: tk.Path, rttm_dir: tk.Path | None = None, max_piece_seconds: float = 120.0):
+        """
+        :param corpus: bliss corpus, recordings named by their audio file
+        :param rttm_dir: `rttm/<name>.rttm` per recording
+        :param max_piece_seconds: maximal piece duration
+        """
+        self.corpus = corpus
+        self.rttm_dir = rttm_dir
+        self.max_piece_seconds = max_piece_seconds
+
+        self.out_corpus = self.output_path("corpus.xml.gz")
+        self.out_summary = self.output_path("summary.json")
+
+    def tasks(self):
+        yield Task("run", mini_task=True)
+
+    def run(self):
+        import soundfile
+
+        from i6_core.lib.intervals import split_pieces
+        from i6_core.lib.rttm import parse_rttm
+
+        c = corpus.Corpus()
+        c.load(self.corpus.get_path())
+        summary = {"recordings": 0, "pieces": 0, "hours": 0.0}
+        for recording in c.all_recordings():
+            duration = soundfile.info(recording.audio).duration
+            speech = []
+            if self.rttm_dir is not None:
+                rttm = os.path.join(self.rttm_dir.get_path(), f"{recording.name}.rttm")
+                speech = [(start, end) for _, start, end in parse_rttm(rttm)]
+            pieces = split_pieces(speech, duration=duration, max_seconds=self.max_piece_seconds)
+            recording.segments = []
+            for i, (start, end) in enumerate(pieces, start=1):
+                segment = corpus.Segment(start=start, end=end)
+                segment.name = str(i)
+                recording.add_segment(segment)
+            summary["recordings"] += 1
+            summary["pieces"] += len(pieces)
+            summary["hours"] += duration / 3600
+        c.dump(self.out_corpus.get_path())
+        with open(self.out_summary.get_path(), "wt") as f:
+            json.dump(summary, f, indent=2)
+
+
 class LexiconStrategy(enum.Enum):
     PICK_FIRST = 0
 

lib/intervals.py

@@ -0,0 +1,38 @@
+__all__ = ["Interval", "split_pieces", "union_intervals"]
+
+Interval = tuple[float, float]
+
+
+def union_intervals(intervals: list[Interval]) -> list[Interval]:
+    """
+    :return: the merged intervals, disjoint and sorted
+    """
+    merged = []
+    for start, end in sorted(intervals):
+        if merged and start <= merged[-1][1]:
+            merged[-1] = (merged[-1][0], max(merged[-1][1], end))
+        else:
+            merged.append((start, end))
+    return merged
+
+
+def split_pieces(intervals: list[Interval], *, duration: float, max_seconds: float) -> list[Interval]:
+    """
+    Splits `[0, duration]` into consecutive pieces of at most `max_seconds`. Each piece ends in the
+    middle of the longest gap between the intervals in its second half, or at `max_seconds` if
+    there is no gap there, so all pieces but the last are at least `max_seconds / 2` long.
+
+    :param intervals: `(start, end)` not to be cut if possible, e.g. the speech of a diarization
+    :return: `(start, end)` per piece, covering `[0, duration]`
+    """
+    occupied = union_intervals(intervals)
+    gaps = [(a_end, b_start) for (_, a_end), (b_start, _) in zip(occupied, occupied[1:])]
+    pieces, start = [], 0.0
+    while duration - start > max_seconds:
+        lower, upper = start + max_seconds / 2, start + max_seconds
+        candidates = [(b - a, (a + b) / 2) for a, b in gaps if lower <= (a + b) / 2 <= upper]
+        cut = max(candidates)[1] if candidates else upper
+        pieces.append((start, cut))
+        start = cut
+    pieces.append((start, duration))
+    return pieces