treble_tsdk.scene.audio_dataset

Classes

AudioDataset

A parquet-backed table of audio references, indexable by id/int/slice/pl.Expr, and consumed by TrackGenerator to drive scene/track generation.
class treble_tsdk.scene.audio_dataset.AudioDataset

A parquet-backed table of audio references, indexable by id/int/slice/pl.Expr, and consumed by TrackGenerator to drive scene/track generation.

Three ways to build one, depending on where the audio lives:

__init__(parquet_url: str | Path | list[str | Path], audio_loader: Callable, schema_mapping: dict[str, str] | None = None, drop_columns: list[str] | None = None, metadata_cache_path: Path | None = None)

Initialize an audio dataset from one or more parquet shards.

Parameters:

  • parquet_url (str | Path | list[str | Path]) – Path/URL (or list of paths/URLs) to parquet shards containing metadata rows.

  • audio_loader (Callable) – Loader used to resolve and decode audio by sample identifier.

  • schema_mapping (dict[str, str] | None) – Mapping from logical column names (e.g. id, transcript, length_s) to the parquet’s actual column names.

  • drop_columns (list[str] | None) – Columns to drop from this dataset’s own metadata table only (typically a large embedded-audio-bytes column, to avoid holding it in memory for filtering/indexing). Doesn’t affect audio loading: audio_loader keeps its own separate, lazy handle on the same parquet.

  • metadata_cache_path (Path | None) – For a multi-shard parquet_url, avoids re-scanning and reassembling the full metadata table from every shard on each instantiation: the assembled result is written here on first use and read directly from it thereafter. Delete the file manually to force a rebuild (e.g. after the source shards change).

add_column(name: str, mapper: Callable, dtype: pl.DataType | None = None) → AudioDataset

Add a column to the dataset by applying a mapper function to each sample.

Parameters:

  • name (str) – Name of the new column.

  • mapper (Callable) – Function called with each AudioSampleReference that returns the value for that row; return type must be consistent across all rows.

  • dtype (pl.DataType | None) – Polars data type of the new column. If None, inferred from the first row’s return value.

Return AudioDataset:

This dataset after adding the column (mutates self.dataframe in place). AudioDataset itself has no write_parquet; persist the added column with dataset.dataframe.write_parquet(path).

enrich_with_spl(use_active_speech_level: bool = False) → None

Compute octave-band SPL values for dataset items and add them as columns in place: one list column of SPL values and one of the corresponding octave-band labels, per row.

Required before predicted-SNR estimation (see predict_snr()) on every AudioDataset feeding a track that participates in the estimate — and it must run before generating the tracks that will read this metadata (e.g. before this dataset is passed to TrackGenerator), not merely before predict_snr is called. Each generated AudioTrack/ RepeatedAudioTrack embeds a snapshot AudioSampleReference (with its own, independent metadata dict) at the moment it’s pulled from the dataset, so enriching afterward does not retroactively add spl to samples already captured inside tracks generated before the enrichment ran.

Slow for a large dataset: it decodes and processes every audio file in the dataset — a load_audio_signal() call plus an octave-band SPL pass per row, not just metadata. If you only need SPL for a subset of samples, narrow the dataset first (e.g. via filter_collection()) and call this on that smaller dataset. If the same dataset is reused across sessions, enrich once and persist the result (dataset.dataframe.write_parquet(path)) instead of re-running this every time.

Parameters:

use_active_speech_level (bool) – If True, restrict the SPL computation to samples flagged as active speech by the ITU-T P.56 detector, so leading/trailing silence doesn’t depress the measured SPL for speech datasets.

filter_collection(filter_expression: polars.Expr)

Filter the collection using a polars expression.

Parameters:

filter_expression (pl.Expr) – The polars expression to filter the collection by.

Return AudioDataset:

The filtered collection.

classmethod from_huggingface(repo_id: str, audio_loader_class: type, config: str | None = None, split: str | None = None, schema_mapping: dict[str, str] | None = None, drop_columns: list[str] | None = None, cache_directory: str | Path | None = None, audio_loader_kwargs: dict[str, object] | None = None, max_parquet_files: int | None = None, token: str | None = None) → AudioDataset

Create an AudioDataset from a HuggingFace dataset repository.

Discovers parquet files via the HuggingFace Hub API and loads them.

Parameters:

  • repo_id (str) – HuggingFace dataset repository ID (e.g. “openslr/librispeech_asr”).

  • audio_loader_class (type) – Audio loader class to use (e.g. LibriSpeechAudioLoader). An instance is created internally using the first discovered URL and cache_directory.

  • config (str | None) – Dataset configuration/subset to load (e.g. “clean”). If None, all configs.

  • split (str | None) – Dataset split to load (e.g. “train.100”). If None, all splits.

  • schema_mapping (dict | None) – Column name mapping (logical_name -> actual_column_name).

  • drop_columns (list | None) – Columns to drop from the parquet files.

  • cache_directory (str | Path | None) – Local directory for caching audio data and sample metadata (the latter under a metadata subfolder). If None, uses the default TSDK cache directory under an audio_datasets subfolder; pass None to keep the dataset fully portable for serialization and loading elsewhere.

  • audio_loader_kwargs (dict[str, object] | None) – Optional keyword arguments passed to audio_loader_class. If omitted, loaders implementing get_default_loader_kwargs (e.g. AudioSetAudioLoader) infer sensible defaults from schema_mapping automatically (id_column defaults to schema_mapping.get("id", "video_id")); pass explicitly to override.

  • max_parquet_files (int | None) – Maximum number of parquet files to load after config/split filtering. Useful for faster iteration on large datasets.

  • token (str | None) – HuggingFace auth token for private datasets.

Return AudioDataset:

Loaded collection built from discovered parquet shards.

static from_local_speech_directory(speech_dir: Path, parquet_path: Path, load_transcripts: bool = False, talker_identifier: str = 'speaker_id') → AudioDataset

Create a speech AudioDataset from a directory of per-talker WAV folders.

Unlike from_local_wav_directory() (suited to flat noise datasets), this produces a dataset carrying a talker identifier, a per-clip duration, and a transcript, usable with conversation rules and talker-aware scene generation just like LibriSpeechAudioLoader-backed datasets.

Directory structure:

speech_dir/
    talker_a/      # subfolder name acts as the talker ID
        001.wav
        001.txt    # transcript of 001.wav
        002.wav
        002.txt
    talker_b/
        001.wav
        001.txt
        ...

Each immediate subfolder is treated as a single talker; every *.wav inside it (recursively) is attributed to that talker. WAV files directly in speech_dir (outside any subfolder) are ignored.

Each *.wav should have a sibling *.txt file with the same stem (e.g. 001.wav -> 001.txt) holding its transcript. This populates a transcript column, consumed by AudioDataset and Scene Generation the same way as the embedded transcripts of HuggingFace speech loaders.

If the parquet file already exists, it is loaded instead of scanning the directory and rebuilding the manifest.

Parameters:

  • speech_dir (Path) – Directory containing one subfolder per talker.

  • parquet_path (Path) – Parquet manifest to write.

  • load_transcripts (bool) – If True, raise when a WAV has no sibling *.txt transcript. If False, log a warning and store an empty string.

  • talker_identifier (str) – Name of the column holding the talker ID (the subfolder name); pass this same value as talker_identifier to TrackGenerator. Defaults to "speaker_id".

Return AudioDataset:

Dataset backed by the written Parquet and AudioPathAudioLoader, carrying talker IDs, clip durations, and transcripts.

static from_local_wav_directory(wav_dir: Path, parquet_path: Path, search_subdir: bool = True) → AudioDataset

Create an AudioDataset by scanning a directory for *.wav files and writing a small Parquet manifest that AudioPathAudioLoader can consume.

Manifest columns:

  • id: each file’s path relative to wav_dir.

  • path: each WAV file’s path relative to the parquet file’s parent directory, so the dataset stays portable when wav_dir and the parquet are moved together.

If the parquet file already exists, it is loaded instead of scanning the directory and rebuilding the manifest.

Parameters:

  • wav_dir (Path) – Directory to scan for *.wav files.

  • parquet_path (Path) – Parquet manifest to write.

  • search_subdir (bool) – If True (default), include WAVs in subfolders via recursive search; if False, only WAVs directly in wav_dir.

Return AudioDataset:

Dataset backed by the written Parquet and AudioPathAudioLoader.

head(n: int = 10) → AudioDataset

Return a new collection with the first n rows.

Parameters:

n (int) – Number of rows to keep.

Return AudioDataset:

A new collection with only the first rows.

property dataframe: polars.DataFrame

Underlying materialized Polars DataFrame.

property ids

Return the list of sample identifiers.

Return list[str]:

IDs from the id column.

property required_columns: list[str]

Columns reserved for core dataset semantics and reference building.