From 90251a8a0f3cfb3bd29b56bbb7313ed03fd98340 Mon Sep 17 00:00:00 2001 From: Saveliy Yusufov Date: Sun, 9 Jun 2024 14:46:17 -0400 Subject: [PATCH] Add python bindings for Real FFT --- pyphastft/src/lib.rs | 23 ++++- pyphastft/vis_qt.py | 208 +++++++++++++++++++++---------------------- 2 files changed, 123 insertions(+), 108 deletions(-) diff --git a/pyphastft/src/lib.rs b/pyphastft/src/lib.rs index f69b64f..ae1ebbc 100644 --- a/pyphastft/src/lib.rs +++ b/pyphastft/src/lib.rs @@ -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] @@ -18,9 +18,28 @@ fn fft(mut reals: PyReadwriteArray1, mut imags: PyReadwriteArray1, dir ); } +#[pyfunction] +fn rfft(reals: PyReadonlyArray1, direction: char) -> (Vec, Vec) { + 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(()) } diff --git a/pyphastft/vis_qt.py b/pyphastft/vis_qt.py index 2a5c5e7..2396ebf 100644 --- a/pyphastft/vis_qt.py +++ b/pyphastft/vis_qt.py @@ -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_()) + +