generate_diarization_datasets.py

import argparse
import random
from collections import defaultdict
from pathlib import Path
from typing import List, Tuple, Dict

import numpy as np
import soundfile as sf
from scipy.signal import resample_poly


SAMPLE_RATE = 16000

N_SET_A          = 10
N_PER_OVERLAP    = 10
N_PER_SNR        = 10

OVERLAP_RATIOS   = [0.10, 0.20, 0.30]
SNR_LEVELS_DB    = [5, 10, 15]

TURNS_PER_FILE       = (4, 6)
INTER_TURN_SILENCE   = (0.2, 0.8)
LEAD_IN_SILENCE      = 0.5

MIN_UTT_SEC = 3.0
MAX_UTT_SEC = 8.0

STATIONARY_NOISE_PREFIXES = (
    "airconditioner",
    "airportannouncements",
    "copymachine",
    "vacuumcleaner",
    "washer",
)


def parse_speakers_file(speakers_file: Path) -> Dict[str, str]:
    gender = {}
    with open(speakers_file, "r", encoding="utf-8") as f:
        for line in f:
            if line.startswith(";") or not line.strip():
                continue
            parts = [p.strip() for p in line.split("|")]
            if len(parts) >= 2:
                gender[parts[0]] = parts[1].upper()
    return gender


def index_librispeech(root: Path, gender_map: Dict[str, str]) -> Dict[str, List[Path]]:
    by_speaker: Dict[str, List[Path]] = defaultdict(list)
    for flac in root.rglob("*.flac"):
        spk_id = flac.stem.split("-")[0]
        if spk_id in gender_map:
            by_speaker[spk_id].append(flac)
    return dict(by_speaker)


def load_utterance(path: Path) -> np.ndarray:
    audio, sr = sf.read(str(path), dtype="float32", always_2d=False)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    if sr != SAMPLE_RATE:
        from math import gcd
        g = gcd(sr, SAMPLE_RATE)
        audio = resample_poly(audio, SAMPLE_RATE // g, sr // g).astype(np.float32)
    return audio


def pick_utterance_for_speaker(
    speaker_utts: List[Path],
    rng: random.Random,
    max_attempts: int = 20,
) -> np.ndarray:
    pool = speaker_utts.copy()
    rng.shuffle(pool)
    for path in pool[:max_attempts]:
        audio = load_utterance(path)
        dur = len(audio) / SAMPLE_RATE
        if MIN_UTT_SEC <= dur <= MAX_UTT_SEC:
            return audio
    candidates = [(abs(len(load_utterance(p)) / SAMPLE_RATE -
                       (MIN_UTT_SEC + MAX_UTT_SEC) / 2), p) for p in pool[:max_attempts]]
    candidates.sort()
    return load_utterance(candidates[0][1])


def index_msnsd_noise(noise_dir: Path) -> List[Path]:
    keep = []
    for wav in noise_dir.rglob("*.wav"):
        name = wav.stem.lower().replace("_", "").replace("-", "")
        if any(name.startswith(p) for p in STATIONARY_NOISE_PREFIXES):
            keep.append(wav)
    if not keep:
        raise RuntimeError(
            f"No stationary noise files found in {noise_dir}. "
            f"Expected files matching one of: {STATIONARY_NOISE_PREFIXES}"
        )
    return keep


def place_segment(buffer: np.ndarray, segment: np.ndarray, start_sec: float) -> np.ndarray:
    start_sample = int(round(start_sec * SAMPLE_RATE))
    end_sample = start_sample + len(segment)
    if end_sample > len(buffer):
        buffer = np.pad(buffer, (0, end_sample - len(buffer)))
    buffer[start_sample:end_sample] += segment
    return buffer


def normalize_peak(x: np.ndarray, target_peak: float = 0.9) -> np.ndarray:
    peak = np.max(np.abs(x))
    if peak < 1e-9:
        return x
    return x * (target_peak / peak)


def mix_at_snr(speech: np.ndarray, noise: np.ndarray, target_snr_db: float) -> np.ndarray:
    if len(noise) < len(speech):
        repeats = int(np.ceil(len(speech) / len(noise)))
        noise = np.tile(noise, repeats)
    noise = noise[: len(speech)].astype(np.float32)

    speech_power = float(np.mean(speech ** 2))
    noise_power  = float(np.mean(noise  ** 2))
    if noise_power < 1e-12 or speech_power < 1e-12:
        return speech.copy()

    target_noise_power = speech_power / (10.0 ** (target_snr_db / 10.0))
    scale = float(np.sqrt(target_noise_power / noise_power))
    return speech + scale * noise


Timeline = List[Tuple[float, float, str]]


def build_alternating_conversation(
    utterances_A: List[np.ndarray],
    utterances_B: List[np.ndarray],
    rng: random.Random,
    leading_silence: float = 0.0,
) -> Tuple[np.ndarray, Timeline]:
    buf = np.zeros(int(leading_silence * SAMPLE_RATE), dtype=np.float32)
    cursor = leading_silence
    timeline: Timeline = []

    turns = list(zip(utterances_A, utterances_B))
    for utt_a, utt_b in turns:
        dur_a = len(utt_a) / SAMPLE_RATE
        start_a, end_a = cursor, cursor + dur_a
        buf = place_segment(buf, utt_a, start_a)
        timeline.append((start_a, end_a, "A"))
        cursor = end_a + rng.uniform(*INTER_TURN_SILENCE)

        dur_b = len(utt_b) / SAMPLE_RATE
        start_b, end_b = cursor, cursor + dur_b
        buf = place_segment(buf, utt_b, start_b)
        timeline.append((start_b, end_b, "B"))
        cursor = end_b + rng.uniform(*INTER_TURN_SILENCE)

    return buf, timeline


def build_overlapping_conversation(
    utterances_A: List[np.ndarray],
    utterances_B: List[np.ndarray],
    overlap_ratio: float,
    rng: random.Random,
) -> Tuple[np.ndarray, Timeline]:
    buf = np.zeros(1, dtype=np.float32)
    cursor = 0.0
    timeline: Timeline = []

    for utt_a, utt_b in zip(utterances_A, utterances_B):
        dur_a = len(utt_a) / SAMPLE_RATE
        dur_b = len(utt_b) / SAMPLE_RATE

        start_a = cursor
        end_a   = start_a + dur_a

        overlap_dur = min(dur_a, dur_b) * overlap_ratio
        start_b = end_a - overlap_dur
        end_b   = start_b + dur_b

        buf = place_segment(buf, utt_a, start_a)
        buf = place_segment(buf, utt_b, start_b)

        timeline.append((start_a, end_a, "A"))
        timeline.append((start_b, end_b, "B"))

        cursor = end_b + rng.uniform(*INTER_TURN_SILENCE)

    return buf, timeline


def write_rttm(timeline: Timeline, file_id: str, out_path: Path) -> None:
    with open(out_path, "w", encoding="utf-8") as f:
        for start, end, spk in sorted(timeline, key=lambda t: t[0]):
            dur = end - start
            if dur <= 0:
                continue
            f.write(
                f"SPEAKER {file_id} 1 {start:.3f} {dur:.3f} "
                f"<NA> <NA> {spk} <NA> <NA>\n"
            )


def write_uem(file_id: str, total_duration: float, out_path: Path) -> None:
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(f"{file_id} 1 0.000 {total_duration:.3f}\n")


def sample_speaker_pairs(
    speakers_by_gender: Dict[str, List[str]],
    n_pairs: int,
    rng: random.Random,
) -> List[Tuple[str, str]]:
    males   = list(speakers_by_gender.get("M", []))
    females = list(speakers_by_gender.get("F", []))
    pairs = []
    n_cross = n_pairs // 2
    n_same  = n_pairs - n_cross

    for _ in range(n_cross):
        if males and females:
            pairs.append((rng.choice(males), rng.choice(females)))

    for i in range(n_same):
        pool = males if (i % 2 == 0 and len(males) >= 2) else females
        if len(pool) >= 2:
            a, b = rng.sample(pool, 2)
            pairs.append((a, b))

    rng.shuffle(pairs)
    while len(pairs) < n_pairs and males and females:
        pairs.append((rng.choice(males), rng.choice(females)))
    return pairs[:n_pairs]


def save_file(
    audio: np.ndarray,
    timeline: Timeline,
    file_id: str,
    out_dir: Path,
) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    audio = normalize_peak(audio, target_peak=0.9)
    sf.write(out_dir / f"{file_id}.wav", audio, SAMPLE_RATE, subtype="PCM_16")
    write_rttm(timeline, file_id, out_dir / f"{file_id}.rttm")
    total_dur = len(audio) / SAMPLE_RATE
    write_uem(file_id, total_dur, out_dir / f"{file_id}.uem")


def build_pair_utterances(
    spk_A: str,
    spk_B: str,
    utts_by_speaker: Dict[str, List[Path]],
    n_turns: int,
    rng: random.Random,
) -> Tuple[List[np.ndarray], List[np.ndarray]]:
    utts_A = [pick_utterance_for_speaker(utts_by_speaker[spk_A], rng) for _ in range(n_turns)]
    utts_B = [pick_utterance_for_speaker(utts_by_speaker[spk_B], rng) for _ in range(n_turns)]
    return utts_A, utts_B


def generate_set_a(
    out_dir: Path,
    utts_by_speaker: Dict[str, List[Path]],
    speakers_by_gender: Dict[str, List[str]],
    rng: random.Random,
) -> None:
    print(f"\n[Set A] Generating {N_SET_A} clean alternating files...")
    pairs = sample_speaker_pairs(speakers_by_gender, N_SET_A, rng)
    for i, (spk_A, spk_B) in enumerate(pairs):
        n_turns = rng.randint(*TURNS_PER_FILE)
        utts_A, utts_B = build_pair_utterances(spk_A, spk_B, utts_by_speaker, n_turns, rng)
        audio, tl = build_alternating_conversation(utts_A, utts_B, rng, leading_silence=0.0)
        file_id = f"setA_{i:03d}_clean_{spk_A}-{spk_B}"
        save_file(audio, tl, file_id, out_dir)
        print(f"  ✓ {file_id}.wav  ({len(audio) / SAMPLE_RATE:.1f}s, {n_turns} turns)")


def generate_set_b(
    out_dir: Path,
    utts_by_speaker: Dict[str, List[Path]],
    speakers_by_gender: Dict[str, List[str]],
    rng: random.Random,
) -> None:
    print(f"\n[Set B] Generating {N_PER_OVERLAP * len(OVERLAP_RATIOS)} overlapping files "
          f"({N_PER_OVERLAP} per overlap level)...")
    for ratio in OVERLAP_RATIOS:
        sub_dir = out_dir / f"overlap_{int(ratio * 100):02d}pct"
        pairs = sample_speaker_pairs(speakers_by_gender, N_PER_OVERLAP, rng)
        for i, (spk_A, spk_B) in enumerate(pairs):
            n_turns = rng.randint(*TURNS_PER_FILE)
            utts_A, utts_B = build_pair_utterances(spk_A, spk_B, utts_by_speaker, n_turns, rng)
            audio, tl = build_overlapping_conversation(utts_A, utts_B, ratio, rng)
            file_id = f"setB_{int(ratio * 100):02d}pct_{i:03d}_{spk_A}-{spk_B}"
            save_file(audio, tl, file_id, sub_dir)
        print(f"  ✓ {int(ratio * 100)}% overlap → {N_PER_OVERLAP} files in {sub_dir.name}/")


def generate_set_c(
    out_dir: Path,
    utts_by_speaker: Dict[str, List[Path]],
    speakers_by_gender: Dict[str, List[str]],
    noise_files: List[Path],
    rng: random.Random,
) -> None:
    print(f"\n[Set C] Generating {N_PER_SNR * len(SNR_LEVELS_DB)} noisy files "
          f"({N_PER_SNR} per SNR level)...")
    for snr_db in SNR_LEVELS_DB:
        sub_dir = out_dir / f"snr_{snr_db:02d}dB"
        pairs = sample_speaker_pairs(speakers_by_gender, N_PER_SNR, rng)
        for i, (spk_A, spk_B) in enumerate(pairs):
            n_turns = rng.randint(*TURNS_PER_FILE)
            utts_A, utts_B = build_pair_utterances(spk_A, spk_B, utts_by_speaker, n_turns, rng)

            clean_audio, tl = build_alternating_conversation(
                utts_A, utts_B, rng, leading_silence=LEAD_IN_SILENCE
            )

            noise_path = rng.choice(noise_files)
            noise = load_utterance(noise_path)
            noisy_audio = mix_at_snr(clean_audio, noise, target_snr_db=snr_db)

            file_id = (f"setC_{snr_db:02d}dB_{i:03d}_{spk_A}-{spk_B}"
                       f"_{noise_path.stem}")
            save_file(noisy_audio, tl, file_id, sub_dir)
        print(f"  ✓ {snr_db}dB SNR → {N_PER_SNR} files in {sub_dir.name}/")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--librispeech_root", type=Path, required=True,
                        help="Path to LibriSpeech/test-clean directory")
    parser.add_argument("--speakers_file",    type=Path, required=True,
                        help="Path to LibriSpeech/SPEAKERS.TXT")
    parser.add_argument("--msnsd_noise_dir",  type=Path, required=True,
                        help="Path to MS-SNSD/noise_test (or noise_train)")
    parser.add_argument("--output_dir",       type=Path, required=True,
                        help="Where to write the generated sets")
    parser.add_argument("--seed", type=int, default=42,
                        help="Random seed for full reproducibility")
    args = parser.parse_args()

    rng = random.Random(args.seed)
    np.random.seed(args.seed)

    print("Indexing LibriSpeech...")
    gender_map = parse_speakers_file(args.speakers_file)
    utts_by_speaker = index_librispeech(args.librispeech_root, gender_map)
    if not utts_by_speaker:
        raise RuntimeError(f"No FLAC files found under {args.librispeech_root}")
    print(f"  Found {len(utts_by_speaker)} speakers, "
          f"{sum(len(v) for v in utts_by_speaker.values())} utterances")

    speakers_by_gender: Dict[str, List[str]] = defaultdict(list)
    for spk, paths in utts_by_speaker.items():
        if len(paths) >= 2:
            speakers_by_gender[gender_map[spk]].append(spk)
    print(f"  Usable speakers: {len(speakers_by_gender.get('M', []))} M, "
          f"{len(speakers_by_gender.get('F', []))} F")

    print("\nIndexing MS-SNSD stationary noise files...")
    noise_files = index_msnsd_noise(args.msnsd_noise_dir)
    print(f"  Found {len(noise_files)} stationary noise files")
    for nf in noise_files[:5]:
        print(f"    · {nf.name}")
    if len(noise_files) > 5:
        print(f"    · ... ({len(noise_files) - 5} more)")

    out_root = args.output_dir
    generate_set_a(out_root / "setA_clean",    utts_by_speaker, speakers_by_gender, rng)
    generate_set_b(out_root / "setB_overlap",  utts_by_speaker, speakers_by_gender, rng)
    generate_set_c(out_root / "setC_noise",    utts_by_speaker, speakers_by_gender, noise_files, rng)

    print(f"\n✓ Done. All files written to {out_root}")
    print(f"  Set A: {N_SET_A} files")
    print(f"  Set B: {N_PER_OVERLAP * len(OVERLAP_RATIOS)} files "
          f"({N_PER_OVERLAP} per level × {len(OVERLAP_RATIOS)} levels)")
    print(f"  Set C: {N_PER_SNR * len(SNR_LEVELS_DB)} files "
          f"({N_PER_SNR} per level × {len(SNR_LEVELS_DB)} levels)")


if __name__ == "__main__":
    main()