Plot a colored spectrum analyzer

This commit is contained in:
Correl Roush 2020-09-04 23:04:56 -04:00
parent 777ba43d93
commit f2a2424405

View file

@ -1,3 +1,4 @@
from bisect import bisect
import logging import logging
from multiprocessing import Queue from multiprocessing import Queue
import os import os
@ -24,8 +25,6 @@ class Plot:
width: int, width: int,
height: int, height: int,
bars: int = 20, bars: int = 20,
bar_width: int = 40,
color: Tuple[int, int, int] = (255, 255, 255),
) -> None: ) -> None:
self.screen = screen self.screen = screen
self.x = x self.x = x
@ -33,8 +32,6 @@ class Plot:
self.width = width self.width = width
self.height = height self.height = height
self.bars = bars self.bars = bars
self.bar_width = bar_width
self.color = color
self.audio = b"" self.audio = b""
def draw(self) -> None: def draw(self) -> None:
@ -43,19 +40,33 @@ class Plot:
return return
fft = abs(np.fft.fft(data).real) fft = abs(np.fft.fft(data).real)
fft = fft[: len(fft) // 2] fft = fft[: len(fft) // 2]
heights = abs(scipy.signal.resample(fft, self.bars) * self.height / 2 ** 16) fft = abs(scipy.signal.resample(fft, self.bars))
for i, height in enumerate(heights):
pygame.draw.rect( light_width = self.width // (2 * self.bars - 1)
self.screen, light_height = self.height // 2 // light_width
self.color, light_steps = self.height // light_height // 2
( lights = fft * light_steps / 2 ** 16
self.x + i / self.bars * self.width,
self.height, colors = [
self.bar_width, (0, (0, 255, 0)),
-height, (50, (255, 255, 0)),
), (75, (255, 0, 0)),
0, ]
) color_keys = [k for k, v in colors]
color_values = [v for k, v in colors]
for i, steps in enumerate(lights):
for step in range(int(steps)):
color = color_values[bisect(color_keys, step / light_steps * 100) - 1]
pygame.draw.rect(
self.screen,
color,
(
self.x + i * light_width * 2,
self.height - step * light_height * 2,
light_width,
-light_height,
),
)
def main(): def main():
@ -107,9 +118,7 @@ def main():
y=0, y=0,
width=screen.get_width(), width=screen.get_width(),
height=screen.get_height() - 50, height=screen.get_height() - 50,
bars=40, bars=15,
bar_width=screen.get_width() // 45,
color=(139, 0, 139),
) )
try: try: