diart

repository·main·Indexed 24 days ago

https://github.com/juanmc2005/diart

A Python framework for building AI-powered real-time audio applications, specifically focusing on speaker diarization. It supports streaming audio from microphones or files, integration with pyannote.audio models from Hugging Face, and custom model integration via ONNX or loader functions. Features include a CLI for streaming and hyper-parameter tuning via Optuna, WebSocket support for network-based diarization, and benchmarking tools for measuring real-time latency.

Tokens
9.7K
Snippets
11
Records
55
Agent score
84%

What's inside diart

  1. Tune pipeline hyper-parameters

    main

    Diart uses an optimizer based on optuna to tune pipeline hyper-parameters.

    From the command line

    diart.tune /wav/dir --reference /rttm/dir --output /output/dir

    From Python

    from diart.optim import Optimizer
    
    optimizer = Optimizer("/wav/dir", "/rttm/dir", "/output/dir")
    optimizer(num_iter=100)

    This writes results to an SQLite database in the specified output directory.

    Distributed Tuning

    To run multiple optimization processes in parallel, use a shared DBMS (e.g., MySQL or PostgreSQL) and pass the --storage flag in CLI or the study object in Python.

    CLI Example:

    diart.tune /wav/dir --reference /rttm/dir --storage mysql://root@localhost/example

    Python Example:

    from diart.optim import Optimizer
    from optuna.samplers import TPESampler
    import optuna
    
    db = "mysql://root@localhost/example"
    study = optuna.load_study("example", db, TPESampler())
    optimizer = Optimizer("/wav/dir", "/rttm/dir", study)
    optimizer(num_iter=100)
  2. Install diart

    main

    System Dependencies

    Before installing the package, ensure your system has the following dependencies installed:

    • ffmpeg < 4.4
    • portaudio == 19.6.X
    • libsndfile >= 1.2.2

    Alternatively, use the provided environment.yml to create a pre-configured conda environment:

    conda env create -f diart/environment.yml
    conda activate diart

    Package Installation

    Install the package via pip:

    pip install diart
  3. Serve and use Diart via WebSockets

    main

    Diart supports the WebSocket protocol to serve diarization pipelines over a network. You can run a server from the command line and connect a client, or implement a custom server in Python using WebSocketAudioSource.

    Important: Ensure the client uses the same step and sample_rate as the server by using the --step and -sr flags.

  4. How SpeakerMap and MappingMatrixObjective work together

    main

    A SpeakerMap represents the relationship between two sets of speakers using a mapping_matrix and an objective.

    The Mental Model

    1. The Matrix: A 2D array where matrix[i, j] represents the strength or cost of mapping source speaker i to target speaker j.
    2. The Objective: Since the matrix can represent either 'scores' (higher is better) or 'costs' (lower is better), the MappingMatrixObjective defines how to interpret the values:
      • MaximizationObjective: Used when higher values in the matrix indicate a better match (e.g., correlation). It uses np.max to find best values.
      • MinimizationObjective: Used when lower values indicate a better match (e.g., MSE, distance, or cost). It uses np.min to find best values.
    3. Optimal Assignment: The SpeakerMap uses the optimal_assignments method (via the objective) to solve the assignment problem (typically using scipy.optimize.linear_sum_assignment) to find the best one-to-one mapping between speakers.
  5. Audio source formats for the diart client

    main

    When using the diart client, the source argument accepts two main types of inputs:

    1. File Path: A string representing the path to an audio file (e.g., path/to/audio.wav). This uses FileAudioSource internally.
    2. Microphone:
      • 'microphone': Uses the default microphone device.
      • 'microphone:<DEVICE_ID>': Uses a specific microphone device identified by its integer ID (e.g., microphone:1). This uses MicrophoneAudioSource internally.
  6. Aggregation strategies for combining streaming results

    main

    Diart provides several strategies to handle overlapping buffers in a streaming context. These strategies are implemented via the AggregationStrategy interface and are typically used within a DelayedAggregation block.

    Available Strategies

    Strategy NameClassDescription
    "mean"AverageStrategyComputes a simple arithmetic mean of all overlapping buffer regions.
    "hamming"HammingWeightedAverageStrategyComputes a weighted average where weights are determined by a Hamming window aligned to each buffer. This helps smooth transitions.
    "first"FirstOnlyStrategyDoes not perform mathematical aggregation; it simply returns the first available buffer region in the list.

    Cropping Modes

    All strategies support cropping_mode to handle how buffers are sliced to fit the target Segment. These modes are compatible with pyannote.core:

    • "strict"
    • "loose" (default)
    • "center"
  7. How AudioSource works and how to consume its stream

    main

    An AudioSource is an abstract base class representing a source of audio. Every AudioSource provides a stream property, which is an rx.subject.Subject.

    To use an audio source, you typically:

    1. Instantiate the specific source (e.g., FileAudioSource).
    2. Subscribe to its stream to receive audio chunks (as NumPy arrays).
    3. Call .read() to start the data flow.
    4. Call .close() to stop the source.

    All sources emit audio chunks through the stream subject using on_next(waveform).

  8. How LazyModel handles model loading

    main

    Both SegmentationModel and EmbeddingModel inherit from LazyModel. This means the actual model weights are not loaded into memory until the first time the model is called or explicitly loaded.

    To manage memory and device placement, you can use:

    • .load(): Manually triggers the loading of the model into memory.
    • .to(device): Loads the model (if not already loaded) and moves it to the specified torch.device.
    • .eval(): Loads the model and sets it to evaluation mode (if it is a nn.Module).
    • .is_in_memory(): Checks if the model has been loaded yet.
  9. Implement a custom PipelineConfig

    main

    To define the operational parameters of a pipeline, you must implement the PipelineConfig abstract base class. A configuration must provide the following properties:

    • duration: The duration of an input audio chunk in seconds.
    • step: The step between two consecutive input audio chunks in seconds.
    • latency: The algorithmic latency in seconds. At time t of the audio stream, the pipeline outputs predictions for time t - latency.
    • sample_rate: The sample rate of the input audio stream in Hz.
  10. Build custom pipelines using blocks and RxPY

    main

    For advanced usage, you can combine building blocks from the diart.blocks module using RxPY operators. The blocks module is independent of the streaming engine and can be used separately.

    Example: Obtain overlap-aware speaker embeddings from a microphone stream

    import rx.operators as ops
    import diart.operators as dops
    from diart.sources import MicrophoneAudioSource, FileAudioSource
    from diart.blocks import SpeakerSegmentation, OverlapAwareSpeakerEmbedding
    
    segmentation = SpeakerSegmentation.from_pretrained("pyannote/segmentation")
    embedding = OverlapAwareSpeakerEmbedding.from_pretrained("pyannote/embedding")
    
    source = MicrophoneAudioSource()
    # To take input from file:
    # source = FileAudioSource("<filename>", sample_rate=16000)
    
    # Make sure the models have been trained with this sample rate
    print(source.sample_rate)
    
    stream = source.stream.pipe(
        # Reformat stream to 5s duration and 500ms shift
        dops.rearrange_audio_stream(sample_rate=source.sample_rate),
        ops.map(lambda wav: (wav, segmentation(wav))),
        ops.starmap(embedding)
    ).subscribe(on_next=lambda emb: print(emb.shape))
    
    source.read()

    Output shape is (batch_size, num_speakers, embedding_dim) (e.g., torch.Size([1, 3, 512])).

    import rx.operators as ops
    import diart.operators as dops
    from diart.sources import MicrophoneAudioSource, FileAudioSource
    from diart.blocks import SpeakerSegmentation, OverlapAwareSpeakerEmbedding
    
    segmentation = SpeakerSegmentation.from_pretrained("pyannote/segmentation")
    embedding = OverlapAwareSpeakerEmbedding.from_pretrained("pyannote/embedding")
    
    source = MicrophoneAudioSource()
    
    stream = source.stream.pipe(
        dops.rearrange_audio_stream(sample_rate=source.sample_rate),
        ops.map(lambda wav: (wav, segmentation(wav))),
        ops.starmap(embedding)
    ).subscribe(on_next=lambda emb: print(emb.shape))
    
    source.read()