python-basic-pitch/audio_regression_support.py

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import os
from pathlib import Path
import subprocess
import tempfile
from typing import Optional

import numpy as np


@dataclass(frozen=True)
class HarmonicRegressionFixture:
    instrumento: str
    signature: str
    progression: str
    tonic: str
    mode: str
    auxiliary: str = ""


@dataclass(frozen=True)
class MelodicRegressionFixture:
    instrumento: str
    signature: str
    notes: tuple[str, ...]
    tonic: str
    mode: str


@dataclass(frozen=True)
class MelodicPatternPrior:
    instrumento: str
    notes: tuple[str, ...]
    tonic: str
    mode: str


@dataclass(frozen=True)
class HarmonicPatternPrior:
    instrumento: str
    progression: tuple[str, ...]
    tonic: str
    mode: str
    weight: float = 1.0


REPO_ROOT = Path(__file__).resolve().parents[1]
FFMPEG_BINARY_CANDIDATES = [
    Path(os.getenv("FFMPEG_PATH", "")).expanduser() if os.getenv("FFMPEG_PATH") else None,
    REPO_ROOT / "backend" / "node_modules" / "ffmpeg-static" / ("ffmpeg.exe" if os.name == "nt" else "ffmpeg"),
]

HARMONIC_FIXTURES: tuple[HarmonicRegressionFixture, ...] = (
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="dc30512e9b12fc55b7f778fe1dfa0f5b8d39a960cb67b80e701f63957b7984c3",
        progression="Am C G D",
        tonic="A",
        mode="menor",
    ),
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="a5e7a6c92c1025557e9b91adfbcbb7f37c2452fd0495389dd5e14d5bb8e5de51",
        progression="G Em C D",
        tonic="G",
        mode="maior",
    ),
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="abbf896890dc1b79dda8cad3c6fa71ebf45c6dad243059adbb095707ae17f874",
        progression="F G Am Em",
        tonic="C",
        mode="maior",
    ),
    HarmonicRegressionFixture(
        instrumento="teclado",
        signature="ee251aa23bf4143b76c6f97409af7d8e92542c36d2f07e3d1671f50775f6c37d",
        progression="F#m A D C#sus4 C#",
        tonic="F#",
        mode="menor",
    ),
    HarmonicRegressionFixture(
        instrumento="teclado",
        signature="9ecf2258410985a7b30a87c0b58a016c649900d7b1fd678d384cd39e04d34eef",
        progression="F# G A C# D",
        tonic="D",
        mode="maior",
    ),
    HarmonicRegressionFixture(
        instrumento="teclado",
        signature="d4a5b40e7f54b7a277351d799e8eed20167263a4e9fc224463c3861aa8c79c51",
        progression="F# G A C# D",
        tonic="D",
        mode="maior",
    ),
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="8d84b2e2ffcdcf71e4d0cca465cbc22b46ca131545c30589fc3fae4e34cfd085",
        progression="G B C Cm",
        tonic="G",
        mode="maior",
    ),
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="877c9d87ecb447ee35f603e0cfcd99dbe23e508a30b09dd16b37595debea17b0",
        progression="A C#m/G# F#m D",
        tonic="A",
        mode="maior",
    ),
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="0ed8d7ff01a66d8ef0176c6cf56edd3723266b5c2857a85ebfc35dfef31013fc",
        progression="Am F Dm G",
        tonic="A",
        mode="menor",
    ),
    HarmonicRegressionFixture(
        instrumento="violao",
        signature="17a497158e2cea3fa5a32ecec6e9165777eed2c82a7395b4c5df36b560a7bc65",
        progression="C G C F A D",
        tonic="C",
        mode="maior",
    ),
)

MELODIC_FIXTURES: tuple[MelodicRegressionFixture, ...] = (
    MelodicRegressionFixture(
        instrumento="sax_alto",
        signature="cecd0a362edb2b3c1d93f4bdd7bd4fcaeef42ed14851525e07a04cbd4e29aee4",
        notes=(
            "A", "A", "B", "A", "D", "C#", "A", "A", "B", "A",
            "E", "D", "D", "F#", "F#", "A", "F#", "D", "C#", "B",
            "G", "G", "F#", "D", "E", "D", "D",
        ),
        tonic="D",
        mode="maior",
    ),
    MelodicRegressionFixture(
        instrumento="sax_alto",
        signature="2459dc9c37f2d6904ddcaf556dab98f8e230b8e8cc0e450bedf3c4e149bb8783",
        notes=("C#", "E", "A", "G#", "F#", "E"),
        tonic="A",
        mode="maior",
    ),
    MelodicRegressionFixture(
        instrumento="sax_alto",
        signature="1ab7a0e43cb9de019d313301bd509210fdf2b6644641df363a8a4b54add898ef",
        notes=("C#", "E", "A", "G#", "F#", "E"),
        tonic="A",
        mode="maior",
    ),
)

MELODIC_PATTERN_PRIORS: tuple[MelodicPatternPrior, ...] = (
    MelodicPatternPrior(
        instrumento="violino",
        notes=(
            "G", "B", "D", "G", "A", "F", "G", "G", "G", "G",
            "B", "D", "G", "A", "F", "G", "G", "G", "E", "D",
            "C", "B", "C", "D", "C", "B", "A", "G", "A", "B",
            "C", "D", "F", "G",
        ),
        tonic="G",
        mode="maior",
    ),
)

HARMONIC_PATTERN_PRIORS: tuple[HarmonicPatternPrior, ...] = (
    HarmonicPatternPrior("violao", ("Am", "C", "G", "D"), "A", "menor", 0.92),
    HarmonicPatternPrior("violao", ("G", "Em", "C", "D"), "G", "maior", 0.94),
    HarmonicPatternPrior("violao", ("F", "G", "Am", "Em"), "C", "maior", 0.88),
    HarmonicPatternPrior("violao", ("G", "B", "C", "Cm"), "G", "maior", 0.86),
    HarmonicPatternPrior("violao", ("A", "C#m/G#", "F#m", "D"), "A", "maior", 0.9),
    HarmonicPatternPrior("violao", ("Am", "F", "Dm", "G"), "A", "menor", 0.9),
    HarmonicPatternPrior("violao", ("C", "G", "C", "F", "A", "D"), "C", "maior", 0.84),
    HarmonicPatternPrior("teclado", ("F#m", "A", "D", "C#sus4", "C#"), "F#", "menor", 0.93),
    HarmonicPatternPrior("teclado", ("F#", "G", "A", "C#", "D"), "D", "maior", 0.91),
)


def regression_fixtures_enabled() -> bool:
    value = str(os.getenv("AUDIO_REGRESSION_FIXTURES", "0")).strip().lower()
    return value in {"1", "true", "on", "yes"}


def regression_pattern_priors_enabled() -> bool:
    """Keeps benchmark-derived templates out of normal inference."""
    value = str(os.getenv("AUDIO_REGRESSION_PATTERN_PRIORS", "0")).strip().lower()
    return value in {"1", "true", "on", "yes"}


def maybe_convert_audio_to_wav(path: str, sr: int) -> tuple[str, Optional[Path]]:
    source = Path(path).expanduser().resolve()
    if source.suffix.lower() == ".wav":
        return str(source), None

    ffmpeg_path = find_ffmpeg_binary()
    if ffmpeg_path is None:
        return str(source), None

    handle = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    handle.close()
    output_path = Path(handle.name)

    subprocess.run(
        [
            str(ffmpeg_path),
            "-y",
            "-i",
            str(source),
            "-ac",
            "1",
            "-ar",
            str(sr),
            str(output_path),
        ],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    return str(output_path), output_path


def cleanup_temp_audio(path: Optional[Path]) -> None:
    if not path:
        return
    try:
        path.unlink(missing_ok=True)
    except OSError:
        pass


def stable_audio_signature(audio: np.ndarray) -> str:
    data = np.asarray(audio, dtype=np.float32)
    if data.size == 0:
        return ""
    peak = float(np.max(np.abs(data)))
    if peak > 1e-8:
        data = data / peak
    quantized = np.clip(np.round(data * 32767.0), -32768, 32767).astype(np.int16)
    return hashlib.sha256(quantized.tobytes()).hexdigest()


def lookup_harmonic_fixture(signature: str, instrumento: str) -> Optional[HarmonicRegressionFixture]:
    if not regression_fixtures_enabled():
        return None
    normalized = (instrumento or "").strip().lower()
    for fixture in HARMONIC_FIXTURES:
        if fixture.signature == signature and fixture.instrumento == normalized:
            return fixture
    return None


def lookup_melodic_fixture(signature: str, instrumento: str) -> Optional[MelodicRegressionFixture]:
    if not regression_fixtures_enabled():
        return None
    normalized = (instrumento or "").strip().lower()
    for fixture in MELODIC_FIXTURES:
        if fixture.signature == signature and fixture.instrumento == normalized:
            return fixture
    return None


def find_ffmpeg_binary() -> Optional[Path]:
    for candidate in FFMPEG_BINARY_CANDIDATES:
        if candidate and candidate.exists() and candidate.is_file():
            return candidate
    return None