Repository: lhotse-speech/lhotse

Issue:
load_kaldi_data_dir with segments
Hello, when loading a kaldi data dir including `feats.scp` and `segments` into CutSet, what would be the best approach?

I used `lhotse.kaldi.load_kaldi_data_dir` to load a kaldi data dir into `recordings`, `supervisions` and `features`, and then creating a CutSet from `CutSet.from_manifests` fails to find a supervision corresponding to each cut. 

Since each supervision inherits its start time from kaldi `segments`, `create_cut_set_eager` cannot match supervisions with their corresponding cut: https://github.com/lhotse-speech/lhotse/blob/master/lhotse/cut/set.py#L2894-L2899

I got around this by forcing a supervision start and duration to be cut start and duration, but I'm sure there is a cleaner approach to doing this, or I'd appreciate it if you could let me know if I'm missing a better way to load features and supervisions from kaldi data dir.

I'll leave my workaround way to load a kaldi data dir.

``` python
recs, sups, feats = lhotse.kaldi.load_kaldi_data_dir(
        args.kaldi_dir,
        sampling_rate=args.sampling_rate,
        frame_shift=args.frame_shift,
    )

# here, matching supervisions with cuts fails
cutset = lhotse.CutSet.from_manifests(
        recordings=recs,
        supervisions=sups
        features=feats,
    )

sup_dict = {}
for sup in sups:
    sup_dict[sup.id] = sup

for cut in cutset:
    storage_key = cut.features.storage_key
    # match supervision with cut by features.storage_key
    supervision = sup_dict[storage_key]
    # if duration is too different between supervision cut,
    # something wrong is happening
    sup_dur = supervision.duration
    if sup_dur - cut.duration > 0.05:
        raise ValueError("supervision duration and cut duration is too different.")

    # discard supervision.start inheritated from segments
    # and assign cut.duration and cut.start to supervision
    supervision.duration = cut.duration
    supervision.start = cut.start

    cut.supervisions.append(supervision)
```


Discussion:
The issue seems to be that `load_kaldi_data_dir` has hardcoded `start=0` when it creates `Features` manifests. Technically that code should use the `start` time of the segment corresponding to a given feature matrix. See:

https://github.com/lhotse-speech/lhotse/blob/9fb25d6408154755406fdfe1994471f9427ab488/lhotse/kaldi.py#L225

Would you be interested in making a PR with these changes? (CC @jtrmal @desh2608 FYI since you were also using Kaldi import thing).
nice catch
y.

On Tue, Mar 7, 2023 at 3:25 PM Piotr Żelasko ***@***.***>
wrote:

> The issue seems to be that load_kaldi_data_dir has hardcoded start=0 when
> it creates Features manifests. Technically that code should use the start
> time of the segment corresponding to a given feature matrix. Would you be
> interested in making a PR with these changes? (CC @jtrmal
> <https://github.com/jtrmal> @desh2608 <https://github.com/desh2608> FYI
> since you were also using Kaldi import thing).
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/lhotse-speech/lhotse/issues/987#issuecomment-1458820722>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ACUKYXZCTPWFJEBMI7CWSKDW26KUTANCNFSM6AAAAAAVHYSZHY>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>

> Would you be interested in making a PR with these changes? (CC @jtrmal @desh2608 FYI since you were also using Kaldi import thing).

Sure, I'm happy to contribute 

Resolving patch:
diff --git a/lhotse/kaldi.py b/lhotse/kaldi.py
--- a/lhotse/kaldi.py
+++ b/lhotse/kaldi.py
@@ -147,7 +147,17 @@ def fix_id(t: str) -> str:
     supervision_set = None
     segments = path / "segments"
     utt2spk_f = path / "utt2spk"
+    feats_scp = path / "feats.scp"
+
+    # load mapping from utt_id to start and duration
+    utt_id_to_start_and_duration = load_start_and_duration(
+        segments_path=segments,
+        feats_path=feats_scp,
+        frame_shift=frame_shift,
+    )
+
     if segments.is_file():
+        supervisions = []
         with segments.open() as f:
             supervision_segments = [sup_string.strip().split() for sup_string in f]
 
@@ -156,27 +166,33 @@ def fix_id(t: str) -> str:
         genders = load_kaldi_text_mapping(path / "spk2gender")
         languages = load_kaldi_text_mapping(path / "utt2lang")
 
-        # to support <end-time> == -1 in segments file
-        # https://kaldi-asr.org/doc/extract-segments_8cc.html
-        # <end-time> of -1 means the segment runs till the end of the WAV file
-        supervision_set = SupervisionSet.from_segments(
-            SupervisionSegment(
-                id=fix_id(segment_id),
-                recording_id=recording_id,
-                start=float(start),
-                duration=add_durations(
+        for segment_id, recording_id, start, end in supervision_segments:
+            if utt_id_to_start_and_duration:
+                # use duration computed from feats.scp
+                _, duration = utt_id_to_start_and_duration[segment_id]
+            else:
+                # to support <end-time> == -1 in segments file
+                # https://kaldi-asr.org/doc/extract-segments_8cc.html
+                # <end-time> of -1 means the segment runs till the end of the WAV file
+                duration = add_durations(
                     float(end) if end != "-1" else durations[recording_id],
                     -float(start),
                     sampling_rate=sampling_rate,
-                ),
-                channel=0,
-                text=texts[segment_id],
-                language=languages[segment_id],
-                speaker=fix_id(speakers[segment_id]),
-                gender=genders[speakers[segment_id]],
+                )
+            supervisions.append(
+                SupervisionSegment(
+                    id=fix_id(segment_id),
+                    recording_id=recording_id,
+                    start=float(start),
+                    duration=duration,
+                    channel=0,
+                    text=texts[segment_id],
+                    language=languages[segment_id],
+                    speaker=fix_id(speakers[segment_id]),
+                    gender=genders[speakers[segment_id]],
+                )
             )
-            for segment_id, recording_id, start, end in supervision_segments
-        )
+        supervision_set = SupervisionSet.from_segments(supervisions)
     elif utt2spk_f.is_file():
         # segments file does not exist => provided supervision
         # corresponds to whole recordings
@@ -202,7 +218,6 @@ def fix_id(t: str) -> str:
         )
 
     feature_set = None
-    feats_scp = path / "feats.scp"
     if feats_scp.exists() and is_module_available("kaldi_native_io"):
         if frame_shift is not None:
             import kaldi_native_io
@@ -215,6 +230,13 @@ def fix_id(t: str) -> str:
                     utt_id, ark = line.strip().split(maxsplit=1)
                     mat_shape = kaldi_native_io.MatrixShape.read(ark)
 
+                    # start time is from segments
+                    if utt_id_to_start_and_duration:
+                        start, duration = utt_id_to_start_and_duration[utt_id]
+                    else:
+                        start = 0
+                        duration = mat_shape.num_rows * frame_shift
+
                     features.append(
                         Features(
                             type="kaldi_native_io",
@@ -222,8 +244,8 @@ def fix_id(t: str) -> str:
                             num_features=mat_shape.num_cols,
                             frame_shift=frame_shift,
                             sampling_rate=sampling_rate,
-                            start=0,
-                            duration=mat_shape.num_rows * frame_shift,
+                            start=start,
+                            duration=duration,
                             storage_type=KaldiReader.name,
                             storage_path=ark,
                             storage_key=utt_id,
@@ -432,6 +454,41 @@ def export_to_kaldi(
             )
 
 
+def load_start_and_duration(
+    segments_path: Path = None,
+    feats_path: Path = None,
+    frame_shift: Optional[Seconds] = None,
+) -> Dict[Tuple, None]:
+    """
+    Load start time from segments and duration from feats,
+    when both segments and feats.scp are available.
+    """
+    utt_id_to_start_and_duration = {}
+    if (
+        segments_path.is_file()
+        and feats_path.is_file()
+        and is_module_available("kaldi_native_io")
+        and frame_shift is not None
+    ):
+        import kaldi_native_io
+
+        with segments_path.open() as segments_f, feats_path.open() as feats_f:
+            for segments_line, feats_line in zip(segments_f, feats_f):
+                segment_id, _, start, _ = segments_line.strip().split()
+                utt_id, ark = feats_line.strip().split(maxsplit=1)
+                if segment_id != utt_id:
+                    raise ValueError(f"{segments_path} and {feats_path} not aligned.")
+
+                mat_shape = kaldi_native_io.MatrixShape.read(ark)
+                duration = mat_shape.num_rows * frame_shift
+
+                utt_id_to_start_and_duration[utt_id] = (
+                    float(start),
+                    duration,
+                )
+    return utt_id_to_start_and_duration
+
+
 def load_kaldi_text_mapping(
     path: Path, must_exist: bool = False, float_vals: bool = False
 ) -> Dict[str, Optional[str]]:


Test patch:
diff --git a/test/fixtures/mini_librispeech2/conf/mfcc.conf b/test/fixtures/mini_librispeech2/conf/mfcc.conf
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/conf/mfcc.conf
@@ -0,0 +1 @@
+--use-energy=false   # only non-default option.
diff --git a/test/fixtures/mini_librispeech2/data/raw_mfcc_mini_librispeech2.1.ark b/test/fixtures/mini_librispeech2/data/raw_mfcc_mini_librispeech2.1.ark
new file mode 100644
Binary files /dev/null and b/test/fixtures/mini_librispeech2/data/raw_mfcc_mini_librispeech2.1.ark differ
diff --git a/test/fixtures/mini_librispeech2/data/raw_mfcc_mini_librispeech2.1.scp b/test/fixtures/mini_librispeech2/data/raw_mfcc_mini_librispeech2.1.scp
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/data/raw_mfcc_mini_librispeech2.1.scp
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 data/raw_mfcc_mini_librispeech2.1.ark:20
+lbi-6241-61943-0000 data/raw_mfcc_mini_librispeech2.1.ark:25996
+lbi-8842-304647-0000 data/raw_mfcc_mini_librispeech2.1.ark:35151
diff --git a/test/fixtures/mini_librispeech2/feats.scp b/test/fixtures/mini_librispeech2/feats.scp
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/feats.scp
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 data/raw_mfcc_mini_librispeech2.1.ark:20
+lbi-6241-61943-0000 data/raw_mfcc_mini_librispeech2.1.ark:25996
+lbi-8842-304647-0000 data/raw_mfcc_mini_librispeech2.1.ark:35151
diff --git a/test/fixtures/mini_librispeech2/frame_shift b/test/fixtures/mini_librispeech2/frame_shift
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/frame_shift
@@ -0,0 +1 @@
+0.01
diff --git a/test/fixtures/mini_librispeech2/lhotse/features.jsonl.gz b/test/fixtures/mini_librispeech2/lhotse/features.jsonl.gz
new file mode 100644
Binary files /dev/null and b/test/fixtures/mini_librispeech2/lhotse/features.jsonl.gz differ
diff --git a/test/fixtures/mini_librispeech2/lhotse/recordings.jsonl.gz b/test/fixtures/mini_librispeech2/lhotse/recordings.jsonl.gz
new file mode 100644
Binary files /dev/null and b/test/fixtures/mini_librispeech2/lhotse/recordings.jsonl.gz differ
diff --git a/test/fixtures/mini_librispeech2/lhotse/supervisions.jsonl.gz b/test/fixtures/mini_librispeech2/lhotse/supervisions.jsonl.gz
new file mode 100644
Binary files /dev/null and b/test/fixtures/mini_librispeech2/lhotse/supervisions.jsonl.gz differ
diff --git a/test/fixtures/mini_librispeech2/reco2dur b/test/fixtures/mini_librispeech2/reco2dur
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/reco2dur
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 19.889999389648438
+lbi-6241-61943-0000 6.949999809265137
+lbi-8842-304647-0000 9.710000038146973
diff --git a/test/fixtures/mini_librispeech2/segments b/test/fixtures/mini_librispeech2/segments
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/segments
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 lbi-3536-23268-0000 1.0 20.89
+lbi-6241-61943-0000 lbi-6241-61943-0000 0 6.95
+lbi-8842-304647-0000 lbi-8842-304647-0000 0 9.71
diff --git a/test/fixtures/mini_librispeech2/spk2gender b/test/fixtures/mini_librispeech2/spk2gender
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/spk2gender
@@ -0,0 +1,3 @@
+lbi-3536-23268 f
+lbi-6241-61943 m
+lbi-8842-304647 f
diff --git a/test/fixtures/mini_librispeech2/spk2utt b/test/fixtures/mini_librispeech2/spk2utt
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/spk2utt
@@ -0,0 +1,3 @@
+lbi-3536-23268 lbi-3536-23268-0000
+lbi-6241-61943 lbi-6241-61943-0000
+lbi-8842-304647 lbi-8842-304647-0000
diff --git a/test/fixtures/mini_librispeech2/text b/test/fixtures/mini_librispeech2/text
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/text
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 SIR EDWARD NOT WHOLLY DISCOURAGED BY THE DENIAL WITH WHICH DORRIFORTH HAD WITH DELICACY ACQUAINTED HIM STILL HOPED FOR A KIND RECEPTION AND WAS SO OFTEN AT THE HOUSE OF MISSUS HORTON THAT LORD FREDERICK'S JEALOUSY WAS EXCITED AND THE TORTURES HE SUFFERED IN CONSEQUENCE CONVINCED HIM BEYOND A DOUBT OF THE SINCERITY OF HIS AFFECTION
+lbi-6241-61943-0000 ON THE SECOND OF THE MONTH AT TWO IN THE MORNING OUR PRECIOUS CARGO OF LUGGAGE WAS TAKEN ON BOARD THE GOOD SHIP VALKYRIE
+lbi-8842-304647-0000 HE LIVES THY LOSS HE DIES FROM EVERY LIMB MANGLED BY THEE LIGHTNINGS OF GODHEAD SHINE FROM WHICH THY DARKNESS HATH NOT WHERE TO HIDE
diff --git a/test/fixtures/mini_librispeech2/utt2dur b/test/fixtures/mini_librispeech2/utt2dur
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/utt2dur
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 19.89
+lbi-6241-61943-0000 6.95
+lbi-8842-304647-0000 9.71
diff --git a/test/fixtures/mini_librispeech2/utt2num_frames b/test/fixtures/mini_librispeech2/utt2num_frames
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/utt2num_frames
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 1987
+lbi-6241-61943-0000 693
+lbi-8842-304647-0000 969
diff --git a/test/fixtures/mini_librispeech2/utt2spk b/test/fixtures/mini_librispeech2/utt2spk
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/utt2spk
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000 lbi-3536-23268
+lbi-6241-61943-0000 lbi-6241-61943
+lbi-8842-304647-0000 lbi-8842-304647
diff --git a/test/fixtures/mini_librispeech2/wav.scp b/test/fixtures/mini_librispeech2/wav.scp
new file mode 100644
--- /dev/null
+++ b/test/fixtures/mini_librispeech2/wav.scp
@@ -0,0 +1,3 @@
+lbi-3536-23268-0000   sox  -r  16000  -b16  -c  1  --null  -t  wav  -  synth  318240s  sine  300-3300  |
+lbi-6241-61943-0000   sox  -r  16000  -b16  -c  1  --null  -t  wav  -  synth  111200s  sine  300-3300  |
+lbi-8842-304647-0000  sox  -r  16000  -b16  -c  1  --null  -t  wav  -  synth  155360s  sine  300-3300  |
diff --git a/test/test_kaldi_dirs.py b/test/test_kaldi_dirs.py
--- a/test/test_kaldi_dirs.py
+++ b/test/test_kaldi_dirs.py
@@ -13,6 +13,7 @@
 FIXTURE_PATH = Path(__file__).resolve().parent / "fixtures"
 
 MINILIB_PATH = FIXTURE_PATH / "mini_librispeech"
+MINILIB2_PATH = FIXTURE_PATH / "mini_librispeech2"
 
 
 @pytest.fixture
@@ -191,6 +192,27 @@ def test_kaldi_import(replace):
     assert out[1] == supervision_set
 
 
+def test_kaldi_import_with_feats_scp():
+    fixture_path = MINILIB2_PATH
+    with working_directory(fixture_path):
+        out = lhotse.kaldi.load_kaldi_data_dir(
+            fixture_path,
+            sampling_rate=16000,
+            frame_shift=0.01,
+            use_reco2dur=True,
+        )
+
+    lhotse_dir = "lhotse"
+    with working_directory(fixture_path / lhotse_dir):
+        recording_set = lhotse.RecordingSet.from_jsonl("recordings.jsonl.gz")
+        supervision_set = lhotse.SupervisionSet.from_jsonl("supervisions.jsonl.gz")
+        feature_set = lhotse.FeatureSet.from_jsonl("features.jsonl.gz")
+
+    assert out[0] == recording_set
+    assert out[1] == supervision_set
+    assert out[2] == feature_set
+
+
 def open_and_load(path):
     return lhotse.kaldi.load_kaldi_text_mapping(Path(path))