Refactoring, fixup of missing variables in animation callback

.gitignore

@@ -1,3 +1,4 @@
-env/*
+env/
 *.swp
-
+__pycache__/
+.idea/

mic_read.py

@@ -7,63 +7,61 @@
 
 Dependencies: pyaudio, numpy and matplotlib
 """
-############### Import Libraries ###############
+
+from typing import Tuple
+
 import pyaudio
 import numpy as np
 import matplotlib.pyplot as plt
 
-############### Constants ###############
-#RATE = 44100 #sample rate
-RATE = 16000
-FORMAT = pyaudio.paInt16 #conversion format for PyAudio stream
-CHANNELS = 1 #microphone audio channels
-CHUNK_SIZE = 8192 #number of samples to take per read
-SAMPLE_LENGTH = int(CHUNK_SIZE*1000/RATE) #length of each sample in ms
 
-############### Functions ###############
-"""
-open_mic:
-creates a PyAudio object and initializes the mic stream
-inputs: none
-ouputs: stream, PyAudio object
-"""
-def open_mic():
+RATE = 16_000  # sample rate
+FORMAT = pyaudio.paInt16  # conversion format for PyAudio stream
+CHANNELS = 1  # microphone audio channels
+CHUNK_SIZE = 8_192  # number of samples to take per read
+SAMPLE_LENGTH = int(CHUNK_SIZE*1_000/RATE)  # length of each sample in ms
+
+
+def open_mic() -> Tuple[pyaudio.Stream, pyaudio.PyAudio]:
+    """
+    creates a PyAudio object and initializes the mic stream
+    inputs: none
+    ouputs: stream, PyAudio object
+    """
     pa = pyaudio.PyAudio()
-    stream = pa.open(format = FORMAT,
-                     channels = CHANNELS,
-                     rate = RATE,
-                     input = True,
-                     frames_per_buffer = CHUNK_SIZE)
-    return stream,pa
+    stream = pa.open(input=True,
+                     format=FORMAT,
+                     channels=CHANNELS,
+                     rate=RATE,
+                     frames_per_buffer=2*CHUNK_SIZE)
+    return stream, pa
 
-"""
-get_data:
-reads from the audio stream for a constant length of time, converts it to data
-inputs: stream, PyAudio object
-outputs: int16 data array
-"""
-def get_data(stream,pa):
+
+def get_data(stream: pyaudio.Stream) -> np.ndarray:
+    """
+    reads from the audio stream for a constant length of time, converts it to data
+    inputs: stream, PyAudio object
+    outputs: int16 data array
+    """
     input_data = stream.read(CHUNK_SIZE)
-    data = np.fromstring(input_data,np.int16)
+    data = np.frombuffer(input_data, np.int16)
     return data
 
-############### Test Functions ###############
-"""
-make_10k:
-creates a 10kHz test tone
-"""
-def make_10k():
-    x = np.linspace(-2*np.pi,2*np.pi,21000)
-    x = np.tile(x,int(SAMPLE_LENGTH/(4*np.pi)))
-    y = np.sin(2*np.pi*5000*x)
-    return x,y
 
-"""
-show_freq:
-plots the test tone for a sanity check
-"""
+def make_10k() -> Tuple[np.ndarray, np.ndarray]:
+    """
+    creates a 10kHz test tone
+    """
+    x = np.linspace(-2*np.pi, 2*np.pi, 21_000)
+    x = np.tile(x, int(SAMPLE_LENGTH/(4*np.pi)))
+    y = np.sin(2*np.pi*5_000*x)
+    return x, y
+
+
 def show_freq():
-    x,y = make_10k()
-    plt.plot(x,y)
+    """
+    plots the test tone for a sanity check
+    """
+    x, y = make_10k()
+    plt.plot(x, y)
     plt.show()
-

requirements.txt

@@ -1,11 +1,3 @@
-funcsigs==0.4
-matplotlib==1.4.3
-mock==1.3.0
-nose==1.3.7
-numpy==1.9.2
-pbr==1.4.0
-PyAudio==0.2.8
-pyparsing==2.0.3
-python-dateutil==2.4.2
-pytz==2015.4
-six==1.9.0
+matplotlib
+numpy
+PyAudio

run_specgram.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python3
+
 """
 run_specgram.py
 Created By Alexander Yared (akyared@gmail.com)
@@ -7,103 +9,119 @@
 
 Dependencies: matplotlib, numpy and the mic_read.py module
 """
-############### Import Libraries ###############
-from matplotlib.mlab import window_hanning,specgram
+
 import matplotlib.pyplot as plt
-import matplotlib.animation as animation
-from matplotlib.colors import LogNorm
 import numpy as np
+import pyaudio
+from matplotlib.animation import FuncAnimation
+from matplotlib.colors import LogNorm
+from matplotlib.image import AxesImage
+from matplotlib.mlab import window_hanning, specgram
+from typing import Tuple
 
-############### Import Modules ###############
-import mic_read
+from mic_read import SAMPLE_LENGTH, RATE, get_data, open_mic
 
-############### Constants ###############
-#SAMPLES_PER_FRAME = 10 #Number of mic reads concatenated within a single window
-SAMPLES_PER_FRAME = 4
-nfft = 1024#256#1024 #NFFT value for spectrogram
-overlap = 1000#512 #overlap value for spectrogram
-rate = mic_read.RATE #sampling rate
+SAMPLES_PER_FRAME = 4  # Number of mic reads concatenated within a single window
+N_FFT = 1_024  # NFFT value for spectrogram
+OVERLAP = 1_000  # overlap value for spectrogram
 
-############### Functions ###############
-"""
-get_sample:
-gets the audio data from the microphone
-inputs: audio stream and PyAudio object
-outputs: int16 array
-"""
-def get_sample(stream,pa):
-    data = mic_read.get_data(stream,pa)
-    return data
-"""
-get_specgram:
-takes the FFT to create a spectrogram of the given audio signal
-input: audio signal, sampling rate
-output: 2D Spectrogram Array, Frequency Array, Bin Array
-see matplotlib.mlab.specgram documentation for help
-"""
-def get_specgram(signal,rate):
-    arr2D,freqs,bins = specgram(signal,window=window_hanning,
-                                Fs = rate,NFFT=nfft,noverlap=overlap)
-    return arr2D,freqs,bins
 
-"""
-update_fig:
-updates the image, just adds on samples at the start until the maximum size is
-reached, at which point it 'scrolls' horizontally by determining how much of the
-data needs to stay, shifting it left, and appending the new data. 
-inputs: iteration number
-outputs: updated image
-"""
-def update_fig(n):
-    data = get_sample(stream,pa)
-    arr2D,freqs,bins = get_specgram(data,rate)
+def get_sample(stream: pyaudio.Stream) -> np.ndarray:
+    """
+    gets the audio data from the microphone
+    inputs: audio stream and PyAudio object
+    outputs: int16 array
+    """
+    return get_data(stream)
+
+
+def get_specgram(signal: np.ndarray) -> Tuple[
+    np.ndarray,  # 2D spectrum
+    np.ndarray,  # Frequency axis
+    np.ndarray,  # Time axis
+]:
+    """
+    takes the FFT to create a spectrogram of the given audio signal
+    input: audio signal, sampling rate
+    output: 2D Spectrogram Array, Frequency Array, Bin Array
+    see matplotlib.mlab.specgram documentation for help
+    """
+    return specgram(
+        signal,
+        Fs=RATE, NFFT=N_FFT, noverlap=OVERLAP,
+        window=window_hanning,
+    )
+
+
+def update_fig(frame: int, im: AxesImage, stream: pyaudio.Stream) -> Tuple[AxesImage]:
+    """
+    updates the image, just adds on samples at the start until the maximum size is
+    reached, at which point it 'scrolls' horizontally by determining how much of the
+    data needs to stay, shifting it left, and appending the new data.
+    inputs: iteration number
+    outputs: updated image
+    """
+    data = get_sample(stream)
+    arr_2d, freqs, times = get_specgram(data)
     im_data = im.get_array()
-    if n < SAMPLES_PER_FRAME:
-        im_data = np.hstack((im_data,arr2D))
+
+    # frame cannot be relied upon: we're called multiple times with 0 before it
+    # starts to increment.
+    frame = im_data.shape[1] // len(times)
+
+    if frame < SAMPLES_PER_FRAME:
+        im_data = np.hstack((im_data, arr_2d))
         im.set_array(im_data)
     else:
-        keep_block = arr2D.shape[1]*(SAMPLES_PER_FRAME - 1)
-        im_data = np.delete(im_data,np.s_[:-keep_block],1)
-        im_data = np.hstack((im_data,arr2D))
+        im_data = np.hstack((
+            im_data[:, len(times):],
+            arr_2d,
+        ))
         im.set_array(im_data)
+
     return im,
 
-def main():
-    ############### Initialize Plot ###############
+
+def make_plot(stream: pyaudio.Stream) -> FuncAnimation:
+    # Initialize Plot
     fig = plt.figure()
-    """
-    Launch the stream and the original spectrogram
-    """
-    stream,pa = mic_read.open_mic()
-    data = get_sample(stream,pa)
-    arr2D,freqs,bins = get_specgram(data,rate)
-    """
-    Setup the plot paramters
-    """
-    extent = (bins[0],bins[-1]*SAMPLES_PER_FRAME,freqs[-1],freqs[0])
-    im = plt.imshow(arr2D,aspect='auto',extent = extent,interpolation="none",
-                    cmap = 'jet',norm = LogNorm(vmin=.01,vmax=1))
-    plt.xlabel('Time (s)')
-    plt.ylabel('Frequency (Hz)')
-    plt.title('Real Time Spectogram')
-    plt.gca().invert_yaxis()
-    ##plt.colorbar() #enable if you want to display a color bar
-
-    ############### Animate ###############
-    anim = animation.FuncAnimation(fig,update_fig,blit = False,
-                                interval=mic_read.CHUNK_SIZE/1000)
-
-
+    ax = fig.gca()
+
+    # Data for first frame
+    data = get_sample(stream)
+    arr_2d, freqs, times = get_specgram(data)
+
+    # Set up the plot parameters
+    extent = (times[0], times[-1]*SAMPLES_PER_FRAME, freqs[-1], freqs[0])
+    im = ax.imshow(arr_2d, aspect='auto', extent=extent, interpolation='none',
+                   cmap='jet', norm=LogNorm(vmin=.01, vmax=1))
+    ax.set_xlabel('Time (s)')
+    ax.set_ylabel('Frequency (Hz)')
+    ax.set_title('Real-Time Spectogram')
+    ax.invert_yaxis()
+    # fig.colorbar(im)  # enable if you want to display a color bar
+
+    # Animate
+    return FuncAnimation(
+        fig,
+        func=update_fig, fargs=(im, stream),
+        interval=SAMPLE_LENGTH,
+        blit=True,
+    )
+
+
+def main():
+    # Launch the stream and the original spectrogram
+    stream, pa = open_mic()
+
     try:
+        animation = make_plot(stream)
         plt.show()
-    except:
-        print("Plot Closed")
+    finally:
+        stream.stop_stream()
+        stream.close()
+        pa.terminate()
 
-    ############### Terminate ###############
-    stream.stop_stream()
-    stream.close()
-    pa.terminate()
-    print("Program Terminated")
 
-if __name__ == "__main__":
+if __name__ == '__main__':
     main()