pysimplegui_plot_spec_envelop.py

#!/usr/bin/env python3

"""PyAudioとPySimpleGUIを使ったリアルタイムスペクトル包絡表示."""

# Copyright (C) 2022 by Akira TAMAMORI

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

# 参考: https://gist.github.com/KenoLeon/c4c32ee7d0b133b24211a2250d56fdce

import numpy as np
import pyaudio
import PySimpleGUI as sg
import pyworld

CHUNK = 1024
RATE = 16000
CHANNEL = 1

pAud = pyaudio.PyAudio()
SHORT_MAX = 32767
SHORT_MIN = -32768
EPSIRON = 1.0e-6

FONT = "Any 16"
WINSIZE = (512, 300)
LAYOUT = [
    [
        sg.Graph(
            canvas_size=WINSIZE,
            graph_bottom_left=(0, -10),
            graph_top_right=(int(CHUNK / 2), 200),
            background_color="white",
            key="-GRAPH-",
        )
    ],
    [
        sg.Button("収録開始", key="-LISTEN-", font=FONT),
        sg.Button("収録停止", key="-STOP-", font=FONT, disabled=True),
        sg.Button("終了", key="-EXIT-", font=FONT),
    ],
]
WINDOW = sg.Window("Waveform plot", LAYOUT, finalize=True)
GRAPH = WINDOW["-GRAPH-"]
STREAM = None
AUDIODATA = np.array([])
TIMEOUT = 10  # 10 msec for redraw


def stop():
    """音声収録を中止する."""
    if STREAM is not None:
        STREAM.stop_stream()
        STREAM.close()
        WINDOW.FindElement("-STOP-").Update(disabled=True)
        WINDOW.FindElement("-LISTEN-").Update(disabled=False)


def callback(in_data, frame_count, time_info, status):
    """PyAudioのstream経由でデータを取得し、プロット用変数に格納する."""
    global AUDIODATA
    AUDIODATA = np.frombuffer(in_data, dtype=np.int16)
    return (in_data, pyaudio.paContinue)


def listen():
    """音声収録を開始する."""

    WINDOW.FindElement("-STOP-").Update(disabled=False)
    WINDOW.FindElement("-LISTEN-").Update(disabled=True)

    stream = pAud.open(
        format=pyaudio.paInt16,
        channels=CHANNEL,
        rate=RATE,
        input=True,
        frames_per_buffer=CHUNK,
        stream_callback=callback,
    )
    stream.start_stream()
    return stream


def plotWaveform():
    """スペクトル包絡をグラフに描画する."""
    GRAPH.Erase()  # 再描画

    # 軸を描く
    GRAPH.DrawLine((0, 0), (1024, 0))
    GRAPH.DrawLine((0, SHORT_MIN), (0, SHORT_MIN))

    # 軸のラベルを描く
    GRAPH.DrawText("kHz", location=(500, 5), color="black")

    # スペクトル包絡の計算
    _, sp, _ = pyworld.wav2world(AUDIODATA.astype(np.float64), RATE)
    envelop = np.mean(sp, axis=0)
    envelop = 20 * np.log10(envelop + EPSIRON)

    prev_x = prev_y = None
    for x, y in enumerate(envelop):
        if prev_x is not None:
            # 前の点と現在の点を線で結ぶ
            GRAPH.draw_line((prev_x, prev_y), (x, y), color="red")
        prev_x, prev_y = x, y


while True:
    event, values = WINDOW.read(timeout=TIMEOUT)

    if event in (sg.WIN_CLOSED, "-EXIT-"):
        stop()
        pAud.terminate()
        break
    if event == "-LISTEN-":
        STREAM = listen()
    elif event == "-STOP-":
        stop()
    elif AUDIODATA.size != 0:
        plotWaveform()

WINDOW.close()