visualise.py

from pydub import AudioSegment
from PyQt5 import QtCore, QtWidgets
from PyQt5.QtMultimedia import *
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import sys
import time
import math


class Visualise(QtCore.QThread):

    def __init__(self, song, player):
        super().__init__()
        self.player = player
        self.samples = song.get_array_of_samples()
        self.song = song

    def run(self):
        interval = 0.05

        fig = plt.figure()
        fig.patch.set_facecolor((0, 0, 0))
        ax1 = fig.add_subplot(1,1,1)
        ax1.set_facecolor((0, 0, 0))

        
        bars_n = 30
        max_bar_height = 35
        frequency_range = 0.3
        
        bars = [0]*bars_n
        for timestamp in range(0, math.ceil(self.song.duration_seconds/interval)):
            start = time.time()
            
            
            timestamp *= interval
            sample_count = int(self.song.frame_rate * interval)
            start_index = int((self.player.position()/1000) * self.song.frame_rate)
            v_sample = self.samples[start_index:start_index+sample_count]

            zero_crossings = np.where(np.diff(np.sign(v_sample)))[0]
            data = []
            for i in range(len(zero_crossings)-1):
                x1 = zero_crossings[i]
                x2 = zero_crossings[i+1]
                diff = abs(x1 - x2)
                frequency = diff / len(v_sample)
                amp = max(max(v_sample[x1:x2]), abs(min(v_sample[x1:x2])))
                data.append((frequency, amp))
            
            
            bar_width_range = frequency_range / bars_n
            bars_samples = []

            if not data:
                time.sleep(max(interval - time.time() + start, 0))
                continue
            
            for f in np.arange(0, frequency_range, bar_width_range):
                amps = [x[1] for x in data if f-bar_width_range<x[0]<f]
                if not amps:
                    bars_samples.append(0)
                else:
                    bars_samples.append(max(amps))

            highest_amp = 18000
            step = highest_amp // max_bar_height
            if step:
                for n, amp in enumerate(bars_samples):
                    if bars[n] > 0 and amp // step < bars[n]:
                            bars[n] -= bars[n] / 3
                            if bars[n] < 1:
                                bars[n] = 0
                    else:
                        bars[n] = amp // step
                        if bars[n] < 1:
                                bars[n] = 0
            else:
                bars = [0]*bars

            ax1.clear()
            ax1.plot(range(len(bars)), bars, color='r')
            ax1.fill_between(range(len(bars)), bars, color='#c60303')
            plt.ylim(top=70, bottom=0)
            plt.draw()
            plt.pause(0.001)
            time.sleep(max(interval - time.time() + start, 0))


class Window(QtWidgets.QWidget):

    def __init__(self):
        super().__init__()
        song = AudioSegment.from_mp3('song.mp3')
        song = song.set_channels(1)

        content = QMediaContent(QtCore.QUrl('song.mp3'))
        self.player = QMediaPlayer()

        self.vis = Visualise(song, self.player)
        
        self.player.setMedia(content)
        self.player.mediaStatusChanged.connect(self.start)

    def start(self):
        self.player.play()
        self.vis.start()
        plt.show()

app = QtWidgets.QApplication(sys.argv)
window = Window()
window.showMinimized()
sys.exit(app.exec_())