helpers.py

"""Helpers for the Deep Voice multi-model bioacoustic Gradio app.

Wraps soundbay's `inference_to_file` so we can run inference programmatically
(without a Hydra CLI shell-out), and bundles file-validation / Raven-merge
utilities used by the UI layer in `app.py`.
"""
from __future__ import annotations

import os
import zipfile
from functools import lru_cache
from pathlib import Path
from typing import Optional

import pandas as pd
import soundfile as sf
import torch
from hydra import compose, initialize
from hydra.core.global_hydra import GlobalHydra
from huggingface_hub import hf_hub_download

from soundbay.inference import inference_to_file
from soundbay.utils.checkpoint_utils import merge_with_checkpoint

HF_CKPT_REPO = "deepvoice1/bioacoustic-checkpoints"

MAX_TOTAL_BYTES = 500 * 1024 * 1024  # 500 MB — only hard cap; runtime estimate guards the rest

# Raven TSV column order, used when we have to emit an empty merged file.
_RAVEN_COLS = [
    "Selection", "View", "Channel", "Begin Time (s)", "End Time (s)",
    "Low Freq (Hz)", "High Freq (Hz)", "Annotation", "Class Name",
    "Probability", "Begin File",
]


@lru_cache(maxsize=16)
def download_checkpoint(filename: str) -> str:
    """Resolve a checkpoint path.

    If `DV_LOCAL_CKPT_DIR` is set, look it up locally (for development).
    Otherwise pull from the HF model repo, cached on disk.
    """
    local_dir = os.environ.get("DV_LOCAL_CKPT_DIR")
    if local_dir:
        local_path = Path(local_dir) / filename
        if local_path.exists():
            return str(local_path)
        raise FileNotFoundError(
            f"DV_LOCAL_CKPT_DIR is set but {local_path} does not exist."
        )
    return hf_hub_download(
        repo_id=HF_CKPT_REPO,
        filename=filename,
        token=os.environ.get("HF_TOKEN"),
    )


def _compose_inference_cfg(wav_path: str, ckpt_path: str, data_sr: int, threshold: float):
    """Build the OmegaConf cfg for a single-file inference run."""
    if GlobalHydra.instance().is_initialized():
        GlobalHydra.instance().clear()
    # config_path is relative to *this file*; soundbay/ is a sibling of helpers.py
    with initialize(config_path="soundbay/conf", version_base="1.2"):
        cfg = compose(
            config_name="runs/inference_single_audio",
            overrides=[
                f"experiment.checkpoint.path={ckpt_path}",
                f"data.test_dataset.file_path={wav_path}",
                f"data.data_sample_rate={data_sr}",
                f"experiment.threshold={threshold}",
                "experiment.save_raven=true",
            ],
        )
    return cfg


def run_inference(
    wav_path: str,
    ckpt_path: str,
    threshold: float,
    output_dir: str | Path,
) -> tuple[Path, Optional[Path]]:
    """Run inference on one wav and return (csv_path, raven_path-or-None).

    Picks `data_sample_rate` from the wav file itself (not whatever the model
    was trained on). Only constraint: wav SR must be >= the model's internal
    sample rate, otherwise we'd be making up information that isn't there.
    """
    ckpt_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    model_sr = int(ckpt_dict["args"].data.sample_rate)
    wav_sr = int(sf.info(wav_path).samplerate)
    if wav_sr < model_sr:
        raise ValueError(
            f"{Path(wav_path).name}: sample rate {wav_sr} Hz is below this "
            f"model's required minimum of {model_sr} Hz — upload a higher-rate recording."
        )

    cfg = _compose_inference_cfg(wav_path, ckpt_path, wav_sr, threshold)
    cfg = merge_with_checkpoint(cfg, ckpt_dict["args"])
    state_dict = ckpt_dict["model"]

    default_norm = "softmax" if cfg.data.label_type == "single_label" else "sigmoid"
    out = Path(output_dir)
    out.mkdir(exist_ok=True, parents=True)
    model_name = Path(ckpt_path).parent.stem

    inference_to_file(
        device=torch.device("cpu"),
        batch_size=cfg.data.batch_size,
        dataset_args=cfg.data.test_dataset,
        model_args=cfg.model.model,
        checkpoint_state_dict=state_dict,
        output_path=out,
        model_name=model_name,
        save_raven=cfg.experiment.save_raven,
        threshold=cfg.experiment.threshold,
        label_names=cfg.data.label_names,
        raven_max_freq=cfg.experiment.raven_max_freq,
        proba_norm_func=cfg.data.get("proba_norm_func", default_norm),
        label_type=cfg.data.label_type,
    )

    wav_stem = Path(wav_path).stem
    csvs = sorted(out.glob(f"Inference_results-*-{model_name}-{wav_stem}.csv"))
    ravens = sorted(out.glob(f"{wav_stem}-Raven-inference_results-*-{model_name}.txt"))
    if not csvs:
        raise RuntimeError(f"No CSV output produced for {wav_path}")
    # Rewrite the CSV's `filename` column to just the basename — the original
    # value is the gradio temp upload path (/tmp/gradio/<hash>/...) which is
    # noisy and useless to the user.
    wav_basename = Path(wav_path).name
    df = pd.read_csv(csvs[-1])
    df["filename"] = wav_basename
    df.to_csv(csvs[-1], index=False)
    return csvs[-1], (ravens[-1] if ravens else None)


def merge_ravens(
    raven_paths: list[Optional[Path]],
    audio_paths: list[Path],
    output_path: Path,
) -> Path:
    """Concatenate per-file Raven TSVs into one, with offset times and a
    `Begin File` column. Mirrors `scripts/merge_multiple_ravens_to_one_file.py`
    but without its CLI / strict-count assertions."""
    df_list: list[pd.DataFrame] = []
    seconds_offset = 0.0
    entries_offset = 0
    for raven_p, audio_p in zip(raven_paths, audio_paths):
        if raven_p is None or not Path(raven_p).exists():
            # Still need to advance the time offset by this file's duration.
            seconds_offset += sf.info(str(audio_p)).duration
            continue
        df = pd.read_csv(raven_p, sep="\t")
        if len(df):
            df["Begin Time (s)"] = df["Begin Time (s)"] + seconds_offset
            df["End Time (s)"] = df["End Time (s)"] + seconds_offset
            df["Selection"] = df["Selection"] + entries_offset
        df["Begin File"] = [Path(audio_p).name] * df.shape[0]
        df_list.append(df)
        entries_offset += df.shape[0]
        seconds_offset += sf.info(str(audio_p)).duration

    if df_list:
        pd.concat(df_list).to_csv(output_path, sep="\t", index=False)
    else:
        pd.DataFrame(columns=_RAVEN_COLS).to_csv(output_path, sep="\t", index=False)
    return output_path


def validate_uploads(file_paths: list[str]) -> dict:
    """Inspect uploads; return summary dict. Raises gr.Error on cap violation."""
    import gradio as gr  # local import so helpers stays usable in non-Gradio contexts

    if not file_paths:
        raise gr.Error("Please upload at least one .wav file.")

    total_bytes = 0
    durations: list[float] = []
    for fp in file_paths:
        p = Path(fp)
        if p.suffix.lower() != ".wav":
            raise gr.Error(f"Not a .wav file: {p.name}")
        total_bytes += p.stat().st_size
        try:
            durations.append(sf.info(str(p)).duration)
        except Exception as e:
            raise gr.Error(f"Cannot read {p.name}: {e}")

    if total_bytes > MAX_TOTAL_BYTES:
        raise gr.Error(
            f"Total size {total_bytes/1e6:.1f} MB > {MAX_TOTAL_BYTES/1e6:.0f} MB."
        )

    return {
        "n_files": len(file_paths),
        "total_mb": total_bytes / 1e6,
        "total_min": sum(durations) / 60,
        "durations": durations,
    }


def estimate_minutes(total_audio_seconds: float, coef: float = 0.2) -> float:
    """Upper-bound runtime estimate assuming `coef` x real-time on CPU."""
    return (total_audio_seconds * coef) / 60.0


def zip_files(paths: list[Path], zip_path: Path) -> Path:
    """Bundle the given paths into a single zip (flat layout)."""
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        for p in paths:
            zf.write(p, arcname=Path(p).name)
    return zip_path