Spectrum Analyzer

main.py

from math import sqrt
import numpy as np
import pygame
import pyaudio

pygame.init()

# CD quality typically 44.1kHz
RATE = 44100
# Update screen 60fps
CHUNKS_PER_SECOND = 60  
# Size of each chunk
CHUNK = int(1/CHUNKS_PER_SECOND) * RATE
# 16 bit samples
FORMAT = pyaudio.paInt16

# open up stream from pyaudio
p = pyaudio.PyAudio()

stream = p.open(format=FORMAT, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)

SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((CHUNK, SCREEN_HEIGHT))

done = False
while not done:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            done = True
            break
    # clear screen to black
    screen.fill((0,0,0))
    # draw
    buffer = stream.read(CHUNK)
    waveform = np.frombuffer(buffer, dtype=np.int16)
    fft_complex = np.fft.fft(waveform, n=CHUNK)
    color = (0, 128, 0)
    max_val = sqrt(max(v.real * v.real + v.imag * v.imag for v in fft_complex))
    scale_value = SCREEN_HEIGHT / max_val
    for i, v in enumerate(fft_complex):
        dist = sqrt(v.real * v.real + v.imag * v.imag)
        mapped_dist = dist*scale_value
        pygame.draw.line(screen, color, (i, SCREEN_HEIGHT), (i, SCREEN_HEIGHT-mapped_dist))
    # flip buffer
    pygame.display.flip()