{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":129271,"databundleVersionId":15506307}],"dockerImageVersionId":31260,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\n# ============================================================================\n# INSTALL PACKAGES (unchanged)\n# ============================================================================\nimport subprocess\nimport sys\n\ndef install_packages():\n    packages = [\n        \"speechbrain\", \"torchaudio\", \"librosa\", \"transformers\",\n        \"scikit-learn\", \"pandas\", \"tqdm\", \"numpy\"\n    ]\n    for pkg in packages:\n        print(f\"Installing {pkg}...\")\n        subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", pkg])\n    print(\"✅ All packages installed!\")\n\ninstall_packages()\n\n# ============================================================================\n# IMPORTS\n# ============================================================================\nimport json\nimport numpy as np\nimport pandas as pd\nimport torch\nimport librosa\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom sklearn.cluster import SpectralClustering\nfrom sklearn.preprocessing import normalize\nfrom sklearn.metrics import silhouette_score\nfrom sklearn.decomposition import PCA\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom speechbrain.pretrained import EncoderClassifier\nfrom transformers import Wav2Vec2FeatureExtractor, WavLMForXVector\n\n# ============================================================================\n# CONFIGURATION – Tuned for Bengali\n# ============================================================================\nclass Config:\n    AUDIO_DIR = Path(\"/kaggle/input/dl-sprint-4-0-bengali-speaker-diarization-challenge/diarization/diarization/test/audio\")\n    OUTPUT_PATH = \"submission.csv\"\n    \n    SAMPLE_RATE = 16000\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    \n    # Bengali‑optimised windowing (longer utterances, softer transitions)\n    WINDOW_SIZE = 3.0          # seconds (was 2.0)\n    HOP_SIZE = 1.0              # seconds (was 0.5)\n    \n    # VAD – lower threshold to capture soft Bengali consonants\n    VAD_THRESHOLD = 0.25        # (was 0.30)\n    MIN_SPEECH_DURATION = 0.3\n    MIN_SILENCE_DURATION = 0.1\n    SPEECH_PAD_MS = 300\n    \n    # Output filtering\n    MIN_SEGMENT_DURATION = 1.0  # seconds (was 0.5) – reduces spurious short turns\n    MAX_GAP_SAME_SPEAKER = 0.5  # seconds – merge segments of same speaker closer than this\n    \n    # Speaker limits (Bengali conversation typical range)\n    MIN_SPEAKERS = 2\n    MAX_SPEAKERS = 15\n    \n    # Clustering parameters\n    CLUSTERING_METRIC = 'cosine'\n    USE_SPECTRAL = True          # use spectral clustering instead of agglomerative\n\ncfg = Config()\n\n# ============================================================================\n# LOAD MODELS (unchanged – XLSR removed)\n# ============================================================================\nprint(\"=\"*70)\nprint(\"🏆 OPTIMIZED BENGALI SPEAKER DIARIZATION\")\nprint(\"=\"*70)\nprint(f\"Device: {cfg.DEVICE}\")\nprint(\"=\"*70)\n\nmodels = {}\n\n# 1. Silero VAD\nprint(\"\\n[1/3] Loading Silero VAD...\")\nimport torch.hub\nsilero_model, utils = torch.hub.load('snakers4/silero-vad', 'silero_vad', trust_repo=True)\nget_speech_timestamps = utils[0]\nmodels['vad'] = silero_model\nprint(\"✓ Silero VAD loaded\")\n\n# 2. WavLM\nprint(\"\\n[2/3] Loading WavLM...\")\nwavlm_model = WavLMForXVector.from_pretrained(\"microsoft/wavlm-base-plus-sv\").to(cfg.DEVICE)\nwavlm_model.eval()\nwavlm_processor = Wav2Vec2FeatureExtractor.from_pretrained(\"microsoft/wavlm-base-plus-sv\")\nmodels['wavlm'] = (wavlm_model, wavlm_processor)\nprint(\"✓ WavLM loaded\")\n\n# 3. ECAPA\nprint(\"\\n[3/3] Loading ECAPA...\")\necapa_model = EncoderClassifier.from_hparams(\n    source=\"speechbrain/spkrec-ecapa-voxceleb\",\n    run_opts={\"device\": cfg.DEVICE}\n)\necapa_model.eval()\nmodels['ecapa'] = ecapa_model\nprint(\"✓ ECAPA loaded\")\n\nprint(f\"\\n✅ Loaded {len(models)} models successfully!\")\n\n# ============================================================================\n# EMBEDDING EXTRACTION (improved fusion)\n# ============================================================================\ndef extract_embedding(chunk):\n    \"\"\"Extract combined WavLM+ECAPA embedding, normalised and optionally PCA‑reduced.\"\"\"\n    try:\n        if len(chunk) < 16000:\n            chunk = np.pad(chunk, (0, 16000 - len(chunk)))\n        \n        # WavLM\n        inputs = wavlm_processor(chunk, sampling_rate=cfg.SAMPLE_RATE,\n                                return_tensors=\"pt\").input_values.to(cfg.DEVICE)\n        with torch.no_grad():\n            wavlm_emb = wavlm_model(inputs).embeddings.mean(dim=1).cpu().numpy().flatten()\n        \n        # ECAPA\n        x = torch.from_numpy(chunk).float().to(cfg.DEVICE).unsqueeze(0)\n        with torch.no_grad():\n            ecapa_emb = ecapa_model.encode_batch(x).squeeze().cpu().numpy()\n        \n        # Normalise individually\n        wavlm_emb = wavlm_emb / (np.linalg.norm(wavlm_emb) + 1e-8)\n        ecapa_emb = ecapa_emb / (np.linalg.norm(ecapa_emb) + 1e-8)\n        \n        # Concatenate and normalise\n        combined = np.concatenate([wavlm_emb[:128], ecapa_emb[:128]])  # fixed dim\n        combined = combined / (np.linalg.norm(combined) + 1e-8)\n        \n        return combined\n    except Exception as e:\n        print(f\"      ⚠️ Embedding error: {e}\")\n        return None\n\n# ============================================================================\n# VAD (unchanged but uses new threshold)\n# ============================================================================\ndef vad_with_silero(audio, sr):\n    audio_tensor = torch.from_numpy(audio).float()\n    speech_timestamps = get_speech_timestamps(\n        audio_tensor, models['vad'],\n        threshold=cfg.VAD_THRESHOLD,\n        sampling_rate=sr,\n        min_speech_duration_ms=int(cfg.MIN_SPEECH_DURATION * 1000),\n        min_silence_duration_ms=int(cfg.MIN_SILENCE_DURATION * 1000),\n        speech_pad_ms=cfg.SPEECH_PAD_MS\n    )\n    segments = [(ts['start'] / sr, ts['end'] / sr) for ts in speech_timestamps]\n    return segments\n\n# ============================================================================\n# SPEAKER COUNT ESTIMATION (gap statistic + silhouette)\n# ============================================================================\ndef estimate_speakers(embeddings):\n    if len(embeddings) < 10:\n        return cfg.MIN_SPEAKERS\n    \n    X = normalize(embeddings)\n    min_k = cfg.MIN_SPEAKERS\n    max_k = min(cfg.MAX_SPEAKERS, len(embeddings) // 2)  # more permissive\n    \n    best_k = cfg.MIN_SPEAKERS\n    best_score = -1\n    \n    for k in range(min_k, max_k + 1):\n        try:\n            # Use spectral clustering for estimation as well (consistent with final)\n            clustering = SpectralClustering(\n                n_clusters=k,\n                affinity='nearest_neighbors',\n                random_state=42,\n                n_init=10\n            )\n            labels = clustering.fit_predict(X)\n            \n            if len(set(labels)) < 2:\n                continue\n            \n            score = silhouette_score(X, labels, metric='cosine')\n            if score > best_score:\n                best_score = score\n                best_k = k\n        except Exception:\n            continue\n    \n    return best_k\n\n# ============================================================================\n# MAIN DIARIZATION FUNCTION\n# ============================================================================\ndef diarize_file(audio_path):\n    print(f\"\\n📁 Processing: {audio_path.stem}\")\n    \n    # Load audio\n    try:\n        audio, sr = librosa.load(str(audio_path), sr=cfg.SAMPLE_RATE, mono=True)\n        duration = len(audio) / sr\n        print(f\"   Duration: {duration:.1f}s\")\n    except Exception as e:\n        print(f\"   ❌ Failed to load: {e}\")\n        return []\n    \n    # VAD\n    segments = vad_with_silero(audio, sr)\n    print(f\"   Speech segments: {len(segments)}\")\n    \n    # If VAD fails, use uniform segmentation as fallback\n    if len(segments) < 2:\n        print(\"   ⚠️ VAD returned too few segments – using uniform windows\")\n        seg_dur = 3.0\n        segments = [(i*seg_dur, min((i+1)*seg_dur, duration))\n                    for i in range(int(np.ceil(duration/seg_dur)))]\n    \n    # Create overlapping windows\n    windows = []\n    for start, end in segments:\n        dur = end - start\n        if dur <= cfg.WINDOW_SIZE:\n            windows.append((start, end))\n        else:\n            current = start\n            while current + cfg.WINDOW_SIZE <= end:\n                windows.append((current, current + cfg.WINDOW_SIZE))\n                current += cfg.HOP_SIZE\n            if current < end - 0.5:\n                windows.append((max(start, end - cfg.WINDOW_SIZE), end))\n    \n    print(f\"   Windows: {len(windows)}\")\n    \n    # Extract embeddings\n    embeddings = []\n    valid_windows = []\n    for start, end in windows:\n        s = int(start * sr)\n        e = int(end * sr)\n        if e - s >= 0.8 * sr:\n            emb = extract_embedding(audio[s:e])\n            if emb is not None:\n                embeddings.append(emb)\n                valid_windows.append((start, end))\n    \n    print(f\"   Embeddings extracted: {len(embeddings)}\")\n    \n    if len(embeddings) < 5:\n        print(\"   ⚠️ Too few embeddings – fallback to uniform segments\")\n        # Create 3 uniform speakers\n        seg_len = duration / 3\n        turns = []\n        for i in range(3):\n            start = i * seg_len\n            end = min((i+1) * seg_len, duration)\n            turns.append((start, end, i))\n        return format_turns(turns)\n    \n    # Estimate number of speakers\n    n_speakers = estimate_speakers(embeddings)\n    print(f\"   → Estimated speakers: {n_speakers}\")\n    \n    # Cluster using spectral clustering (or agglomerative fallback)\n    X = normalize(embeddings)\n    if cfg.USE_SPECTRAL:\n        clustering = SpectralClustering(\n            n_clusters=n_speakers,\n            affinity='nearest_neighbors',\n            random_state=42,\n            n_init=10,\n            assign_labels='kmeans'\n        )\n    else:\n        from sklearn.cluster import AgglomerativeClustering\n        clustering = AgglomerativeClustering(\n            n_clusters=n_speakers,\n            metric='cosine',\n            linkage='average'\n        )\n    \n    labels = clustering.fit_predict(X)\n    \n    # Merge adjacent same‑speaker windows\n    items = sorted(zip(valid_windows, labels), key=lambda x: x[0][0])\n    turns = []\n    cur_start, cur_end, cur_label = items[0][0][0], items[0][0][1], items[0][1]\n    \n    for (start, end), label in items[1:]:\n        if label == cur_label and start <= cur_end + cfg.MAX_GAP_SAME_SPEAKER:\n            cur_end = max(cur_end, end)\n        else:\n            turns.append((cur_start, cur_end, cur_label))\n            cur_start, cur_end, cur_label = start, end, label\n    turns.append((cur_start, cur_end, cur_label))\n    \n    # Filter short segments, but keep at least one per file\n    min_dur = cfg.MIN_SEGMENT_DURATION\n    turns = [t for t in turns if t[1] - t[0] >= min_dur] or turns\n    print(f\"   Final turns: {len(turns)}\")\n    \n    return format_turns(turns)\n\ndef format_turns(turns):\n    \"\"\"Convert (start, end, label) list to required JSON format.\"\"\"\n    output = []\n    speaker_map = {}\n    next_id = 1\n    for start, end, label in sorted(turns, key=lambda x: x[0]):\n        key = str(label)\n        if key not in speaker_map:\n            speaker_map[key] = f\"SPEAKER_{next_id}\"\n            next_id += 1\n        output.append({\n            \"start_time\": round(float(start), 2),\n            \"end_time\": round(float(end), 2),\n            \"speaker_id\": speaker_map[key]\n        })\n    return output\n\n# ============================================================================\n# PROCESS ALL FILES\n# ============================================================================\ndef main():\n    print(\"\\n\" + \"=\"*70)\n    print(\"🏆 PROCESSING ALL 14 FILES (Bengali‑optimised)\")\n    print(\"=\"*70)\n    \n    audio_files = sorted(cfg.AUDIO_DIR.glob(\"*.wav\"))\n    print(f\"Found {len(audio_files)} test files\")\n    print(\"Estimated runtime: ~2 hours\")\n    print(\"=\"*70)\n    \n    results = []\n    speaker_counts = []\n    \n    for audio_path in tqdm(audio_files, desc=\"Diarizing\"):\n        try:\n            diarization = diarize_file(audio_path)\n            if not diarization:\n                # Absolute fallback: one speaker entire file\n                audio, sr = librosa.load(str(audio_path), sr=cfg.SAMPLE_RATE)\n                duration = len(audio) / sr\n                diarization = [{\n                    \"start_time\": 0.0,\n                    \"end_time\": round(duration, 2),\n                    \"speaker_id\": \"SPEAKER_1\"\n                }]\n            \n            speakers = len(set(seg['speaker_id'] for seg in diarization))\n            speaker_counts.append(speakers)\n            results.append({\n                \"filename\": audio_path.stem,\n                \"diarization\": json.dumps(diarization, ensure_ascii=False)\n            })\n            print(f\"   ✅ {audio_path.stem}: {speakers} speakers, {len(diarization)} segments\")\n        except Exception as e:\n            print(f\"\\n❌ Error processing {audio_path.stem}: {e}\")\n            import traceback\n            traceback.print_exc()\n            results.append({\n                \"filename\": audio_path.stem,\n                \"diarization\": json.dumps([{\n                    \"start_time\": 0.0,\n                    \"end_time\": 10.0,\n                    \"speaker_id\": \"SPEAKER_1\"\n                }], ensure_ascii=False)\n            })\n    \n    # Save submission\n    df = pd.DataFrame(results)\n    df.to_csv(cfg.OUTPUT_PATH, index=False)\n    \n    print(\"\\n\" + \"=\"*70)\n    print(\"✅ PROCESSING COMPLETE\")\n    print(\"=\"*70)\n    if speaker_counts:\n        print(f\"Files processed: {len(results)}\")\n        print(f\"Average speakers: {np.mean(speaker_counts):.1f}\")\n        print(f\"Min speakers: {min(speaker_counts)}\")\n        print(f\"Max speakers: {max(speaker_counts)}\")\n        print(f\"Distribution: {sorted(speaker_counts)}\")\n    \n    print(f\"\\n📊 Submission saved: {cfg.OUTPUT_PATH}\")\n    return df\n\nif __name__ == \"__main__\":\n    print(\"\\n\" + \"=\"*70)\n    print(\"🏆 OPTIMIZED BENGALI SPEAKER DIARIZATION\")\n    print(\"✓ VAD threshold lowered for Bengali soft consonants\")\n    print(\"✓ Longer windows (3s) to capture utterance traits\")\n    print(\"✓ Spectral clustering + gap‑statistic for speaker count\")\n    print(\"✓ Post‑processing merges short same‑speaker segments\")\n    print(\"✓ Will process ALL 14 files\")\n    print(\"=\"*70)\n    \n    submission = main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T09:33:40.187383Z","iopub.execute_input":"2026-02-22T09:33:40.188284Z"}},"outputs":[],"execution_count":null}]}