Pythonでリアルタイム信号処理(2)

import pyaudio
import numpy as np
import matplotlib.pyplot as plt
import time
import sys
from scipy.signal import hilbert
from scipy.signal import medfilt
from sklearn.cluster import KMeans

CHUNK    = 8001
FORMAT   = pyaudio.paInt16
CHANNELS = 1
RATE     = 8000

x  = range   (CHUNK)
y  = np.zeros(CHUNK)
binary = np.zeros(CHUNK)
xx = range   (CHUNK*2)
yy = np.zeros(CHUNK*2)

fig = plt.figure(1, figsize=(8, 16))
plt.subplots_adjust(hspace=0.9, wspace=0.1)

ax1  = plt.subplot(9,1,1)
ax1.set_xlim(0, CHUNK)
lines1, = plt.plot(x,y)
plt.title('input signal')

ax2  = plt.subplot(9,1,2)
ax2.set_xlim(CHUNK/2,CHUNK+CHUNK/2)
lines2, = plt.plot(xx,yy, color="green")
plt.title('buffered input signal')

ax3  = plt.subplot(9,1,3)
ax3.set_xlim(CHUNK/2,CHUNK+CHUNK/2)
lines3, = plt.plot(xx,yy, color="orange")
plt.title('envelope')

ax4  = plt.subplot(9,1,4)
ax4.set_xlim(CHUNK/2,CHUNK+CHUNK/2)
lines4, = plt.plot(xx,yy, color="firebrick")
plt.title('median filtered')

ax5  = plt.subplot(9,1,5)
ax5.set_xlim(0, CHUNK)
lines5, = plt.plot(x,y, color="maroon")
plt.title('binary signal')

y_buf  = np.zeros(CHUNK*2)

time_base = 0
count = 0
trig_count = 0
t_up       = 0

p = pyaudio.PyAudio()

def callback(in_data, frame_count, time_info, status):
    global y_buf, z_buf, window, yy
    global time_base, count, trig_count, t_up, t_down, t_b4, up, down

    y = np.frombuffer(in_data, dtype="int16")

    for i in range(CHUNK):
        y_buf[i      ] = y_buf[i+CHUNK]
        y_buf[i+CHUNK] = y    [i]

    analytic_signal = hilbert(y_buf)
    envelope = np.abs(analytic_signal)
    nmedian = 51
    medsignal = medfilt(envelope, nmedian)

    y_max          = max(1.1 * max(abs(y.min()),     abs(y.max())),     100.0)
    y_buf_max      = max(1.1 * max(abs(y_buf.min()), abs(y_buf.max())), 100.0)
    envelope_max   = max(1.1 * abs(envelope.max()) , 100.0)
    medsignal_max  = max(1.1 * abs(medsignal.max()), 100.0)
    th = 0.4*medsignal_max

    for i in range(CHUNK):
        if medsignal[i+CHUNK//2] > th:
            binary[i] = 1
        else:
            binary[i] = 0

    print('count =', count)

    if count%10 == 0:
        ax1.set_ylim(-y_max, y_max)
        ax2.set_ylim(-y_buf_max, y_buf_max)
        ax3.set_ylim(-0.1*envelope_max, envelope_max)
        ax4.set_ylim(-0.1*medsignal_max, medsignal_max)
        ax5.set_ylim(-0.1, 1.1)
        lines1.set_data(x,y) # input signal
        lines2.set_data(xx,y_buf)
        lines3.set_data(xx,envelope)
        lines4.set_data(xx,medsignal)
        lines5.set_data(x,binary)
        print('*')

    count += 1

    return (None, pyaudio.paContinue)

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

stream.start_stream()

while stream.is_active():
	plt.pause (0.01)
	time.sleep(0.01)

stream.stop_stream()
stream.close()
p.terminate()