Skip to content

Commit

Permalink
Add python bindings for Real FFT
Browse files Browse the repository at this point in the history
  • Loading branch information
@smu160
smu160 committed Jun 9, 2024
1 parent 8cfa650 commit 90251a8
Show file tree
Hide file tree
Showing 2 changed files with 123 additions and 108 deletions.
23 changes: 21 additions & 2 deletions pyphastft/src/lib.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use numpy::PyReadwriteArray1;
use phastft::{fft_64 as fft_64_rs, planner::Direction};
use numpy::{PyReadonlyArray1, PyReadwriteArray1};
use phastft::{fft_64 as fft_64_rs, fft::r2c_fft_f64, planner::Direction};
use pyo3::prelude::*;

#[pyfunction]
Expand All @@ -18,9 +18,28 @@ fn fft(mut reals: PyReadwriteArray1<f64>, mut imags: PyReadwriteArray1<f64>, dir
);
}

#[pyfunction]
fn rfft(reals: PyReadonlyArray1<f64>, direction: char) -> (Vec<f64>, Vec<f64>) {
assert!(direction == 'f' || direction == 'r');
let _dir = if direction == 'f' {
Direction::Forward
} else {
Direction::Reverse
};

let big_n = reals.as_slice().unwrap().len();

let mut output_re = vec![0.0; big_n];
let mut output_im = vec![0.0; big_n];
r2c_fft_f64(reals.as_slice().unwrap(), &mut output_re, &mut output_im);
(output_re, output_im)
}


/// A Python module implemented in Rust.
#[pymodule]
fn pyphastft(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(fft, m)?)?;
m.add_function(wrap_pyfunction!(rfft, m)?)?;
Ok(())
}
208 changes: 102 additions & 106 deletions pyphastft/vis_qt.py
View file
Original file line number Diff line number Diff line change
@@ -1,115 +1,111 @@
import sys

import numpy as np
import pyaudio
import pyqtgraph as pg
from pyphastft import fft
from pyphastft import rfft
from pyqtgraph.Qt import QtWidgets, QtCore


class RealTimeAudioSpectrum(QtWidgets.QWidget):
def __init__(self, parent=None):
super(RealTimeAudioSpectrum, self).__init__(parent)
self.n_fft_bins = 1024 # Increased FFT size for better frequency resolution
self.n_display_bins = 32 # Maintain the same number of bars in the display
self.sample_rate = 44100
self.smoothing_factor = 0.1 # Smaller value for more smoothing
self.ema_fft_data = np.zeros(
self.n_display_bins
) # Adjusted to the number of display bins
self.init_ui()
self.init_audio_stream()

def init_ui(self):
self.layout = QtWidgets.QVBoxLayout(self)
self.plot_widget = pg.PlotWidget()
self.layout.addWidget(self.plot_widget)

# Customize plot aesthetics
self.plot_widget.setBackground("k")
self.plot_item = self.plot_widget.getPlotItem()
self.plot_item.setTitle(
"Real-Time Audio Spectrum Visualizer powered by PhastFT",
color="w",
size="16pt",
)

# Hide axis labels
self.plot_item.getAxis("left").hide()
self.plot_item.getAxis("bottom").hide()

# Set fixed ranges for the x and y axes to prevent them from jumping
self.plot_item.setXRange(0, self.sample_rate / 2, padding=0)
self.plot_item.setYRange(0, 1, padding=0)

self.bar_width = (
(self.sample_rate / 2) / self.n_display_bins * 0.90
) # Adjusted width for display bins

# Calculate bar positions so they are centered with respect to their frequency values
self.freqs = np.linspace(
0 + self.bar_width / 2,
self.sample_rate / 2 - self.bar_width / 2,
self.n_display_bins,
)

self.bar_graph = pg.BarGraphItem(
x=self.freqs,
height=np.zeros(self.n_display_bins),
width=self.bar_width,
brush=pg.mkBrush("m"),
)
self.plot_item.addItem(self.bar_graph)

self.timer = QtCore.QTimer()
self.timer.timeout.connect(self.update)
self.timer.start(50) # Update interval in milliseconds

def init_audio_stream(self):
self.p = pyaudio.PyAudio()
self.stream = self.p.open(
format=pyaudio.paFloat32,
channels=1,
rate=self.sample_rate,
input=True,
frames_per_buffer=self.n_fft_bins, # This should match the FFT size
stream_callback=self.audio_callback,
)
self.stream.start_stream()

def audio_callback(self, in_data, frame_count, time_info, status):
audio_data = np.frombuffer(in_data, dtype=np.float32)
reals = np.zeros(self.n_fft_bins)
imags = np.zeros(self.n_fft_bins)
reals[: len(audio_data)] = audio_data # Fill the reals array with audio data
fft(reals, imags, direction="f")
fft_magnitude = np.sqrt(reals**2 + imags**2)[: self.n_fft_bins // 2]

# Aggregate or interpolate FFT data to fit into display bins
new_fft_data = np.interp(
np.linspace(0, len(fft_magnitude), self.n_display_bins),
np.arange(len(fft_magnitude)),
fft_magnitude,
)

# Apply exponential moving average filter
self.ema_fft_data = self.ema_fft_data * self.smoothing_factor + new_fft_data * (
1.0 - self.smoothing_factor
)
return in_data, pyaudio.paContinue

def update(self):
self.bar_graph.setOpts(height=self.ema_fft_data, width=self.bar_width)

def closeEvent(self, event):
self.stream.stop_stream()
self.stream.close()
self.p.terminate()
event.accept()

def __init__(self, parent=None):
super(RealTimeAudioSpectrum, self).__init__(parent)
self.n_fft_bins = 1024
self.n_display_bins = 64
self.sample_rate = 44100
self.smoothing_factor = 0.1 # Fine-tuned smoothing factor
self.ema_fft_data = np.zeros(self.n_display_bins)
self.init_ui()
self.init_audio_stream()

def init_ui(self):
self.layout = QtWidgets.QVBoxLayout(self)
self.plot_widget = pg.PlotWidget()
self.layout.addWidget(self.plot_widget)

self.plot_widget.setBackground("k")
self.plot_item = self.plot_widget.getPlotItem()
self.plot_item.setTitle(
"Real-Time Audio Spectrum Visualizer powered by PhastFT",
color="w",
size="16pt",
)

self.plot_item.getAxis("left").hide()
self.plot_item.getAxis("bottom").hide()

self.plot_item.setXRange(0, self.sample_rate / 2, padding=0)
self.plot_item.setYRange(0, 1, padding=0)

self.bar_width = (self.sample_rate / 2) / self.n_display_bins * 0.8
self.freqs = np.linspace(
0, self.sample_rate / 2, self.n_display_bins, endpoint=False
)

self.bar_graph = pg.BarGraphItem(
x=self.freqs,
height=np.zeros(self.n_display_bins),
width=self.bar_width,
brush=pg.mkBrush("m"),
)
self.plot_item.addItem(self.bar_graph)

self.timer = QtCore.QTimer()
self.timer.timeout.connect(self.update)
self.timer.start(50)

def init_audio_stream(self):
self.p = pyaudio.PyAudio()
self.stream = self.p.open(
format=pyaudio.paFloat32,
channels=1,
rate=self.sample_rate,
input=True,
frames_per_buffer=self.n_fft_bins,
stream_callback=self.audio_callback,
)
self.stream.start_stream()

def audio_callback(self, in_data, frame_count, time_info, status):
audio_data = np.frombuffer(in_data, dtype=np.float32)
audio_data = np.ascontiguousarray(audio_data, dtype=np.float64)
reals, imags = rfft(audio_data, direction="f")
reals = np.ascontiguousarray(reals)
imags = np.ascontiguousarray(imags)
fft_magnitude = np.sqrt(reals**2 + imags**2)[: self.n_fft_bins // 2]

new_fft_data = np.interp(
np.linspace(0, len(fft_magnitude), self.n_display_bins),
np.arange(len(fft_magnitude)),
fft_magnitude,
)

new_fft_data = np.log1p(new_fft_data) # Apply logarithmic scaling

self.ema_fft_data = self.ema_fft_data * self.smoothing_factor + new_fft_data * (
1.0 - self.smoothing_factor
)

return in_data, pyaudio.paContinue

def update(self):
# Normalize the FFT data to ensure it fits within the display range
max_value = np.max(self.ema_fft_data)
if max_value > 0:
normalized_fft_data = self.ema_fft_data / max_value
else:
normalized_fft_data = self.ema_fft_data

self.bar_graph.setOpts(height=normalized_fft_data, width=self.bar_width)

def closeEvent(self, event):
self.stream.stop_stream()
self.stream.close()
self.p.terminate()
event.accept()

if __name__ == "__main__":
app = QtWidgets.QApplication(sys.argv)
window = RealTimeAudioSpectrum()
window.show()
sys.exit(app.exec_())
app = QtWidgets.QApplication(sys.argv)
window = RealTimeAudioSpectrum()
window.show()
sys.exit(app.exec_())


0 comments on commit 90251a8

Please sign in to comment.