uj-mastering-master/audio_visualization_widget.py

"""
Interactive visualization widget built on pyqtgraph.

One persistent PlotItem that is *reused* across renders — never torn down — so
mouse zoom/pan, the view box, and scale toggles all survive redraws. Consumes a
list of `(label, PlotSpec)` pairs and draws them onto the same axes, using a
caller-supplied colour per dataset so a song keeps its colour regardless of which
others are overlaid.

Interaction notes:
- Plain scroll zooms both axes; Ctrl+scroll zooms time only; Shift+scroll zooms
  the value axis only (see `_AxisZoomViewBox`). Scrolling directly over an axis
  also zooms just that axis (pyqtgraph default).
- Reference lines (`set_reference_lines`) are draggable via a triangle handle,
  survive redraws, and write their position back into the GUI-owned RefLineProps;
  the GUI clears them when the metric changes (units change).

Why the spectrogram is special: pyqtgraph's ImageItem is affine-only, so it does
not follow a log-scaled axis. Log frequency is therefore realised by resampling
the STFT rows onto a log-spaced grid and labelling the axis by row index — see
`_render_heatmap`.
"""

import numpy as np
import pyqtgraph as pg
from scipy.interpolate import interp1d
from PyQt5.QtWidgets import QWidget, QVBoxLayout, QLabel
from PyQt5.QtCore import Qt, pyqtSignal

from plotspec import PlotSpec, ViewState, DEFAULT_VIEW, RefLineProps

# White canvas / black ink to match the previous matplotlib aesthetic.
pg.setConfigOption("background", "w")
pg.setConfigOption("foreground", "k")
pg.setConfigOptions(antialias=True)

# Dataset colour cycle for overlay. First colour is the single-dataset default.
_PALETTE = [
  "#3a7ad6", "#e76f51", "#2a9d8f", "#e09f3e",
  "#7251b5", "#c1121f", "#588157", "#9d4edd",
]

# Pen styles for reference lines.
_PEN_STYLE = {"solid": Qt.SolidLine, "dash": Qt.DashLine, "dot": Qt.DotLine}

# "Nice" frequencies to label on a log frequency axis, in Hz.
_LOG_FREQ_TICKS = [20, 50, 100, 200, 500, 1000, 2000, 5000, 10000, 20000]


def dataset_color(index: int) -> str:
  """Stable dataset colour for a given index (e.g. a file's row in the list)."""
  return _PALETTE[index % len(_PALETTE)]


def _colormap(name: str):
  """Fetch a colormap, preferring matplotlib's so 'magma' etc. resolve."""
  try:
    return pg.colormap.getFromMatplotlib(name)
  except Exception:
    return pg.colormap.get(name)


def _fmt_hz(hz: float) -> str:
  return f"{hz / 1000:.0f}k" if hz >= 1000 else f"{hz:.0f}"


class _AxisZoomViewBox(pg.ViewBox):
  """ViewBox whose wheel zoom can be constrained to one axis via a modifier.

  Plain scroll keeps pyqtgraph's both-axes zoom; Ctrl constrains to x (time),
  Shift constrains to y (the metric's value axis). This answers the "scroll
  zooms both axes, I want one" problem without taking away the default.
  """

  def wheelEvent(self, ev, axis=None):
    mods = ev.modifiers()
    if mods & Qt.ControlModifier:
      axis = 0  # x only
    elif mods & Qt.ShiftModifier:
      axis = 1  # y only
    super().wheelEvent(ev, axis=axis)


class _RefLine(pg.InfiniteLine):
  """A draggable horizontal reference line bound to a RefLineProps.

  `props.value` is in the metric's natural units (LUFS, dB, ... or Hz for the
  spectrogram). The drawn y-position may differ from that value — the spectrogram
  maps frequency to a row index — so `to_pos`/`from_pos` convert between the two.
  For curve metrics these are identity. The label and the value written back on
  drag are always in natural units.
  """

  def __init__(self, index, props: RefLineProps, to_pos, from_pos, fmt, on_moved):
    pen = pg.mkPen(props.color, width=1.4,
                   style=_PEN_STYLE.get(props.style, Qt.DashLine))
    super().__init__(
      pos=to_pos(props.value), angle=0, movable=True, pen=pen,
      label="",
      labelOpts={"position": 0.06, "color": props.color,
                 "fill": (255, 255, 255, 180)},
    )
    self._index = index
    self._props = props
    self._from_pos = from_pos
    self._fmt = fmt
    self._on_moved = on_moved
    self.addMarker("|>", position=0.0, size=12)  # triangle handle at the start
    self._update_label()
    self.sigPositionChanged.connect(self._update_label)
    self.sigPositionChangeFinished.connect(self._commit)

  def _update_label(self):
    val = self._from_pos(self.value())
    self.label.setFormat(self._props.label or self._fmt(val))

  def _commit(self):
    self._props.value = float(self._from_pos(self.value()))
    self._on_moved(self._index)


class AudioVisualizationWidget(QWidget):
  """Persistent interactive plot. Call `show_specs` to (re)draw."""

  # Emitted (with the line's index) when a reference line is dragged, so the
  # side-panel list can refresh its displayed value.
  referenceLineMoved = pyqtSignal(int)

  def __init__(self, parent=None):
    super().__init__(parent)
    layout = QVBoxLayout(self)

    self.glw = pg.GraphicsLayoutWidget()
    self.plot = self.glw.addPlot(row=0, col=0, viewBox=_AxisZoomViewBox())
    self.plot.showGrid(x=True, y=True, alpha=0.3)
    self.plot.setMenuEnabled(True)
    self.legend = self.plot.addLegend(offset=(-10, 10))
    layout.addWidget(self.glw)

    self.status_label = QLabel("Ready for audio analysis...")
    layout.addWidget(self.status_label)

    self._colorbar = None
    # Reference lines are owned by the GUI controller (RefLineProps objects) and
    # passed in via set_reference_lines; the line items are rebuilt each render.
    self._ref_props: list[RefLineProps] = []
    self._ref_lines: list[_RefLine] = []
    # value<->drawn-position transforms for ref lines (identity for curve metrics;
    # frequency<->row-index for the spectrogram). Reset each render.
    self._ref_to_pos = lambda v: v
    self._ref_from_pos = lambda p: p
    self._ref_fmt = lambda v: f"{v:.2f}"
    self._show_empty()

  # ---- public API ---------------------------------------------------------

  def show_specs(self, specs, view: ViewState = DEFAULT_VIEW):
    """Render datasets onto the shared axes.

    `specs` is a list of `(label, PlotSpec)` or `(label, PlotSpec, color)`. When
    no colour is given, the dataset's palette colour by position is used. All
    specs are assumed to be the same metric (compare overlays one metric across
    files), so axis labels/ranges come from the first spec.
    """
    self._reset_plot()
    if not specs:
      self._show_empty()
      return

    specs = [self._normalise(s, i) for i, s in enumerate(specs)]
    base_axes = specs[0][1].axes

    # Heatmaps do not overlay: render only the first dataset's heatmap.
    if specs[0][1].is_heatmap:
      label, spec, _ = specs[0]
      self._render_heatmap(spec, view)
      if len(specs) > 1:
        self.set_status(f"{spec.title or label}: spectrogram shows one track at a time")
      self._apply_axes(base_axes, single=True, log_y_image_handled=True)
      self._draw_ref_lines()
      return

    single = len(specs) == 1
    for label, spec, color in specs:
      prefix = "" if single else f"{label}: "
      self._render_curves_and_bands(spec, color, prefix, single=single)

    # Reference lines from the first spec only (identical across same-metric specs).
    for hl in specs[0][1].hlines:
      self._render_hline(hl)

    # Scalar readouts → legend-only proxy entries.
    for label, spec, _ in specs:
      prefix = "" if single else f"{label}: "
      for note in spec.annotations:
        self._legend_note(prefix + note)

    self._apply_axes(base_axes, single=single)
    self._draw_ref_lines()

  def set_reference_lines(self, props: list[RefLineProps]):
    """Set the reference-line set (RefLineProps owned by the GUI) and redraw them."""
    self._ref_props = props
    self._draw_ref_lines()

  def current_view_center_value(self) -> float:
    """Natural-unit value at the current y-view centre — default for a new line.

    Runs through `from_pos`, so on the spectrogram this returns a frequency, not
    a row index.
    """
    (_, _), (y0, y1) = self.plot.viewRange()
    return float(self._ref_from_pos((y0 + y1) / 2.0))

  def set_status(self, message: str):
    self.status_label.setText(message)

  # ---- rendering helpers --------------------------------------------------

  def _normalise(self, spec_tuple, index: int):
    """Coerce a spec tuple to (label, PlotSpec, color), filling colour by index."""
    if len(spec_tuple) == 3:
      return spec_tuple
    label, spec = spec_tuple
    return label, spec, dataset_color(index)

  def _render_curves_and_bands(self, spec: PlotSpec, color: str, prefix: str, single: bool):
    for band in spec.bands:
      lo = np.ascontiguousarray(np.broadcast_to(band.lo, band.x.shape), dtype=float)
      hi = np.ascontiguousarray(np.broadcast_to(band.hi, band.x.shape), dtype=float)
      # FillBetweenItem fills nothing if its child curves have no pen — give them
      # a thin outline in the dataset colour (this is the RMS/Waveform fix).
      edge = pg.mkPen(color, width=1.0)
      c_lo = pg.PlotDataItem(band.x, lo, pen=edge)
      c_hi = pg.PlotDataItem(band.x, hi, pen=edge)
      self.plot.addItem(c_lo)
      self.plot.addItem(c_hi)
      # Build the colour with alpha up front: QBrush.color() returns a copy, so
      # mutating its alpha after mkBrush would be a no-op (opaque overlay bug).
      fill_color = pg.mkColor(color)
      fill_color.setAlpha(200 if single else 90)
      fill = pg.FillBetweenItem(c_lo, c_hi, brush=pg.mkBrush(fill_color))
      self.plot.addItem(fill)
      if band.label:
        self._legend_swatch(prefix + band.label, color)

    for curve in spec.curves:
      pen = pg.mkPen(curve.color or color, width=curve.width)
      item = self.plot.plot(curve.x, curve.y, pen=pen,
                            name=(prefix + curve.label) if curve.label else None,
                            connect="finite")  # gaps at NaN (gated PSR)
      item.setDownsampling(auto=True)  # keep big series smooth under zoom
      item.setClipToView(True)

  def _render_hline(self, hl):
    pen = pg.mkPen(hl.color, width=hl.width, style=_PEN_STYLE.get(hl.style, Qt.DotLine))
    line = pg.InfiniteLine(
      pos=hl.y, angle=0, pen=pen, movable=False,
      label=hl.label or None,
      labelOpts={"position": 0.95, "color": hl.color, "fill": (255, 255, 255, 150)},
    )
    self.plot.addItem(line)

  def _render_heatmap(self, spec: PlotSpec, view: ViewState):
    hm = spec.heatmap
    t0, t1 = float(hm.x[0]), float(hm.x[-1])
    f_lo = max(spec.axes.y_range[0] if spec.axes.y_range else hm.y[0], hm.y[0])
    f_hi = spec.axes.y_range[1] if spec.axes.y_range else hm.y[-1]
    y_log = view.resolve_y_log(default=spec.axes.y_log)

    n_rows = len(hm.y)
    if y_log:
      f_grid = np.logspace(np.log10(max(f_lo, 1e-6)), np.log10(f_hi), n_rows)
    else:
      f_grid = np.linspace(f_lo, f_hi, n_rows)

    # Resample every time column from native linear freq bins onto f_grid in one
    # vectorised pass — this runs on each redraw and lin/log toggle, so the loop
    # version would make the toggle feel laggy on long files.
    interp = interp1d(hm.y, hm.z, axis=0, bounds_error=False,
                      fill_value=(hm.z[0], hm.z[-1]), assume_sorted=True)
    z_grid = interp(f_grid).astype(np.float32)

    img = pg.ImageItem()
    img.setImage(z_grid.T, autoLevels=False)  # ImageItem wants (x, y) -> transpose
    img.setLevels((hm.z_min, hm.z_max))
    img.setColorMap(_colormap(hm.cmap))
    # Map image pixel space (time cols, freq rows) to data coords: x=time, y=row index.
    img.setRect(pg.QtCore.QRectF(t0, 0.0, t1 - t0, float(n_rows)))
    self.plot.addItem(img)

    # Reference lines on the spectrogram are entered/shown in Hz but drawn at a
    # row index — install the frequency<->row transforms for this f_grid.
    rows = np.arange(n_rows)
    self._ref_to_pos = lambda hz, fg=f_grid, r=rows: float(np.interp(hz, fg, r))
    self._ref_from_pos = lambda pos, fg=f_grid, r=rows: float(np.interp(pos, r, fg))
    self._ref_fmt = lambda v: f"{v:.0f} Hz"

    # Label the row-index y-axis with real frequencies.
    ticks = []
    for hz in _LOG_FREQ_TICKS:
      if f_lo <= hz <= f_hi:
        row = float(np.searchsorted(f_grid, hz))
        ticks.append((row, _fmt_hz(hz)))
    self.plot.getAxis("left").setTicks([ticks])
    self.plot.setYRange(0, n_rows, padding=0)
    self.plot.setXRange(t0, t1, padding=0)

    # Place the colourbar at a fixed layout cell and link it to the image. We
    # add/remove it ourselves (rather than insert_in=) so it can't stack across
    # repeated spectrogram renders.
    self._colorbar = pg.ColorBarItem(values=(hm.z_min, hm.z_max),
                                      colorMap=_colormap(hm.cmap), label=hm.label)
    self._colorbar.setImageItem(img)
    self.glw.addItem(self._colorbar, row=0, col=1)

  def _apply_axes(self, axes, single: bool, log_y_image_handled: bool = False):
    self.plot.setLabel("bottom", axes.x_label)
    self.plot.setLabel("left", axes.y_label)
    # Frame x exactly only for a single dataset; overlaid tracks of different
    # lengths (absolute mode) should autorange to their union rather than clip to
    # the first one's span. In relative mode every spec is 0-100, so either works.
    if axes.x_range and single:
      self.plot.setXRange(*axes.x_range, padding=0)
    elif not single:
      self.plot.enableAutoRange(axis=pg.ViewBox.XAxis)
    if axes.y_range and not log_y_image_handled:
      self.plot.setYRange(*axes.y_range, padding=0)
    if not log_y_image_handled:
      # Curve metrics: honour log mode if a spec ever opts in (none do today).
      self.plot.setLogMode(x=axes.x_log, y=axes.y_log)

  # ---- user reference lines -----------------------------------------------

  def _draw_ref_lines(self):
    """(Re)create draggable lines from the current RefLineProps set."""
    self._remove_ref_line_items()
    for idx, props in enumerate(self._ref_props):
      line = _RefLine(idx, props, self._ref_to_pos, self._ref_from_pos,
                      self._ref_fmt, on_moved=self.referenceLineMoved.emit)
      self.plot.addItem(line)
      self._ref_lines.append(line)

  def _remove_ref_line_items(self):
    for line in self._ref_lines:
      self.plot.removeItem(line)
    self._ref_lines.clear()

  # ---- legend / lifecycle -------------------------------------------------

  def _legend_swatch(self, name: str, color: str):
    self.legend.addItem(pg.PlotDataItem(pen=pg.mkPen(color, width=3)), name)

  def _legend_note(self, text: str):
    self.legend.addItem(pg.PlotDataItem(pen=None), text)

  def _reset_plot(self):
    self._remove_ref_line_items()  # cleared from scene; props persist for redraw
    self.plot.clear()
    if self._colorbar is not None:
      try:
        self.glw.removeItem(self._colorbar)
      except Exception:
        pass
      self._colorbar = None
    self.legend.clear()
    self.plot.getAxis("left").setTicks(None)  # drop heatmap freq ticks
    self.plot.setLogMode(x=False, y=False)
    # Back to identity; the heatmap path reinstalls Hz<->row if needed.
    self._ref_to_pos = lambda v: v
    self._ref_from_pos = lambda p: p
    self._ref_fmt = lambda v: f"{v:.2f}"

  def _show_empty(self):
    text = pg.TextItem("Drop an audio file to see analysis", anchor=(0.5, 0.5),
                       color=(120, 120, 120))
    self.plot.addItem(text)
    self.plot.setXRange(0, 1)
    self.plot.setYRange(0, 1)
    text.setPos(0.5, 0.5)
    self.set_status("Ready for audio analysis...")