Display waveform of an audio file using PyQt/QPainter

waveformqt.py

#!/usr/bin/env python
"""Display waveform of an audio file."""

import argparse
import logging
import sys
from statistics import fmean

from PyQt6 import QtCore, QtGui, QtWidgets
from PyQt6.QtCore import Qt

import soundfile

PROG = "waveform"
log = logging.getLogger(PROG)


class WaveformDisplay(QtWidgets.QWidget):
    """Custom widget for waveform representation of a digital audio signal."""

    def __init__(self, frames=None, channels=1, samplerate=48000, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._sampleframes = frames
        self._channels = channels
        self._samplerate = samplerate
        self.waveform_color = QtGui.QColor(255, 165, 0, 200)  # tranparent orange
        # ~ self.waveform_color = QtGui.QColor(255, 255, 255, 160)
        self.background_color = QtGui.QColor('black')
        self.background_gradient = QtGui.QLinearGradient()
        self.background_gradient.setColorAt(0, QtGui.QColor('black'))
        self.background_gradient.setColorAt(1, QtGui.QColor('#434343'))
        self.background_gradient.setSpread(QtGui.QGradient.Spread.ReflectSpread)
        self.foreground_color = QtGui.QColor('white')
        self._zoom = 1.0
        self._startframe = 0
        self._endframe = -1

    def paintEvent(self, event):
        painter = QtGui.QPainter()
        painter.begin(self)
        self.draw_waveform(painter)
        painter.end()

    def draw_waveform(self, painter):
        pen = painter.pen()
        height = painter.device().height()
        zero_y = float(height) / 2
        width = painter.device().width()
        num_frames = len(self._sampleframes[self._startframe:self._endframe])
        samples_per_pixel = num_frames / float(width)

        # draw background
        # ~ brush = QtGui.QBrush()
        # ~ brush.setColor(self.background_color)
        # ~ brush.setStyle(Qt.BrushStyle.SolidPattern)
        self.background_gradient.setStart(0.0, zero_y)
        self.background_gradient.setFinalStop(0.0, 0.0)
        rect = QtCore.QRect(0, 0, width, height)
        painter.fillRect(rect, self.background_gradient)
        startframe = self._startframe

        #breakpoint()
        # draw waveform
        if self._sampleframes is not None:
            pen.setColor(self.waveform_color)
            painter.setPen(pen)
            for pixel in range(width):
                offset = round(pixel * samples_per_pixel)

                if 0 <= offset < num_frames:
                    start = startframe + offset
                    end = max(start + 1, start + int(samples_per_pixel))
                    values = self._sampleframes[start:end]
                    max_value = max(values)
                    min_value = min(values)

                    if max_value > 0:
                        y = zero_y - zero_y * max_value
                        painter.drawLine(QtCore.QLineF(pixel, zero_y, pixel, y))
                    if min_value < 0:
                        y = zero_y - zero_y * min_value
                        painter.drawLine(QtCore.QLineF(pixel, zero_y, pixel, y))

        # draw zero line
        pen.setColor(self.foreground_color)
        pen.setStyle(Qt.PenStyle.DotLine)
        painter.setPen(pen)
        painter.drawLine(QtCore.QLineF(0.0, zero_y, float(width), zero_y))

    def set_samples(self, frames, channels=1, samplerate=48000):
        self._sampleframes = frames
        self._channels = channels
        self._samplerate = samplerate
        self.update()


class SampleManglerMainWindow(QtWidgets.QMainWindow):
    def __init__(self, *args, **kw):
        super().__init__(*args, **kw)
        self.waveform = WaveformDisplay()
        self.setCentralWidget(self.waveform)
        self.setWindowTitle(PROG)
        self.setMinimumSize(600, 400)

    def set_samples(self, frames, channels=1, samplerate=48000):
        self.waveform.set_samples(frames=frames, channels=channels, samplerate=samplerate)


def main(args=None):
    ap = argparse.ArgumentParser(prog=PROG, usage=__doc__.splitlines()[0])
    ap.add_argument("-v", "--verbose", action="store_true", help="Display verbose messages")
    ap.add_argument("audiofile", help="Audio file to display")
    options, args = ap.parse_known_args(args)

    logging.basicConfig(level=logging.DEBUG if options.verbose else logging.INFO,
        format="%(levelname)s: %(message)s")

    app = QtWidgets.QApplication(args)
    mainwin = SampleManglerMainWindow()

    sf = soundfile.SoundFile(options.audiofile, "rb")

    mainwin.set_samples(frames=sf.read(), channels=sf.channels, samplerate=sf.samplerate)
    mainwin.show()
    app.exec()


if __name__ == '__main__':
    sys.exit(main() or 0)