diff --git a/examples/lhotse_usage.py b/examples/lhotse_usage.py new file mode 100644 index 0000000..ca5413e --- /dev/null +++ b/examples/lhotse_usage.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python +# encoding: utf-8 + +""" +Example: Using pyannote-database with Lhotse + +This example demonstrates how to use the LhotseProtocol to integrate +Lhotse datasets into pyannote.database. + +Usage: + python lhotse_usage.py /path/to/ami/directory + +Requirements: + pip install lhotse pyannote.database +""" + +import sys +from pathlib import Path + +from pyannote.database import LhotseProtocol + + +def example_basic_usage(ami_dir): + """Basic usage of LhotseProtocol + + This example shows how to create a protocol from Lhotse data + and iterate over the files. + + Args: + ami_dir: Path to the AMI directory containing Lhotse JSONL files + """ + import lhotse + + ami_path = Path(ami_dir) + + # Load train/dev/test recordings and supervisions from the compressed JSONL files + # Using ami-ihm subset as an example + # Map from protocol split names to file names + split_mapping = {"train": "train", "development": "dev", "test": "test"} + file_splits = ["train", "dev", "test"] + + # Load all recordings (needed to match supervisions) + all_recordings = [] + all_supervisions = [] + + for file_split in file_splits: + rec_file = ami_path / f"ami-ihm_recordings_{file_split}.jsonl.gz" + sup_file = ami_path / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + + if rec_file.exists() and sup_file.exists(): + print(f" Loading {file_split}: {rec_file.name} and {sup_file.name}...") + all_recordings.extend(lhotse.load_manifest(str(rec_file))) + all_supervisions.extend(lhotse.load_manifest(str(sup_file))) + else: + print(f" Skipping {file_split}: files not found") + + if not all_recordings: + raise FileNotFoundError(f"No ami-ihm recordings files found in {ami_dir}") + if not all_supervisions: + raise FileNotFoundError(f"No ami-ihm supervisions files found in {ami_dir}") + + recording_set = lhotse.RecordingSet.from_recordings(all_recordings) + supervision_set = lhotse.SupervisionSet(all_supervisions) + print(f" Loaded {len(recording_set)} recordings and {len(supervision_set)} supervisions") + + # Define train/development/test splits by reading the actual split files + subset_split = {} + for protocol_split, file_split in split_mapping.items(): + sup_file = ami_path / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + if sup_file.exists(): + sups = lhotse.load_manifest(str(sup_file)) + subset_split[protocol_split] = sorted(set(sup.recording_id for sup in sups)) + + # Create the protocol + protocol = LhotseProtocol( + recording_set=recording_set, + supervision_set=supervision_set, + subset_split=subset_split, + ) + + # Iterate over training files + for file in protocol.train(): + print(f"URI: {file['uri']}") + print(f"Annotation: {file['annotation']}") + print(f"Annotated regions: {file['annotated']}") + # file also contains: + # - file['recording']: Lhotse Recording object + # - file['supervisions']: Lhotse SupervisionSet + + +def example_with_preprocessors(ami_dir): + """Using preprocessors with LhotseProtocol + + Preprocessors allow you to add custom fields to each file + that are computed on-the-fly. + + Args: + ami_dir: Path to the AMI directory containing Lhotse JSONL files + """ + import lhotse + from pathlib import Path + + ami_path = Path(ami_dir) + + # Load train/dev/test recordings and supervisions from the compressed JSONL files + # Map from protocol split names to file names + split_mapping = {"train": "train", "development": "dev", "test": "test"} + file_splits = ["train", "dev", "test"] + + # Load all recordings and supervisions + all_recordings = [] + all_supervisions = [] + + for file_split in file_splits: + rec_file = ami_path / f"ami-ihm_recordings_{file_split}.jsonl.gz" + sup_file = ami_path / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + + if rec_file.exists() and sup_file.exists(): + print(f" Loading {file_split}: {rec_file.name} and {sup_file.name}...") + all_recordings.extend(lhotse.load_manifest(str(rec_file))) + all_supervisions.extend(lhotse.load_manifest(str(sup_file))) + + recording_set = lhotse.RecordingSet.from_recordings(all_recordings) + supervision_set = lhotse.SupervisionSet(all_supervisions) + print(f" Loaded {len(recording_set)} recordings and {len(supervision_set)} supervisions") + + # Define train/development/test splits by reading the actual split files + subset_split = {} + for protocol_split, file_split in split_mapping.items(): + sup_file = ami_path / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + if sup_file.exists(): + sups = lhotse.load_manifest(str(sup_file)) + subset_split[protocol_split] = sorted(set(sup.recording_id for sup in sups)) + + # Define preprocessors (only use callable preprocessors with Lhotse data) + # Note: Template-based preprocessors should reference keys that exist in the file dict + preprocessors = { + "duration": lambda file: file["recording"].duration, + "num_speakers": lambda file: len(set( + s.speaker for s in file["supervisions"] if s.speaker is not None + )), + "audio": lambda file: file["recording"].sources[0].audio, + } + + protocol = LhotseProtocol( + recording_set=recording_set, + supervision_set=supervision_set, + subset_split=subset_split, + preprocessors=preprocessors, + ) + + # Now each file also has 'duration' and 'num_speakers' keys + for file in protocol.train(): + print(f"Duration: {file['duration']}") + print(f"Num speakers: {file['num_speakers']}") + + +def example_with_registry(ami_dir): + """Programmatic protocol registration + + You can also register LhotseProtocol instances with the registry + for use with existing pyannote tools and workflows. + + Args: + ami_dir: Path to the AMI directory containing Lhotse JSONL files + """ + import lhotse + from pathlib import Path + from pyannote.database import registry, Database + + ami_path = Path(ami_dir) + + # Load train/dev/test recordings and supervisions from the compressed JSONL files + # Map from protocol split names to file names + split_mapping = {"train": "train", "development": "dev", "test": "test"} + file_splits = ["train", "dev", "test"] + + # Load all recordings and supervisions + all_recordings = [] + all_supervisions = [] + + for file_split in file_splits: + rec_file = ami_path / f"ami-ihm_recordings_{file_split}.jsonl.gz" + sup_file = ami_path / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + + if rec_file.exists() and sup_file.exists(): + print(f" Loading {file_split}: {rec_file.name} and {sup_file.name}...") + all_recordings.extend(lhotse.load_manifest(str(rec_file))) + all_supervisions.extend(lhotse.load_manifest(str(sup_file))) + + recording_set = lhotse.RecordingSet.from_recordings(all_recordings) + supervision_set = lhotse.SupervisionSet(all_supervisions) + print(f" Loaded {len(recording_set)} recordings and {len(supervision_set)} supervisions") + + # Define train/development/test splits by reading the actual split files + subset_split = {} + for protocol_split, file_split in split_mapping.items(): + sup_file = ami_path / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + if sup_file.exists(): + sups = lhotse.load_manifest(str(sup_file)) + subset_split[protocol_split] = sorted(set(sup.recording_id for sup in sups)) + + # Create the protocol + protocol = LhotseProtocol( + recording_set=recording_set, + supervision_set=supervision_set, + subset_split=subset_split, + ) + + # Create a custom Database class and register protocols + class LhotseAMIDatabase(Database): + def __init__(self): + super().__init__() + self.register_protocol( + "SpeakerDiarization", + "Benchmark", + lambda preprocessors=None: LhotseProtocol( + recording_set=recording_set, + supervision_set=supervision_set, + subset_split=subset_split, + preprocessors=preprocessors, + ), + ) + + # Register the database + registry.databases["LhotseAMI"] = LhotseAMIDatabase + + # Now you can access it like any other protocol + protocol = registry.get_protocol("LhotseAMI.SpeakerDiarization.Benchmark") + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python lhotse_usage.py /path/to/ami/directory") + print() + print("Example:") + print(" python lhotse_usage.py /Users/samco/datasets/AMI") + sys.exit(1) + + ami_dir = sys.argv[1] + ami_path = Path(ami_dir) + + if not ami_path.exists(): + print(f"Error: Directory does not exist: {ami_dir}") + sys.exit(1) + + print(f"Loading AMI data from: {ami_dir}") + print() + + print("Example 1: Basic usage") + print("=" * 50) + try: + example_basic_usage(ami_dir) + except Exception as e: + print(f"Error: {e}") + + print("\nExample 2: With preprocessors") + print("=" * 50) + try: + example_with_preprocessors(ami_dir) + except Exception as e: + print(f"Error: {e}") + + print("\nExample 3: With registry") + print("=" * 50) + try: + example_with_registry(ami_dir) + except Exception as e: + print(f"Error: {e}") diff --git a/src/pyannote/database/__init__.py b/src/pyannote/database/__init__.py index facc8d6..38ae200 100644 --- a/src/pyannote/database/__init__.py +++ b/src/pyannote/database/__init__.py @@ -47,6 +47,11 @@ from .util import get_unique_identifier from .util import get_label_identifier +try: + from .lhotse import LhotseProtocol +except ImportError: + LhotseProtocol = None + import importlib.metadata __version__ = importlib.metadata.version("pyannote-database") @@ -85,4 +90,5 @@ def get_protocol(name, preprocessors: Optional[Preprocessors] = None) -> Protoco "get_annotated", "get_unique_identifier", "get_label_identifier", + "LhotseProtocol", ] diff --git a/src/pyannote/database/lhotse.py b/src/pyannote/database/lhotse.py new file mode 100644 index 0000000..104e5c2 --- /dev/null +++ b/src/pyannote/database/lhotse.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python +# encoding: utf-8 + +"""Lhotse protocol integration for pyannote.database""" + +import warnings +from typing import Dict, Iterator, Optional + +from pyannote.core import Annotation, Timeline, Segment +from pyannote.database.protocol import SpeakerDiarizationProtocol + + +try: + import lhotse + LHOTSE_IS_AVAILABLE = True +except ImportError: + LHOTSE_IS_AVAILABLE = False + +try: + import meeteval.io + from meeteval.io.seglst import SegLST + MEETEVAL_IS_AVAILABLE = True +except ImportError: + MEETEVAL_IS_AVAILABLE = False + + +def supervision_to_annotation(recording_id: str, supervision_set) -> Annotation: + """Convert Lhotse SupervisionSet to pyannote.core.Annotation + + Parameters + ---------- + recording_id : str + Recording identifier (used as uri in Annotation) + supervision_set : lhotse.SupervisionSet + Lhotse supervision set containing speaker labels + + Returns + ------- + annotation : pyannote.core.Annotation + Speaker diarization annotation + """ + annotation = Annotation(uri=recording_id) + + for supervision in supervision_set: + speaker = supervision.speaker + if speaker is None: + speaker = "unknown" + + segment = Segment(supervision.start, supervision.end) + annotation[segment, speaker] = speaker + + return annotation + + +def supervision_set_to_timeline(supervision_set) -> Timeline: + """Convert Lhotse SupervisionSet to pyannote.core.Timeline + + Parameters + ---------- + supervision_set : lhotse.SupervisionSet + Lhotse supervision set + + Returns + ------- + timeline : pyannote.core.Timeline + Timeline of annotated regions + """ + timeline = Timeline() + + for supervision in supervision_set: + segment = Segment(supervision.start, supervision.end) + timeline.add(segment) + + return timeline + + +def supervisions_to_seglst(recording_id: str, supervision_set) -> Optional["SegLST"]: + """Convert Lhotse SupervisionSet to meeteval SegLST + + Parameters + ---------- + recording_id : str + Recording identifier (used as session_id in SegLST) + supervision_set : lhotse.SupervisionSet + Lhotse supervision set containing transcriptions + + Returns + ------- + seglst : SegLST or None + Transcription as SegLST, or None if meeteval is not available + """ + if not MEETEVAL_IS_AVAILABLE: + return None + + segments = [] + for idx, supervision in enumerate(supervision_set): + if supervision.text is None: + continue + + segment = { + "session_id": recording_id, + "speaker": supervision.speaker or "unknown", + "start": supervision.start, + "end": supervision.end, + "words": supervision.text, + "conf": 1.0, + } + segments.append(segment) + + if not segments: + return SegLST([]) + + return SegLST(segments) + + +class LhotseProtocol(SpeakerDiarizationProtocol): + """Lhotse protocol for pyannote.database + + This protocol integrates Lhotse datasets into pyannote.database, allowing + Lhotse recordings and supervisions to be accessed through the standard + pyannote protocol interface. + + Parameters + ---------- + recording_set : lhotse.RecordingSet + Lhotse recording set + supervision_set : lhotse.SupervisionSet + Lhotse supervision set for the entire dataset + subset_split : dict, optional + Mapping of subset names to recording ids + Example: {"train": ["recording1", "recording2"], + "development": ["recording3"], + "test": ["recording4"]} + audio_dir : str, optional + Path to audio directory (used for preprocessing with preprocessors) + preprocessors : dict, optional + Preprocessors for protocol files + """ + + def __init__( + self, + recording_set, + supervision_set, + subset_split: Optional[Dict[str, list]] = None, + audio_dir: Optional[str] = None, + preprocessors: Optional[Dict] = None, + ): + super().__init__(preprocessors=preprocessors) + + self.recording_set = recording_set + self.supervision_set = supervision_set + self.subset_split = subset_split or {"train": [], "development": [], "test": []} + self.audio_dir = audio_dir + + def train_iter(self) -> Iterator[Dict]: + """Iterate over training files""" + recording_ids = self.subset_split.get("train", []) + + for recording_id in recording_ids: + if recording_id not in self.recording_set: + warnings.warn(f"Recording {recording_id} not found in recording set") + continue + + recording = self.recording_set[recording_id] + supervisions = self.supervision_set.filter(lambda s: s.recording_id == recording_id) + + # Convert lazy filter to eager to check if empty + supervisions_list = supervisions.to_eager() if hasattr(supervisions, 'to_eager') else supervisions + + if len(supervisions_list) == 0: + warnings.warn(f"No supervisions found for recording {recording_id}") + continue + + yield { + "uri": recording_id, + "annotation": supervision_to_annotation(recording_id, supervisions_list), + "annotated": supervision_set_to_timeline(supervisions_list), + "recording": recording, + "supervisions": supervisions_list, + } + + def development_iter(self) -> Iterator[Dict]: + """Iterate over development files""" + recording_ids = self.subset_split.get("development", []) + + for recording_id in recording_ids: + if recording_id not in self.recording_set: + warnings.warn(f"Recording {recording_id} not found in recording set") + continue + + recording = self.recording_set[recording_id] + supervisions = self.supervision_set.filter(lambda s: s.recording_id == recording_id) + + # Convert lazy filter to eager to check if empty + supervisions_list = supervisions.to_eager() if hasattr(supervisions, 'to_eager') else supervisions + + if len(supervisions_list) == 0: + warnings.warn(f"No supervisions found for recording {recording_id}") + continue + + yield { + "uri": recording_id, + "annotation": supervision_to_annotation(recording_id, supervisions_list), + "annotated": supervision_set_to_timeline(supervisions_list), + "recording": recording, + "supervisions": supervisions_list, + } + + def test_iter(self) -> Iterator[Dict]: + """Iterate over test files""" + recording_ids = self.subset_split.get("test", []) + + for recording_id in recording_ids: + if recording_id not in self.recording_set: + warnings.warn(f"Recording {recording_id} not found in recording set") + continue + + recording = self.recording_set[recording_id] + supervisions = self.supervision_set.filter(lambda s: s.recording_id == recording_id) + + # Convert lazy filter to eager to check if empty + supervisions_list = supervisions.to_eager() if hasattr(supervisions, 'to_eager') else supervisions + + yield { + "uri": recording_id, + "annotated": supervision_set_to_timeline(supervisions_list), + "recording": recording, + "supervisions": supervisions_list} diff --git a/tests/test_lhotse.py b/tests/test_lhotse.py new file mode 100644 index 0000000..216706e --- /dev/null +++ b/tests/test_lhotse.py @@ -0,0 +1,262 @@ +#!/usr/bin/env python +# encoding: utf-8 + +"""Test Lhotse protocol integration""" + +import pytest +from unittest.mock import MagicMock, patch +from pyannote.core import Annotation, Timeline, Segment +from pyannote.database.lhotse import ( + LhotseProtocol, + supervision_to_annotation, + supervision_set_to_timeline, +) + + +@pytest.fixture +def mock_recording_set(): + """Create a mock Lhotse RecordingSet""" + recording_set = MagicMock() + recording_set.__contains__ = lambda self, key: key in ["rec1", "rec2", "rec3"] + recording_set.__getitem__ = lambda self, key: MagicMock(id=key) + return recording_set + + +@pytest.fixture +def mock_supervision_set(): + """Create a mock Lhotse SupervisionSet""" + supervision_set = MagicMock() + + # Mock supervisions for rec1 + sup1_1 = MagicMock( + recording_id="rec1", + speaker="spk1", + start=0.0, + end=5.0, + text="hello world", + ) + sup1_2 = MagicMock( + recording_id="rec1", + speaker="spk2", + start=5.0, + end=10.0, + text="goodbye", + ) + + # Mock supervisions for rec2 + sup2_1 = MagicMock( + recording_id="rec2", + speaker="spk1", + start=0.0, + end=3.0, + text="test", + ) + + # Mock supervisions for rec3 + sup3_1 = MagicMock( + recording_id="rec3", + speaker="spk3", + start=0.0, + end=2.0, + text="audio", + ) + + def filter_by_recording(filter_fn): + """Mock filter function that actually filters by recording_id""" + supervisions_by_recording = { + "rec1": [sup1_1, sup1_2], + "rec2": [sup2_1], + "rec3": [sup3_1], + } + + # Apply the filter function to determine which supervisions to return + filtered = [] + for rec_id, supervisions in supervisions_by_recording.items(): + for sup in supervisions: + if filter_fn(sup): + filtered.append(sup) + + mock = MagicMock() + mock.__iter__ = lambda self: iter(filtered) + mock.__len__ = lambda self: len(filtered) + mock.to_eager = lambda: mock # to_eager() returns self for this mock + return mock + + supervision_set.filter = filter_by_recording + supervision_set.__iter__ = lambda self: iter([sup1_1, sup1_2, sup2_1, sup3_1]) + supervision_set.__len__ = lambda self: 4 + + return supervision_set + + +class TestSupervisionConversion: + """Test conversion functions""" + + def test_supervision_to_annotation(self, mock_supervision_set): + """Test converting Lhotse SupervisionSet to pyannote Annotation""" + supervisions = list(mock_supervision_set)[:2] + supervision_set = MagicMock() + supervision_set.__iter__ = lambda self: iter(supervisions) + + annotation = supervision_to_annotation("rec1", supervision_set) + + assert isinstance(annotation, Annotation) + assert annotation.uri == "rec1" + assert len(annotation) > 0 + + def test_supervision_set_to_timeline(self, mock_supervision_set): + """Test converting Lhotse SupervisionSet to pyannote Timeline""" + supervisions = list(mock_supervision_set)[:2] + supervision_set = MagicMock() + supervision_set.__iter__ = lambda self: iter(supervisions) + + timeline = supervision_set_to_timeline(supervision_set) + + assert isinstance(timeline, Timeline) + assert len(timeline) == 2 + + +class TestLhotseProtocol: + """Test LhotseProtocol""" + + def test_protocol_initialization(self, mock_recording_set, mock_supervision_set): + """Test protocol initialization""" + subset_split = { + "train": ["rec1"], + "development": ["rec2"], + "test": ["rec3"], + } + + protocol = LhotseProtocol( + recording_set=mock_recording_set, + supervision_set=mock_supervision_set, + subset_split=subset_split, + ) + + assert protocol.recording_set is mock_recording_set + assert protocol.supervision_set is mock_supervision_set + assert protocol.subset_split == subset_split + + def test_train_iter(self, mock_recording_set, mock_supervision_set): + """Test train_iter method""" + subset_split = { + "train": ["rec1"], + "development": [], + "test": [], + } + + protocol = LhotseProtocol( + recording_set=mock_recording_set, + supervision_set=mock_supervision_set, + subset_split=subset_split, + ) + + files = list(protocol.train_iter()) + assert len(files) == 1 + assert files[0]["uri"] == "rec1" + assert "annotation" in files[0] + assert "annotated" in files[0] + # Verify supervisions are correctly filtered for rec1 + assert len(files[0]["supervisions"]) == 2 + + def test_development_iter(self, mock_recording_set, mock_supervision_set): + """Test development_iter method""" + subset_split = { + "train": [], + "development": ["rec2"], + "test": [], + } + + protocol = LhotseProtocol( + recording_set=mock_recording_set, + supervision_set=mock_supervision_set, + subset_split=subset_split, + ) + + files = list(protocol.development_iter()) + assert len(files) == 1 + assert files[0]["uri"] == "rec2" + # Verify supervisions are correctly filtered for rec2 + assert len(files[0]["supervisions"]) == 1 + + def test_test_iter(self, mock_recording_set, mock_supervision_set): + """Test test_iter method""" + subset_split = { + "train": [], + "development": [], + "test": ["rec3"], + } + + protocol = LhotseProtocol( + recording_set=mock_recording_set, + supervision_set=mock_supervision_set, + subset_split=subset_split, + ) + + files = list(protocol.test_iter()) + assert len(files) == 1 + assert files[0]["uri"] == "rec3" + # Verify supervisions are correctly filtered for rec3 + assert len(files[0]["supervisions"]) == 1 + + def test_public_methods(self, mock_recording_set, mock_supervision_set): + """Test public protocol methods (train, development, test)""" + subset_split = { + "train": ["rec1"], + "development": ["rec2"], + "test": ["rec3"], + } + + protocol = LhotseProtocol( + recording_set=mock_recording_set, + supervision_set=mock_supervision_set, + subset_split=subset_split, + ) + + # Test that public methods return iterators + train_files = list(protocol.train()) + assert len(train_files) == 1 + + dev_files = list(protocol.development()) + assert len(dev_files) == 1 + + test_files = list(protocol.test()) + assert len(test_files) == 1 + + +class TestLhotseProtocolEdgeCases: + """Test edge cases""" + + def test_missing_recording(self, mock_supervision_set): + """Test handling of missing recordings""" + recording_set = MagicMock() + recording_set.__contains__ = lambda self, key: False + + subset_split = { + "train": ["nonexistent"], + "development": [], + "test": [], + } + + protocol = LhotseProtocol( + recording_set=recording_set, + supervision_set=mock_supervision_set, + subset_split=subset_split, + ) + + # Should emit warning and skip + with pytest.warns(UserWarning, match="not found in recording set"): + files = list(protocol.train_iter()) + assert len(files) == 0 + + def test_empty_subset_split(self, mock_recording_set, mock_supervision_set): + """Test with empty subset split""" + protocol = LhotseProtocol( + recording_set=mock_recording_set, + supervision_set=mock_supervision_set, + subset_split={}, + ) + + assert list(protocol.train_iter()) == [] + assert list(protocol.development_iter()) == [] + assert list(protocol.test_iter()) == [] diff --git a/tests/test_lhotse_ami_integration.py b/tests/test_lhotse_ami_integration.py new file mode 100644 index 0000000..bd47d05 --- /dev/null +++ b/tests/test_lhotse_ami_integration.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python +# encoding: utf-8 + +"""Integration tests with real AMI data using Lhotse protocol""" + +import pytest +from pathlib import Path +from pyannote.core import Annotation, Timeline +from pyannote.database.lhotse import LhotseProtocol + + +# Path to AMI dataset +AMI_DATA_PATH = Path("/Users/samco/datasets/AMI") + + +@pytest.fixture(scope="module") +def ami_data_available(): + """Check if AMI data is available""" + if not AMI_DATA_PATH.exists(): + pytest.skip(f"AMI dataset not found at {AMI_DATA_PATH}") + return AMI_DATA_PATH + + +@pytest.fixture +def ami_ihm_lhotse_protocol(ami_data_available): + """Create LhotseProtocol for AMI IHM (Individual Headset Microphone)""" + try: + import lhotse + except ImportError: + pytest.skip("lhotse package not installed") + + # Load AMI IHM data using the predefined train/dev/test splits + # Map from protocol split names to file names + split_mapping = {"train": "train", "development": "dev", "test": "test"} + file_splits = ["train", "dev", "test"] + all_recordings = [] + all_supervisions = [] + + for file_split in file_splits: + rec_file = ami_data_available / f"ami-ihm_recordings_{file_split}.jsonl.gz" + sup_file = ami_data_available / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + + if rec_file.exists() and sup_file.exists(): + all_recordings.extend(lhotse.load_manifest(str(rec_file))) + all_supervisions.extend(lhotse.load_manifest(str(sup_file))) + + recordings = lhotse.RecordingSet.from_recordings(all_recordings) + supervisions = lhotse.SupervisionSet(all_supervisions) + + # Define subset_split using the actual split files + subset_split = {} + for protocol_split, file_split in split_mapping.items(): + sup_file = ami_data_available / f"ami-ihm_supervisions_{file_split}.jsonl.gz" + if sup_file.exists(): + sups = lhotse.load_manifest(str(sup_file)) + subset_split[protocol_split] = sorted(set(sup.recording_id for sup in sups)) + + protocol = LhotseProtocol( + recording_set=recordings, + supervision_set=supervisions, + subset_split=subset_split, + ) + + return protocol + + +@pytest.fixture +def ami_mdm_lhotse_protocol(ami_data_available): + """Create LhotseProtocol for AMI MDM (Multiple Distant Microphone) with multi-channel support""" + try: + import lhotse + except ImportError: + raise ImportError("lhotse package not installed") + + # Load AMI MDM data using the predefined train/dev/test splits + # Map from protocol split names to file names + split_mapping = {"train": "train", "development": "dev", "test": "test"} + file_splits = ["train", "dev", "test"] + all_recordings = [] + all_supervisions = [] + + for file_split in file_splits: + rec_file = ami_data_available / f"ami-mdm_recordings_{file_split}.jsonl.gz" + sup_file = ami_data_available / f"ami-mdm_supervisions_{file_split}.jsonl.gz" + + if rec_file.exists() and sup_file.exists(): + all_recordings.extend(lhotse.load_manifest(str(rec_file))) + all_supervisions.extend(lhotse.load_manifest(str(sup_file))) + + supervisions = lhotse.SupervisionSet(all_supervisions) + recordings = lhotse.RecordingSet.from_recordings(all_recordings) + + # Define subset_split using the actual split files + subset_split = {} + for protocol_split, file_split in split_mapping.items(): + sup_file = ami_data_available / f"ami-mdm_supervisions_{file_split}.jsonl.gz" + if sup_file.exists(): + sups = lhotse.load_manifest(str(sup_file)) + subset_split[protocol_split] = sorted(set(s.recording_id for s in sups)) + + protocol = LhotseProtocol( + recording_set=recordings, + supervision_set=supervisions, + subset_split=subset_split, + ) + + return protocol + + +class TestLhotseAMIIntegration: + """Integration tests with real AMI data""" + + def test_ami_ihm_train_iteration(self, ami_ihm_lhotse_protocol): + """Test iterating over AMI IHM training files""" + files = list(ami_ihm_lhotse_protocol.train()) + + assert len(files) > 0, "No training files found" + + # Check first file structure + file = files[0] + assert "uri" in file + assert "annotation" in file + assert "annotated" in file + assert "recording" in file + assert "supervisions" in file + + # Verify annotation is pyannote Annotation + assert isinstance(file["annotation"], Annotation) + assert file["annotation"].uri == file["uri"] + + # Verify annotated is pyannote Timeline + assert isinstance(file["annotated"], Timeline) + + def test_ami_ihm_development_iteration(self, ami_ihm_lhotse_protocol): + """Test iterating over AMI IHM development files""" + files = list(ami_ihm_lhotse_protocol.development()) + + # May be empty depending on split, but should not fail + for file in files: + assert "uri" in file + assert "annotation" in file + assert isinstance(file["annotation"], Annotation) + + def test_ami_ihm_test_iteration(self, ami_ihm_lhotse_protocol): + """Test iterating over AMI IHM test files""" + files = list(ami_ihm_lhotse_protocol.test()) + + # May be empty depending on split, but should not fail + for file in files: + assert "uri" in file + assert "annotation" in file + # Annotation may be None if there are no supervisions for this file + if file["annotation"] is not None: + assert isinstance(file["annotation"], Annotation) + + def test_ami_ihm_supervisions_have_speakers(self, ami_ihm_lhotse_protocol): + """Test that supervisions contain speaker information""" + files = list(ami_ihm_lhotse_protocol.train()) + + if len(files) > 0: + file = files[0] + supervisions = file["supervisions"] + # IHM supervisions should have speaker field + for sup in supervisions: + assert hasattr(sup, "speaker") + assert sup.speaker is not None + + def test_ami_mdm_train_iteration_multi_channel(self, ami_mdm_lhotse_protocol): + """Test iterating over AMI MDM training files with multi-channel support""" + files = list(ami_mdm_lhotse_protocol.train()) + + assert len(files) > 0, "No training files found" + + # Check first file structure + file = files[0] + assert "uri" in file + assert "annotation" in file + assert "annotated" in file + assert "recording" in file + assert "supervisions" in file + + # Verify annotation is pyannote Annotation (if available) + if file["annotation"] is not None: + assert isinstance(file["annotation"], Annotation) + + def test_ami_mdm_supervisions_multi_channel(self, ami_mdm_lhotse_protocol): + """Test that MDM supervisions include multi-channel data""" + files = list(ami_mdm_lhotse_protocol.train()) + + if len(files) > 0: + file = files[0] + supervisions = file["supervisions"] + # All supervisions should have channel information (channel is a list in MDM) + for sup in supervisions: + assert hasattr(sup, "channel") + assert isinstance(sup.channel, list) + assert len(sup.channel) > 0 + + def test_ami_ihm_annotation_content(self, ami_ihm_lhotse_protocol): + """Test that annotations contain actual speaker diarization content""" + files = list(ami_ihm_lhotse_protocol.train()) + + if len(files) > 0: + file = files[0] + annotation = file["annotation"] + + # Annotation may be None if there are no supervisions - only test if available + if annotation is not None: + # Annotation should have segments + assert len(annotation) > 0, "Annotation is empty" + + # Each segment should have a speaker label + for segment, track, label in annotation.itertracks(yield_label=True): + assert segment.start < segment.end + assert label is not None # Speaker label should exist + + def test_ami_ihm_annotated_regions(self, ami_ihm_lhotse_protocol): + """Test that annotated regions (timeline) are correct""" + files = list(ami_ihm_lhotse_protocol.train()) + + if len(files) > 0: + file = files[0] + annotated = file["annotated"] + + # Annotated should be a timeline + assert isinstance(annotated, Timeline) + + # Should have segments + assert len(annotated) > 0, "Annotated timeline is empty" + + # All segments should have valid time ranges + for segment in annotated: + assert segment.start >= 0 + assert segment.start < segment.end + + +class TestLhotseAMIPreprocessors: + """Test using preprocessors with real AMI data""" + + def test_ami_ihm_with_audio_path_preprocessor(self, ami_ihm_lhotse_protocol): + """Test using a preprocessor to add audio paths""" + # Create protocol with preprocessors + protocol = LhotseProtocol( + recording_set=ami_ihm_lhotse_protocol.recording_set, + supervision_set=ami_ihm_lhotse_protocol.supervision_set, + subset_split=ami_ihm_lhotse_protocol.subset_split, + preprocessors={ + # Return audio path from recording + "audio_path": lambda file: str(file["recording"].sources[0].source) + if file["recording"].sources + else None, + }, + ) + + files = list(protocol.train()) + if len(files) > 0: + file = files[0] + assert "audio_path" in file + assert file["audio_path"] is not None + + def test_ami_ihm_with_duration_preprocessor(self, ami_ihm_lhotse_protocol): + """Test using a preprocessor to add duration""" + protocol = LhotseProtocol( + recording_set=ami_ihm_lhotse_protocol.recording_set, + supervision_set=ami_ihm_lhotse_protocol.supervision_set, + subset_split=ami_ihm_lhotse_protocol.subset_split, + preprocessors={ + "duration": lambda file: file["recording"].duration, + }, + ) + + files = list(protocol.train()) + if len(files) > 0: + file = files[0] + assert "duration" in file + assert file["duration"] > 0