segmentation_infer_html.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
segmentation_infer_smooth_segments.py

- Loads WhisperOddEven checkpoint
    /home/user/outs/segmentation_gemini_2p_medium_model_best.pt
  (override via CKPT env var).

- For each audio file in AUDIO_INPUT_DIR:
    * load, resample to 16 kHz mono
    * split into 30 s chunks
    * run segmentation
    * SMOOTH each track so that no segment (incl. background 0) is shorter than
      MIN_SEGMENT_SEC seconds
    * extract per-track segments (odd/even) and cut audio snippets
    * build a MERGED timeline that starts/ends segments whenever either track
      changes label, then smooth that merged timeline so that each merged
      segment is also at least MIN_SEGMENT_SEC long, merging short segments
      with neighbors using the rules described below.

- Writes a single HTML report with:
    * smoothed per-track heatmap
    * merged-timeline heatmap
    * tables of per-track segments (with audio players)
    * tables of merged segments (with audio players)

Merging rule for short merged segments:
- If a merged segment is shorter than MIN_SEGMENT_SEC, merge it with one of its
  immediate neighbors.
- Prefer the neighbor whose (odd_label, even_label) matches this segment best
  (majority vote over the two labels).
- If similarity is equal (or one neighbor is missing), merge with the neighbor
  that has the shorter duration. If still equal, merge with the left neighbor.
"""

from __future__ import annotations
import os
import io
import sys
import time
import math
import base64
import shutil
from pathlib import Path
from typing import List, Dict, Any, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# plotting
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# audio
import soundfile as sf
import librosa
from pydub import AudioSegment  # requires ffmpeg

from transformers import WhisperFeatureExtractor, WhisperModel

# =========================
# ========== CONFIG =======
# =========================

AUDIO_INPUT_DIR = Path(os.getenv("AUDIO_INPUT_DIR", "./infer-audio"))
OUT_DIR         = Path(os.getenv("OUT_DIR", "./outs_infer"))
CKPT_PATH       = Path(os.getenv("CKPT", "/home/user/outs/segmentation_gemini_medium_no_overlap_4epochs_model_best.pt"))
HF_MODEL_ID     = os.getenv("HF_MODEL_ID", "openai/whisper-small")

USE_LOCAL_MODELS    = bool(int(os.getenv("USE_LOCAL_MODELS", "0")))
MODELS_SNAPSHOT_DIR = Path(os.getenv("MODELS_SNAPSHOT_DIR", "")) if USE_LOCAL_MODELS else None
HF_HOME             = Path(os.getenv("HF_HOME", (OUT_DIR / ".hf")))
TRANSFORMERS_CACHE  = Path(os.getenv("TRANSFORMERS_CACHE", (OUT_DIR / ".hf" / "hub")))

MIXED_PRECISION = os.getenv("MIXED_PRECISION", "auto").lower()

# constants (must match training)
SAMPLE_RATE   = 16000
CLIP_SECONDS  = 30.0
NUM_FRAMES    = 1500
NUM_TRACKS    = 2
MAX_SEGMENTS  = 20

# --- MINIMUM SEGMENT LENGTH (seconds) for both per-track and merged segments ---
MIN_SEGMENT_SEC = float(os.getenv("MIN_SEGMENT_SEC", "1.0"))
MIN_SEGMENT_FRAMES = max(1, int(round(MIN_SEGMENT_SEC * NUM_FRAMES / CLIP_SECONDS)))

FFMPEG_AVAILABLE = shutil.which("ffmpeg") is not None
WARNED_NO_FFMPEG = False

# =========================
# ====== BASIC SETUP ======
# =========================

def setup_dirs():
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    (OUT_DIR / ".mplconfig").mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("MPLCONFIGDIR", str((OUT_DIR / ".mplconfig").resolve()))
    HF_HOME.mkdir(parents=True, exist_ok=True)
    os.environ.setdefault("HF_HOME", str(HF_HOME.resolve()))
    os.environ.setdefault("TRANSFORMERS_CACHE", str(TRANSFORMERS_CACHE.resolve()))
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True,max_split_size_mb:128")

def preferred_dtype():
    if MIXED_PRECISION == "bf16":
        return torch.bfloat16
    if MIXED_PRECISION == "fp16":
        return torch.float16
    if MIXED_PRECISION == "fp32":
        return torch.float32
    if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
        return torch.bfloat16
    return torch.float16 if torch.cuda.is_available() else torch.float32

def _model_resolved_name(model_id: str) -> Tuple[str, bool]:
    if USE_LOCAL_MODELS and MODELS_SNAPSHOT_DIR and MODELS_SNAPSHOT_DIR.is_dir():
        local_dirname = model_id.replace("/", "__")
        cand = MODELS_SNAPSHOT_DIR / local_dirname
        if cand.is_dir():
            return str(cand), True
    return model_id, False

# =========================
# ========= MODEL =========
# =========================

class WhisperOddEven(nn.Module):
    def __init__(self, base_id: str, freeze_encoder: bool = False):
        super().__init__()
        resolved, is_local = _model_resolved_name(base_id)
        self.whisper = WhisperModel.from_pretrained(resolved, local_files_only=is_local)

        # decoder unused
        for p in self.whisper.decoder.parameters():
            p.requires_grad = False

        for p in self.whisper.encoder.parameters():
            p.requires_grad = not freeze_encoder

        d_model = self.whisper.config.d_model
        hidden = max(256, d_model // 2)
        self.head = nn.Sequential(
            nn.Linear(d_model, hidden),
            nn.GELU(),
            nn.Linear(hidden, NUM_TRACKS * (MAX_SEGMENTS + 1)),
        )

    def forward(self, input_features: torch.FloatTensor):
        enc = self.whisper.encoder(input_features=input_features).last_hidden_state  # [B,1500,D]
        B, T, D = enc.shape
        logits = self.head(enc)  # [B,T,NUM_TRACKS*(C)]
        C = MAX_SEGMENTS + 1
        logits = logits.view(B, T, NUM_TRACKS, C).permute(0, 2, 1, 3).contiguous()
        return logits  # [B,2,1500,C]

# =========================
# ====== AUDIO UTILS ======
# =========================

def load_audio_mono_16k(path: Path) -> np.ndarray:
    wav, sr = librosa.load(str(path), sr=SAMPLE_RATE, mono=True)
    if wav.ndim > 1:
        wav = wav.mean(axis=0)
    return wav.astype(np.float32, copy=False)

def split_into_chunks(wav: np.ndarray, sr: int, clip_seconds: float):
    chunk_size = int(clip_seconds * sr)
    total = len(wav)
    if total == 0:
        return []
    n_chunks = math.ceil(total / chunk_size)
    chunks = []
    for i in range(n_chunks):
        start = i * chunk_size
        end = min(start + chunk_size, total)
        seg = wav[start:end]
        if len(seg) < chunk_size:
            seg = np.pad(seg, (0, chunk_size - len(seg)), mode="constant")
        chunks.append((i, start, seg.astype(np.float32, copy=False)))
    return chunks

def wav_chunk_to_audio_bytes(wav: np.ndarray, sr: int):
    """
    Try to export as MP3 (if ffmpeg is available). Otherwise fall back to WAV.
    Returns (audio_bytes, mime_type).
    """
    global WARNED_NO_FFMPEG

    buf_wav = io.BytesIO()
    sf.write(buf_wav, wav, sr, format="WAV")
    wav_bytes = buf_wav.getvalue()

    if not FFMPEG_AVAILABLE:
        if not WARNED_NO_FFMPEG:
            print("[audio] ffmpeg not found; embedding WAV instead of MP3.", flush=True)
            WARNED_NO_FFMPEG = True
        return wav_bytes, "audio/wav"

    try:
        buf_wav.seek(0)
        audio = AudioSegment.from_file(buf_wav, format="wav")
        out_buf = io.BytesIO()
        audio.export(out_buf, format="mp3", bitrate="128k")
        out_buf.seek(0)
        return out_buf.read(), "audio/mpeg"
    except Exception as e:
        if not WARNED_NO_FFMPEG:
            print(f"[audio] Failed to encode MP3, falling back to WAV: {e}", flush=True)
            WARNED_NO_FFMPEG = True
        return wav_bytes, "audio/wav"

# =========================
# ====== SEGMENT OPS ======
# =========================

def smooth_min_duration(ids: np.ndarray, min_frames: int, max_iter: int = 10) -> np.ndarray:
    """
    Enforce a minimum run length (in frames) for an ID sequence (1D).
    Shorter runs are reassigned to the longer of their neighbors, iteratively.
    """
    ids = ids.copy()
    n = len(ids)
    if n == 0:
        return ids

    for _ in range(max_iter):
        runs = []
        start = 0
        cur = ids[0]
        for i in range(1, n):
            if ids[i] != cur:
                runs.append((cur, start, i))
                start = i
                cur = ids[i]
        runs.append((cur, start, n))

        changed = False
        for ri, (label, s, e) in enumerate(runs):
            length = e - s
            if length >= min_frames:
                continue

            left = runs[ri - 1] if ri > 0 else None
            right = runs[ri + 1] if ri + 1 < len(runs) else None
            if left is None and right is None:
                continue

            if left is None:
                new_label = right[0]
            elif right is None:
                new_label = left[0]
            else:
                len_left = left[2] - left[1]
                len_right = right[2] - right[1]
                new_label = left[0] if len_left >= len_right else right[0]

            if new_label != label:
                ids[s:e] = new_label
                changed = True

        if not changed:
            break

    return ids

def extract_segments(ids: np.ndarray, include_bg: bool = False):
    """
    Return list of (label, frame_start, frame_end) runs.
    Optionally filter out background label 0.
    """
    n = len(ids)
    if n == 0:
        return []
    runs = []
    start = 0
    cur = ids[0]
    for i in range(1, n):
        if ids[i] != cur:
            runs.append((cur, start, i))
            start = i
            cur = ids[i]
    runs.append((cur, start, n))
    if not include_bg:
        runs = [(lab, s, e) for (lab, s, e) in runs if lab != 0]
    return runs

def frames_to_times(s: int, e: int):
    start_t = s / NUM_FRAMES * CLIP_SECONDS
    end_t = e / NUM_FRAMES * CLIP_SECONDS
    return start_t, end_t

def cut_wav(seg_wav: np.ndarray, start_t: float, end_t: float) -> np.ndarray:
    start_samp = int(round(start_t * SAMPLE_RATE))
    end_samp   = int(round(end_t * SAMPLE_RATE))
    start_samp = max(0, min(start_samp, len(seg_wav)))
    end_samp   = max(start_samp + 1, min(end_samp, len(seg_wav)))
    return seg_wav[start_samp:end_samp]

# =========================
# ==== MERGED TIMELINE ====
# =========================

def smooth_merged_segments(merged: List[Tuple[int,int,int,int]], min_frames: int) -> List[Tuple[int,int,int,int]]:
    """
    Enforce minimum length for merged segments.

    merged: list of (frame_start, frame_end, odd_label, even_label).
    If a segment has length < min_frames, we merge it with a neighbor:
      - If both neighbors exist, choose the one with higher similarity of
        (odd_label, even_label). Similarity is number of matching labels (0..2).
      - If similarity is equal, merge with the neighbor that has shorter
        duration (in frames). If still equal, merge with the left neighbor.
      - If only one neighbor exists, merge with that neighbor.

    Returns a new merged list.
    """
    if len(merged) <= 1:
        return merged

    merged = list(merged)

    def seg_len(seg):
        return seg[1] - seg[0]

    def sim(a, b):
        # a,b: (fs,fe, odd,even)
        score = 0
        if a[2] == b[2]:
            score += 1
        if a[3] == b[3]:
            score += 1
        return score

    changed = True
    while changed:
        changed = False
        n = len(merged)
        if n <= 1:
            break
        for i, seg in enumerate(merged):
            length = seg_len(seg)
            if length >= min_frames:
                continue

            left = merged[i - 1] if i > 0 else None
            right = merged[i + 1] if i + 1 < n else None

            if left is None and right is None:
                continue

            # Decide which neighbor to merge with
            if left is not None and right is not None:
                s_left = sim(seg, left)
                s_right = sim(seg, right)
                if s_left > s_right:
                    target = "left"
                elif s_right > s_left:
                    target = "right"
                else:
                    # similarity tie -> choose shorter neighbor
                    len_left = seg_len(left)
                    len_right = seg_len(right)
                    if len_left < len_right:
                        target = "left"
                    elif len_right < len_left:
                        target = "right"
                    else:
                        target = "left"  # full tie -> left
            elif left is not None:
                target = "left"
            else:
                target = "right"

            if target == "left":
                fs = left[0]
                fe = seg[1]
                odd_label = left[2]
                even_label = left[3]
                merged[i - 1] = (fs, fe, odd_label, even_label)
                del merged[i]
            else:
                fs = seg[0]
                fe = right[1]
                odd_label = right[2]
                even_label = right[3]
                merged[i + 1] = (fs, fe, odd_label, even_label)
                del merged[i]
            changed = True
            break  # restart scanning with new list

    return merged

def build_merged_segments(ids_odd: np.ndarray, ids_even: np.ndarray, min_frames: int):
    """
    Build merged segmentation from two tracks and then smooth merged segments.

    - boundaries are at 0, NUM_FRAMES, and every point where either track changes.
    - for each raw merged segment we set odd/even labels via majority label.
    - then we enforce minimum length for the merged segments via
      smooth_merged_segments.
    """
    assert len(ids_odd) == len(ids_even) == NUM_FRAMES
    n = NUM_FRAMES
    boundaries = {0, n}
    for ids in (ids_odd, ids_even):
        cur = ids[0]
        for i in range(1, n):
            if ids[i] != cur:
                boundaries.add(i)
                cur = ids[i]
    b = sorted(boundaries)
    merged = []
    for i in range(len(b) - 1):
        s = b[i]
        e = b[i + 1]
        if e <= s:
            continue
        slice_odd = ids_odd[s:e]
        slice_even = ids_even[s:e]
        if slice_odd.size == 0 or slice_even.size == 0:
            continue
        odd_vals, odd_counts = np.unique(slice_odd, return_counts=True)
        even_vals, even_counts = np.unique(slice_even, return_counts=True)
        odd_label = int(odd_vals[np.argmax(odd_counts)])
        even_label = int(even_vals[np.argmax(even_counts)])
        merged.append((s, e, odd_label, even_label))

    # Now enforce min length also on merged segments
    merged = smooth_merged_segments(merged, min_frames)
    return merged

# =========================
# ======= PLOTTING ========
# =========================

def _plot_tracks_seconds(pred_ids: torch.Tensor, title: str) -> bytes:
    """
    pred_ids: [2, NUM_FRAMES] LongTensor
    """
    secs = np.linspace(0.0, CLIP_SECONDS, NUM_FRAMES)
    fig = plt.figure(figsize=(10, 2.8))
    ax = plt.gca()
    im = ax.imshow(
        pred_ids.numpy(),
        aspect="auto",
        interpolation="nearest",
        origin="upper",
        extent=[secs[0], secs[-1], -0.5, 1.5],
    )
    ax.set_title(title)
    ax.set_xlabel("Time (s)")
    ax.set_yticks([0, 1])
    ax.set_yticklabels(["odd", "even"])
    cb = plt.colorbar(im, fraction=0.046, pad=0.04)
    cb.set_label("Segment ID")
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    return buf.read()

def _plot_merged_segments(seg_ids: np.ndarray, title: str) -> bytes:
    """
    seg_ids: [NUM_FRAMES] array where each frame holds a merged-segment index.
    """
    secs = np.linspace(0.0, CLIP_SECONDS, NUM_FRAMES)
    fig = plt.figure(figsize=(10, 2.8))
    ax = plt.gca()
    im = ax.imshow(
        seg_ids[np.newaxis, :],
        aspect="auto",
        interpolation="nearest",
        origin="upper",
        extent=[secs[0], secs[-1], -0.5, 0.5],
    )
    ax.set_title(title)
    ax.set_xlabel("Time (s)")
    ax.set_yticks([0])
    ax.set_yticklabels(["merged"])
    cb = plt.colorbar(im, fraction=0.046, pad=0.04)
    cb.set_label("Merged seg ID")
    buf = io.BytesIO()
    fig.savefig(buf, format="png", dpi=150, bbox_inches="tight")
    plt.close(fig)
    buf.seek(0)
    return buf.read()

# =========================
# ========= HTML ==========
# =========================

def write_html_report(out_dir: Path, chunks: List[Dict[str, Any]]) -> Path:
    ts = time.strftime("%Y%m%d_%H%M%S")
    html = [f"""<!doctype html><html><head><meta charset="utf-8">
<style>
 body{{font-family:system-ui,Segoe UI,Roboto,Arial,sans-serif;margin:20px}}
 .card{{border:1px solid #ddd;border-radius:10px;padding:16px;margin:16px 0;
        box-shadow:0 2px 6px rgba(0,0,0,.05)}}
 .grid{{display:grid;grid-template-columns:1fr 1fr;gap:12px}}
 figure{{margin:0}}
 figcaption{{font-size:13px;color:#555;margin-top:6px}}
 audio{{width:100%;min-width:200px;margin-top:4px}}
 .meta{{font-size:13px;color:#666;margin-bottom:4px}}
 table{{border-collapse:collapse;width:100%;margin-top:8px;font-size:13px;table-layout:fixed}}
 th,td{{border:1px solid #ddd;padding:4px 6px;text-align:left;vertical-align:top;overflow:hidden;text-overflow:ellipsis;white-space:nowrap}}
 th{{background:#f5f5f5}}
</style>
<title>Odd/Even Segmentation - Inference {ts}</title></head><body>
<h1>Odd/Even Segmentation - Inference</h1>
<p>
This report shows <b>smoothed</b> segmentations for each 30-second chunk of your audio files.
The model predicts two parallel time tracks ("odd" and "even") that can hold overlapping events.
We first smooth each track so that <b>no segment (including background 0) is shorter than {MIN_SEGMENT_SEC:.2f} seconds</b>.
Then:
</p>
<ul>
<li><b>Per-track segments</b>: segments for each track (odd/even) with duration &gt;= {MIN_SEGMENT_SEC:.2f}s, each with its own audio player.</li>
<li><b>Merged timeline</b>: a single segmentation where a new segment starts or ends whenever either track changes, and each merged segment is also at least {MIN_SEGMENT_SEC:.2f}s long by merging very short segments into their most similar neighbor.</li>
</ul>
"""]

    for ch in chunks:
        html.append(f"""
<section class="card">
  <h2>{ch['file_name']} - chunk {ch['chunk_idx']}</h2>
  <div class="meta">
    Chunk offset in file: {ch['chunk_offset']:.2f} - {ch['chunk_offset'] + CLIP_SECONDS:.2f} s
  </div>
  <div class="grid">
    <figure>
      <img src="data:image/png;base64,{ch['png_tracks']}" alt="smoothed tracks">
      <figcaption>Smoothed per-track predictions (odd/even).</figcaption>
    </figure>
    <figure>
      <img src="data:image/png;base64,{ch['png_merged']}" alt="merged timeline">
      <figcaption>Merged timeline: segment borders whenever odd or even track changes label, then smoothed to enforce a minimum duration.</figcaption>
    </figure>
  </div>

  <h3>Per-track segments (min {MIN_SEGMENT_SEC:.2f} s)</h3>
  <p>Each row is one predicted event on the odd or even track. Times are relative to the start of this 30-second chunk.</p>
  <table class="seg seg-track">
    <colgroup>
      <col style="width:5%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:45%">
    </colgroup>
    <tr><th>#</th><th>Track</th><th>Label ID</th><th>Start (s)</th><th>End (s)</th>
        <th>Duration (s)</th><th>Audio</th></tr>
""")
        # per-track table
        for i, seg in enumerate(ch["track_segments"], start=1):
            audio_cell = ""
            if seg["audio_b64"] and seg["audio_mime"]:
                audio_cell = (
                    '<audio controls preload="none">'
                    f'<source src="data:{seg["audio_mime"]};base64,{seg["audio_b64"]}" '
                    f'type="{seg["audio_mime"]}"></audio>'
                )
            html.append(
                f"<tr><td>{i}</td>"
                f"<td>{seg['track']}</td>"
                f"<td>{seg['label']}</td>"
                f"<td>{seg['start']:.2f}</td>"
                f"<td>{seg['end']:.2f}</td>"
                f"<td>{seg['dur']:.2f}</td>"
                f"<td>{audio_cell}</td></tr>"
            )
        html.append("</table>")

        # merged timeline table
        html.append(f"""
  <h3>Merged timeline segments</h3>
  <p>
The merged timeline splits the 30-second chunk wherever either the odd or even track changes label.
Very short merged segments (shorter than {MIN_SEGMENT_SEC:.2f}s) are merged into their most similar neighbor
based on odd/even labels; if both neighbors are equally similar, they are merged into the shorter neighbor.
This yields a single sequence of non-overlapping segments that cover the entire chunk.
Each row shows the majority label on the odd and even tracks within that merged segment.
  </p>
  <table class="seg seg-merged">
    <colgroup>
      <col style="width:5%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:10%">
      <col style="width:45%">
    </colgroup>
    <tr><th>#</th><th>Start (s)</th><th>End (s)</th><th>Duration (s)</th>
        <th>Odd label</th><th>Even label</th><th>Audio</th></tr>
""")
        for i, seg in enumerate(ch["merged_segments"], start=1):
            audio_cell = ""
            if seg["audio_b64"] and seg["audio_mime"]:
                audio_cell = (
                    '<audio controls preload="none">'
                    f'<source src="data:{seg["audio_mime"]};base64,{seg["audio_b64"]}" '
                    f'type="{seg["audio_mime"]}"></audio>'
                )
            html.append(
                f"<tr><td>{i}</td>"
                f"<td>{seg['start']:.2f}</td>"
                f"<td>{seg['end']:.2f}</td>"
                f"<td>{seg['dur']:.2f}</td>"
                f"<td>{seg['odd_label']}</td>"
                f"<td>{seg['even_label']}</td>"
                f"<td>{audio_cell}</td></tr>"
            )
        html.append("</table></section>")

    html.append("</body></html>")
    out_path = out_dir / f"seg_infer_smooth_{ts}.html"
    out_path.write_text("\n".join(html), encoding="utf-8")
    return out_path

# =========================
# ========= MAIN ==========
# =========================

def main():
    setup_dirs()

    global AUDIO_INPUT_DIR
    if len(sys.argv) > 1:
        AUDIO_INPUT_DIR = Path(sys.argv[1])

    if not AUDIO_INPUT_DIR.is_dir():
        print(f"[ERR] AUDIO_INPUT_DIR not found or not a dir: {AUDIO_INPUT_DIR}", file=sys.stderr)
        sys.exit(1)
    if not CKPT_PATH.is_file():
        print(f"[ERR] Checkpoint not found: {CKPT_PATH}", file=sys.stderr)
        sys.exit(1)

    print(f"[cfg] AUDIO_INPUT_DIR  = {AUDIO_INPUT_DIR}")
    print(f"[cfg] OUT_DIR          = {OUT_DIR}")
    print(f"[cfg] CKPT_PATH       = {CKPT_PATH}")
    print(f"[cfg] HF_MODEL_ID     = {HF_MODEL_ID}")
    print(f"[cfg] ffmpeg available: {FFMPEG_AVAILABLE}")
    print(f"[cfg] MIN_SEGMENT_SEC  = {MIN_SEGMENT_SEC:.2f} (frames >= {MIN_SEGMENT_FRAMES})")

    # find audio files
    exts = {".wav", ".mp3", ".m4a", ".flac", ".ogg"}
    audio_files: List[Path] = []
    for p in AUDIO_INPUT_DIR.rglob("*"):
        if p.is_file() and p.suffix.lower() in exts:
            audio_files.append(p)
    audio_files = sorted(audio_files)

    if not audio_files:
        print("[ERR] No audio files found.", file=sys.stderr)
        sys.exit(1)

    print(f"[scan] Found {len(audio_files)} audio files.")

    # feature extractor
    resolved, is_local = _model_resolved_name(HF_MODEL_ID)
    fe = WhisperFeatureExtractor.from_pretrained(resolved, local_files_only=is_local)

    # model + checkpoint
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = WhisperOddEven(HF_MODEL_ID, freeze_encoder=False).to(device)

    state = torch.load(CKPT_PATH, map_location="cpu")
    # accept full trainer_state dict or plain state_dict
    if isinstance(state, dict) and "model" in state and any(
        k.startswith("whisper.") for k in state["model"].keys()
    ):
        state = state["model"]

    missing, unexpected = model.load_state_dict(state, strict=False)
    print(f"[ckpt] Loaded checkpoint from {CKPT_PATH}")
    if missing:
        print(f"[ckpt] Missing keys: {missing}")
    if unexpected:
        print(f"[ckpt] Unexpected keys: {unexpected}")
    model.eval()

    use_dtype = preferred_dtype()
    amp_enabled = use_dtype in (torch.float16, torch.bfloat16)

    chunk_results: List[Dict[str, Any]] = []

    with torch.no_grad():
        for fpath in audio_files:
            print(f"[file] {fpath}")
            try:
                wav = load_audio_mono_16k(fpath)
            except Exception as e:
                print(f"[file] Failed to load {fpath}: {e}")
                continue

            chunks = split_into_chunks(wav, SAMPLE_RATE, CLIP_SECONDS)
            if not chunks:
                print(f"[file] No audio samples in {fpath}")
                continue

            for chunk_idx, start_sample, seg in chunks:
                chunk_offset_sec = start_sample / SAMPLE_RATE

                # features
                feat = fe(seg, sampling_rate=SAMPLE_RATE, return_tensors="pt")
                x = feat.input_features.to(device)

                # forward
                with torch.autocast(
                    device_type="cuda" if torch.cuda.is_available() else "cpu",
                    enabled=amp_enabled,
                    dtype=use_dtype,
                ):
                    logits = model(x)

                # raw argmax
                raw_ids = logits.argmax(dim=-1).squeeze(0).cpu().numpy()  # [2,1500]

                # aggressive smoothing with min duration per track
                sm_ids = np.zeros_like(raw_ids)
                for tr in range(NUM_TRACKS):
                    sm_ids[tr] = smooth_min_duration(raw_ids[tr], MIN_SEGMENT_FRAMES)

                sm_ids_t = torch.from_numpy(sm_ids)
                png_tracks = base64.b64encode(
                    _plot_tracks_seconds(
                        sm_ids_t,
                        f"Smoothed tracks - {fpath.name} - chunk {chunk_idx}",
                    )
                ).decode("ascii")

                # merged timeline with its own min-duration smoothing
                merged = build_merged_segments(sm_ids[0], sm_ids[1], MIN_SEGMENT_FRAMES)
                merged_index = np.zeros(NUM_FRAMES, dtype=np.int64)
                for idx, (fs, fe_, _ol, _el) in enumerate(merged, start=1):
                    merged_index[fs:fe_] = idx

                png_merged = base64.b64encode(
                    _plot_merged_segments(
                        merged_index,
                        f"Merged segments - {fpath.name} - chunk {chunk_idx}",
                    )
                ).decode("ascii")

                # per-track segments -> audio snippets
                track_segments: List[Dict[str, Any]] = []
                for tr, track_name in enumerate(("odd", "even")):
                    seg_runs = extract_segments(sm_ids[tr], include_bg=False)
                    for (lab, fs, fe_) in seg_runs:
                        start_t, end_t = frames_to_times(fs, fe_)
                        dur = end_t - start_t
                        if dur <= 0:
                            continue
                        sub_wav = cut_wav(seg, start_t, end_t)
                        if sub_wav.size == 0:
                            continue
                        try:
                            audio_bytes, audio_mime = wav_chunk_to_audio_bytes(sub_wav, SAMPLE_RATE)
                            audio_b64 = base64.b64encode(audio_bytes).decode("ascii")
                        except Exception as e:
                            print(f"[audio] Failed per-track snippet for {fpath} chunk {chunk_idx}: {e}")
                            audio_b64 = None
                            audio_mime = None

                        track_segments.append(
                            {
                                "track": track_name,
                                "label": int(lab),
                                "start": float(start_t),
                                "end": float(end_t),
                                "dur": float(dur),
                                "audio_b64": audio_b64,
                                "audio_mime": audio_mime,
                            }
                        )

                # merged segments -> audio snippets
                merged_segments: List[Dict[str, Any]] = []
                for idx, (fs, fe_, odd_label, even_label) in enumerate(merged, start=1):
                    start_t, end_t = frames_to_times(fs, fe_)
                    dur = end_t - start_t
                    if dur <= 0:
                        continue
                    sub_wav = cut_wav(seg, start_t, end_t)
                    if sub_wav.size == 0:
                        continue
                    try:
                        audio_bytes, audio_mime = wav_chunk_to_audio_bytes(sub_wav, SAMPLE_RATE)
                        audio_b64 = base64.b64encode(audio_bytes).decode("ascii")
                    except Exception as e:
                        print(f"[audio] Failed merged snippet for {fpath} chunk {chunk_idx}: {e}")
                        audio_b64 = None
                        audio_mime = None

                    merged_segments.append(
                        {
                            "idx": idx,
                            "start": float(start_t),
                            "end": float(end_t),
                            "dur": float(dur),
                            "odd_label": int(odd_label),
                            "even_label": int(even_label),
                            "audio_b64": audio_b64,
                            "audio_mime": audio_mime,
                        }
                    )

                chunk_results.append(
                    {
                        "file_name": fpath.name,
                        "chunk_idx": int(chunk_idx),
                        "chunk_offset": float(chunk_offset_sec),
                        "png_tracks": png_tracks,
                        "png_merged": png_merged,
                        "track_segments": track_segments,
                        "merged_segments": merged_segments,
                    }
                )

    if not chunk_results:
        print("[ERR] No chunk results; nothing to write.", file=sys.stderr)
        sys.exit(1)

    out_html = write_html_report(OUT_DIR, chunk_results)
    print(f"[done] Wrote HTML report: {out_html}")

if __name__ == "__main__":
    main()