{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":129271,"databundleVersionId":15506307},{"sourceType":"datasetVersion","sourceId":14920384,"datasetId":9546899,"databundleVersionId":15786936}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Install Pyannote and Restart","metadata":{}},{"cell_type":"code","source":"!pip install pyannote-audio","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!apt-get install -y ffmpeg libavcodec-extra","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T14:11:56.101038Z","iopub.execute_input":"2026-02-22T14:11:56.101556Z","iopub.status.idle":"2026-02-22T14:11:58.671121Z","shell.execute_reply.started":"2026-02-22T14:11:56.101508Z","shell.execute_reply":"2026-02-22T14:11:58.670423Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport json\nimport torch\nimport time\nimport re\nimport pandas as pd\nimport soundfile as sf\nfrom multiprocessing import Process, Queue, Lock, Manager\nfrom pathlib import Path\nfrom tqdm import tqdm\nfrom pyannote.audio import Pipeline, Model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T14:09:01.74011Z","iopub.execute_input":"2026-02-22T14:09:01.740453Z","iopub.status.idle":"2026-02-22T14:09:21.589792Z","shell.execute_reply.started":"2026-02-22T14:09:01.740424Z","shell.execute_reply":"2026-02-22T14:09:21.589143Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ================= CONFIG =================\nAUDIO_DIR = \"/kaggle/input/competitions/dl-sprint-4-0-bengali-speaker-diarization-challenge/diarization/diarization/test/audio\"\nOUT_DIR = \"/kaggle/working\"\nHF_TOKEN = \"\" # place your hf token \nSEGMENTATION_MODEL = \"/kaggle/input/datasets/ahmfuad/villagers-segmentation-checkpoint/checkpoint20.ckpt\"\n\nTHRESHOLD = 0.704\nMETHOD = \"centroid\"\n\n# Post-processing Parameters\nGAP_MERGE_TH = 3.80\nMIN_INTER_SPEAKER_GAP = 0.17\nMIN_SEG_DUR = 0.75\nMIN_TOTAL_SPK_DUR = 9.0\n\nos.makedirs(OUT_DIR, exist_ok=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T14:09:45.65521Z","iopub.execute_input":"2026-02-22T14:09:45.656001Z","iopub.status.idle":"2026-02-22T14:09:45.660431Z","shell.execute_reply.started":"2026-02-22T14:09:45.655967Z","shell.execute_reply":"2026-02-22T14:09:45.65975Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ================= PREPROCESS =================\ndef preprocess(path):\n    audio, sr = sf.read(path)\n    duration = len(audio) / sr\n    if audio.ndim > 1:\n        audio = audio.mean(axis=1)\n    return torch.tensor(audio, dtype=torch.float32).unsqueeze(0), sr, duration\n\n# ================= POST-PROCESSING LOGIC =================\ndef apply_post_processing(diar_str):\n    try:\n        segments = json.loads(diar_str)\n    except:\n        return \"[]\"\n    if not segments: return \"[]\"\n\n    # 1. Serial Renaming\n    segments.sort(key=lambda x: x['start_time'])\n    speaker_map, next_num = {}, 0\n    renamed = []\n    for seg in segments:\n        orig = seg['speaker_id']\n        if orig not in speaker_map:\n            speaker_map[orig] = f\"SPEAKER_{next_num}\"\n            next_num += 1\n        new_seg = seg.copy()\n        new_seg['speaker_id'] = speaker_map[orig]\n        renamed.append(new_seg)\n\n    # 2. Overlap Resolution & Inter-Speaker Gap\n    def get_id_num(s_id):\n        m = re.search(r'(\\d+)', str(s_id))\n        return int(m.group(1)) if m else 999\n    \n    renamed.sort(key=lambda x: (x['start_time'], get_id_num(x['speaker_id'])))\n    \n    flat, current_end, last_spk = [], 0.0, None\n    for seg in renamed:\n        start, end, spk = seg['start_time'], seg['end_time'], seg['speaker_id']\n        if start < current_end: start = current_end\n        if last_spk is not None and spk != last_spk:\n            req_start = current_end + MIN_INTER_SPEAKER_GAP\n            if start < req_start: start = req_start\n        \n        if end > start:\n            flat.append({\"start_time\": float(start), \"end_time\": float(end), \"speaker_id\": spk})\n            current_end, last_spk = end, spk\n\n    # 3. Merge Same-Speaker Gaps\n    merged = []\n    if flat:\n        curr = flat[0]\n        for nxt in flat[1:]:\n            if nxt['speaker_id'] == curr['speaker_id'] and (nxt['start_time'] - curr['end_time']) < GAP_MERGE_TH:\n                curr['end_time'] = max(curr['end_time'], nxt['end_time'])\n            else:\n                merged.append(curr)\n                curr = nxt\n        merged.append(curr)\n\n    # 4. Filter by Durations\n    final = []\n    spk_totals = {}\n    # First pass for segment duration\n    dur_filtered = [s for s in merged if (s['end_time'] - s['start_time']) >= MIN_SEG_DUR]\n    # Second pass for total speaker duration\n    for s in dur_filtered:\n        spk_totals[s['speaker_id']] = spk_totals.get(s['speaker_id'], 0.0) + (s['end_time'] - s['start_time'])\n    \n    for s in dur_filtered:\n        if spk_totals.get(s['speaker_id'], 0.0) >= MIN_TOTAL_SPK_DUR:\n            s['start_time'], s['end_time'] = round(s['start_time'], 5), round(s['end_time'], 5)\n            final.append(s)\n\n    return json.dumps(final, separators=(\",\", \":\"))\n\n# ================= WORKER PROCESS =================\ndef diarization_worker(gpu_id, file_queue, result_list, duration_list, pbar, pbar_lock):\n    os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(gpu_id)\n    device = torch.device(\"cuda\")\n\n    pipeline = Pipeline.from_pretrained(\"pyannote/speaker-diarization-community-1\", token=HF_TOKEN)\n    seg_model = Model.from_pretrained(SEGMENTATION_MODEL)\n    pipeline.segmentation = seg_model\n    pipeline.to(device)\n    seg_model.min_duration_on = 0.0\n    seg_model.min_duration_off = 0.0\n    \n    while True:\n        try:\n            path = file_queue.get_nowait()\n        except:\n            break \n\n        file_id = os.path.splitext(os.path.basename(path))[0]\n        try:\n            wav, sr, duration = preprocess(path)\n            duration_list.append(duration)\n            \n            out = pipeline({\"waveform\": wav.to(device), \"sample_rate\": sr}, min_speakers=1, max_speakers=20)\n            diar = out.speaker_diarization if hasattr(out, \"speaker_diarization\") else out\n            \n            segs = [{\"start_time\": round(seg.start, 2), \"end_time\": round(seg.end, 2), \"speaker_id\": label} \n                    for seg, _, label in diar.itertracks(yield_label=True)]\n            \n            # Apply strict gap processing immediately or later; here we store raw for the summary stats\n            raw_json = json.dumps(segs, separators=(\",\", \":\"))\n            processed_json = apply_post_processing(raw_json)\n\n            result_list.append({\"filename\": file_id, \"diarization\": processed_json})\n        except Exception as e:\n            print(f\"\\nError on GPU {gpu_id} with {file_id}: {e}\")\n\n        with pbar_lock:\n            pbar.update(1)\n\n# ================= MAIN EXECUTION =================\nif __name__ == \"__main__\":\n    files = sorted(glob.glob(os.path.join(AUDIO_DIR, \"*.wav\")))\n    manager = Manager()\n    file_queue, result_list, duration_list = Queue(), manager.list(), manager.list()\n    for f in files: file_queue.put(f)\n    \n    pbar, pbar_lock = tqdm(total=len(files), desc=\"Diarizing (2x T4)\"), Lock()\n    start_time = time.time()\n\n    processes = [Process(target=diarization_worker, args=(i, file_queue, result_list, duration_list, pbar, pbar_lock)) for i in range(2)]\n    for p in processes: p.start()\n    for p in processes: p.join()\n\n    total_inference_time = time.time() - start_time\n    pbar.close()\n\n    # Save Final CSV\n    final_df = pd.DataFrame(list(result_list)).sort_values(\"filename\")\n    outname = f\"submission_th_{THRESHOLD}.csv\"\n    final_df.to_csv(os.path.join(OUT_DIR, outname), index=False)\n    \n    # Stats\n    total_audio = sum(duration_list)\n    print(f\"\\n{'='*30}\\nRTF: {total_inference_time/total_audio:.4f}\" if total_audio > 0 else \"\\nDone.\")\n    print(f\"File saved: {outname}\\n{'='*30}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-22T14:16:04.401457Z","iopub.execute_input":"2026-02-22T14:16:04.402355Z","iopub.status.idle":"2026-02-22T14:33:25.749284Z","shell.execute_reply.started":"2026-02-22T14:16:04.402302Z","shell.execute_reply":"2026-02-22T14:33:25.748325Z"}},"outputs":[],"execution_count":null}]}