nc-murray/spectrogram-reconstruction

1"""Spectrogram image → linear amplitude magnitude array."""2 3import numpy as np4from PIL import Image5import matplotlib.cm as mcm6from scipy.spatial import cKDTree7 8 9# Build LUT for a colormap: 256 RGB triples in [0, 255]10def _build_lut(colormap: str) -> tuple[np.ndarray, cKDTree]:11    cmap = mcm._colormaps[colormap]12    values = np.linspace(0.0, 1.0, 256)13    rgba = cmap(values)           # (256, 4)14    rgb = (rgba[:, :3] * 255).astype(np.uint8)15    return values, cKDTree(rgb.astype(np.float32))16 17 18# Cache so we don't rebuild on every call19_LUT_CACHE: dict[str, tuple[np.ndarray, cKDTree]] = {}20 21SUPPORTED_COLORMAPS = ["viridis", "jet", "Greys", "gray", "magma", "plasma", "inferno"]22 23# Normalize user-supplied aliases to the canonical matplotlib name24_CMAP_ALIASES = {"greys": "Greys", "grey": "Greys", "gray": "gray"}25 26 27def _canonical(colormap: str) -> str:28    return _CMAP_ALIASES.get(colormap.lower(), colormap)29 30 31def _get_lut(colormap: str) -> tuple[np.ndarray, cKDTree]:32    key = _canonical(colormap)33    if key not in _LUT_CACHE:34        _LUT_CACHE[key] = _build_lut(key)35    return _LUT_CACHE[key]36 37 38def load_spectrogram_image(path: str) -> np.ndarray:39    """Load image, return RGB uint8 array shaped (H, W, 3)."""40    img = Image.open(path).convert("RGB")41    return np.array(img, dtype=np.uint8)42 43 44def detect_colormap(image: np.ndarray, hint: str = None) -> str:45    """46    Return the best-matching colormap name.47    If hint is given, validate it's supported and return it.48    Otherwise score each candidate by sampling pixel errors on a random subset.49    """50    if hint is not None:51        key = _canonical(hint)52        if key not in SUPPORTED_COLORMAPS:53            raise ValueError(f"Unsupported colormap hint '{hint}'. Choose from {SUPPORTED_COLORMAPS}")54        return key55 56    rng = np.random.default_rng(0)57    h, w = image.shape[:2]58    idx = rng.choice(h * w, size=min(2000, h * w), replace=False)59    pixels = image.reshape(-1, 3)[idx].astype(np.float32)60 61    best_cmap, best_err = None, np.inf62    for name in SUPPORTED_COLORMAPS:63        name = _canonical(name)64        _, tree = _get_lut(name)65        dists, _ = tree.query(pixels, workers=-1)66        err = dists.mean()67        if err < best_err:68            best_err, best_cmap = err, name69    return best_cmap70 71 72def invert_colormap(image: np.ndarray, colormap: str) -> np.ndarray:73    """74    Map every pixel's RGB triple to its nearest scalar value in [0, 1]75    using the colormap's LUT.  Returns float32 array shaped (H, W).76    """77    values, tree = _get_lut(colormap.lower())78    h, w = image.shape[:2]79    pixels = image.reshape(-1, 3).astype(np.float32)80    _, indices = tree.query(pixels, workers=-1)81    return values[indices].reshape(h, w).astype(np.float32)82 83 84def apply_db_scaling(85    magnitude_normalized: np.ndarray,86    db_min: float = -80.0,87    db_max: float = 0.0,88) -> np.ndarray:89    """90    normalized [0, 1]  →  dB  →  linear amplitude.91    Input 0.0 maps to db_min, input 1.0 maps to db_max.92    """93    db = magnitude_normalized * (db_max - db_min) + db_min94    return librosa_db_to_amplitude(db)95 96 97def librosa_db_to_amplitude(db: np.ndarray) -> np.ndarray:98    import librosa99    return librosa.db_to_amplitude(db).astype(np.float32)100 101 102def parse_spectrogram_image(103    path: str,104    colormap: str = None,105    freq_min: float = 0.0,106    freq_max: float = 8000.0,107    db_min: float = -80.0,108    db_max: float = 0.0,109    freq_log_scale: bool = False,110    crop: tuple = None,111) -> dict:112    """113    Full pipeline: image → linear amplitude magnitude array.114 115    Returns:116        {117          'magnitude': np.ndarray (freq_bins × time_frames, linear amplitude),118          'metadata': dict with freq_min, freq_max, db_min, db_max, colormap119        }120    """121    image = load_spectrogram_image(path)122 123    if crop is not None:124        l, t, r, b = crop125        image = image[t:b, l:r]126 127    cmap = detect_colormap(image, hint=colormap)128    normalized = invert_colormap(image, cmap)129 130    # Image rows: top = high freq, bottom = low freq (origin="lower" in synthesizer)131    # Flip vertically so row 0 = lowest frequency bin132    normalized = np.flipud(normalized)133 134    magnitude = apply_db_scaling(normalized, db_min=db_min, db_max=db_max)135 136    return {137        "magnitude": magnitude,138        "metadata": {139            "colormap": cmap,140            "freq_min": freq_min,141            "freq_max": freq_max,142            "db_min": db_min,143            "db_max": db_max,144        },145    }146