xiaozhi-esp32/scripts/acoustic_check/graphic.py

import sys
import numpy as np
import asyncio
import wave
from collections import deque
import qasync

import matplotlib
matplotlib.use('qtagg')  

from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qtagg import NavigationToolbar2QT as NavigationToolbar  # noqa: F401
from matplotlib.figure import Figure

from PyQt6.QtWidgets import (QApplication, QMainWindow, QVBoxLayout, QWidget, 
                             QHBoxLayout, QLineEdit, QPushButton, QLabel, QTextEdit)
from PyQt6.QtCore import QTimer

# 导入解码器
from demod import RealTimeAFSKDecoder


class UDPServerProtocol(asyncio.DatagramProtocol):
    """UDP服务器协议类"""
    def __init__(self, data_queue):
        self.client_address = None
        self.data_queue: deque = data_queue

    def connection_made(self, transport):
        self.transport = transport
        
    def datagram_received(self, data, addr):
        # 如果还没有客户端地址，记录第一个连接的客户端
        if self.client_address is None:
            self.client_address = addr
            print(f"接受来自 {addr} 的连接")
        
        # 只处理来自已记录客户端的数据
        if addr == self.client_address:
            # 将接收到的音频数据添加到队列
            self.data_queue.extend(data)
        else:
            print(f"忽略来自未知地址 {addr} 的数据")


class MatplotlibWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)

        # 创建 Matplotlib 的 Figure 对象
        self.figure = Figure()

        # 创建 FigureCanvas 对象，它是 Figure 的 QWidget 容器
        self.canvas = FigureCanvas(self.figure)

        # 创建 Matplotlib 的导航工具栏
        # self.toolbar = NavigationToolbar(self.canvas, self)
        self.toolbar = None

        # 创建布局
        layout = QVBoxLayout()
        layout.addWidget(self.toolbar)
        layout.addWidget(self.canvas)
        self.setLayout(layout)

        # 初始化音频数据参数
        self.freq = 16000  # 采样频率
        self.time_window = 20  # 显示时间窗口
        self.wave_data = deque(maxlen=self.freq * self.time_window * 2) # 缓冲队列, 用于分发计算/绘图
        self.signals = deque(maxlen=self.freq * self.time_window)  # 双端队列存储信号数据

        # 创建包含两个子图的画布
        self.ax1 = self.figure.add_subplot(2, 1, 1)
        self.ax2 = self.figure.add_subplot(2, 1, 2)
        
        # 时域子图
        self.ax1.set_title('Real-time Audio Waveform')
        self.ax1.set_xlabel('Sample Index')
        self.ax1.set_ylabel('Amplitude')
        self.line_time, = self.ax1.plot([], [])
        self.ax1.grid(True, alpha=0.3)
        
        # 频域子图
        self.ax2.set_title('Real-time Frequency Spectrum')
        self.ax2.set_xlabel('Frequency (Hz)')
        self.ax2.set_ylabel('Magnitude')
        self.line_freq, = self.ax2.plot([], [])
        self.ax2.grid(True, alpha=0.3)
        
        self.figure.tight_layout()

        # 定时器用于更新图表
        self.timer = QTimer(self)
        self.timer.setInterval(100)  # 100毫秒更新一次
        self.timer.timeout.connect(self.update_plot)
        
        # 初始化AFSK解码器
        self.decoder = RealTimeAFSKDecoder(
            f_sample=self.freq,
            mark_freq=1800,
            space_freq=1500,
            bitrate=100,
            s_goertzel=9,
            threshold=0.5
        )
        
        # 解码结果回调
        self.decode_callback = None

    def start_plotting(self):
        """开始绘图"""
        self.timer.start()
        
    def stop_plotting(self):
        """停止绘图"""
        self.timer.stop()

    def update_plot(self):
        """更新绘图数据"""
        if len(self.wave_data) >= 2:
            # 进行实时解码
            # 获取最新的音频数据进行解码
            even = len(self.wave_data) // 2 * 2
            print(f"length of wave_data: {len(self.wave_data)}")
            drained = [self.wave_data.popleft() for _ in range(even)]
            signal = np.frombuffer(bytearray(drained), dtype='<i2') / 32768
            decoded_text_new = self.decoder.process_audio(signal) # 处理新增信号, 返回全量解码文本
            if decoded_text_new and self.decode_callback:
                self.decode_callback(decoded_text_new)
            self.signals.extend(signal.tolist())  # 将波形数据添加到绘图数据

        if len(self.signals) > 0:
            # 只显示最近的一段数据，避免图表过于密集
            signal = np.array(self.signals)
            max_samples = min(len(signal), self.freq * self.time_window)
            if len(signal) > max_samples:
                signal = signal[-max_samples:]
            
            # 更新时域图
            x = np.arange(len(signal))
            self.line_time.set_data(x, signal)
            
            # 自动调整时域坐标轴范围
            if len(signal) > 0:
                self.ax1.set_xlim(0, len(signal))
                y_min, y_max = np.min(signal), np.max(signal)
                if y_min != y_max:
                    margin = (y_max - y_min) * 0.1
                    self.ax1.set_ylim(y_min - margin, y_max + margin)
                else:
                    self.ax1.set_ylim(-1, 1)
            
            # 计算频谱（短时离散傅立叶变换）
            if len(signal) > 1:
                # 计算FFT
                fft_signal = np.abs(np.fft.fft(signal))
                frequencies = np.fft.fftfreq(len(signal), 1/self.freq)
                
                # 只取正频率部分
                positive_freq_idx = frequencies >= 0
                freq_positive = frequencies[positive_freq_idx]
                fft_positive = fft_signal[positive_freq_idx]
                
                # 更新频域图
                self.line_freq.set_data(freq_positive, fft_positive)
                
                # 自动调整频域坐标轴范围
                if len(fft_positive) > 0:
                    # 限制频率显示范围到0-4000Hz，避免过于密集
                    max_freq_show = min(4000, self.freq // 2)
                    freq_mask = freq_positive <= max_freq_show
                    if np.any(freq_mask):
                        self.ax2.set_xlim(0, max_freq_show)
                        fft_masked = fft_positive[freq_mask]
                        if len(fft_masked) > 0:
                            fft_max = np.max(fft_masked)
                            if fft_max > 0:
                                self.ax2.set_ylim(0, fft_max * 1.1)
                            else:
                                self.ax2.set_ylim(0, 1)
            
            self.canvas.draw()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Acoustic Check")
        self.setGeometry(100, 100, 1000, 800)

        # 主窗口部件
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        
        # 主布局
        main_layout = QVBoxLayout(main_widget)
        
        # 绘图区域
        self.matplotlib_widget = MatplotlibWidget()
        main_layout.addWidget(self.matplotlib_widget)
        
        # 控制面板
        control_panel = QWidget()
        control_layout = QHBoxLayout(control_panel)
        
        # 监听地址和端口输入
        control_layout.addWidget(QLabel("监听地址:"))
        self.address_input = QLineEdit("0.0.0.0")
        self.address_input.setFixedWidth(120)
        control_layout.addWidget(self.address_input)
        
        control_layout.addWidget(QLabel("端口:"))
        self.port_input = QLineEdit("8000")
        self.port_input.setFixedWidth(80)
        control_layout.addWidget(self.port_input)
        
        # 监听按钮
        self.listen_button = QPushButton("开始监听")
        self.listen_button.clicked.connect(self.toggle_listening)
        control_layout.addWidget(self.listen_button)
        
        # 状态标签
        self.status_label = QLabel("状态: 未连接")
        control_layout.addWidget(self.status_label)
        
        # 数据统计标签
        self.data_label = QLabel("接收数据: 0 bytes")
        control_layout.addWidget(self.data_label)
        
        # 保存按钮
        self.save_button = QPushButton("保存音频")
        self.save_button.clicked.connect(self.save_audio)
        self.save_button.setEnabled(False)
        control_layout.addWidget(self.save_button)
        
        control_layout.addStretch()  # 添加弹性空间
        
        main_layout.addWidget(control_panel)
        
        # 解码显示区域
        decode_panel = QWidget()
        decode_layout = QVBoxLayout(decode_panel)
        
        # 解码标题
        decode_title = QLabel("实时AFSK解码结果:")
        decode_title.setStyleSheet("font-weight: bold; font-size: 14px;")
        decode_layout.addWidget(decode_title)
        
        # 解码文本显示
        self.decode_text = QTextEdit()
        self.decode_text.setMaximumHeight(150)
        self.decode_text.setReadOnly(True)
        self.decode_text.setStyleSheet("font-family: 'Courier New', monospace; font-size: 12px;")
        decode_layout.addWidget(self.decode_text)
        
        # 解码控制按钮
        decode_control_layout = QHBoxLayout()
        
        # 清空按钮
        self.clear_decode_button = QPushButton("清空解码")
        self.clear_decode_button.clicked.connect(self.clear_decode_text)
        decode_control_layout.addWidget(self.clear_decode_button)
        
        # 解码统计标签
        self.decode_stats_label = QLabel("解码统计: 0 bits, 0 chars")
        decode_control_layout.addWidget(self.decode_stats_label)
        
        decode_control_layout.addStretch()
        decode_layout.addLayout(decode_control_layout)
        
        main_layout.addWidget(decode_panel)
        
        # 设置解码回调
        self.matplotlib_widget.decode_callback = self.on_decode_text
        
        # UDP相关属性
        self.udp_transport = None
        self.is_listening = False
        
        # 数据统计定时器
        self.stats_timer = QTimer(self)
        self.stats_timer.setInterval(1000)  # 每秒更新一次统计
        self.stats_timer.timeout.connect(self.update_stats)
        
    def on_decode_text(self, new_text: str):
        """解码文本回调"""
        if new_text:
            # 添加新解码的文本
            current_text = self.decode_text.toPlainText()
            updated_text = current_text + new_text

            # 限制文本长度，保留最新的1000个字符
            if len(updated_text) > 1000:
                updated_text = updated_text[-1000:]
            
            self.decode_text.setPlainText(updated_text)
            
            # 滚动到底部
            cursor = self.decode_text.textCursor()
            cursor.movePosition(cursor.MoveOperation.End)
            self.decode_text.setTextCursor(cursor)
            
    def clear_decode_text(self):
        """清空解码文本"""
        self.decode_text.clear()
        if hasattr(self.matplotlib_widget, 'decoder'):
            self.matplotlib_widget.decoder.clear()
        self.decode_stats_label.setText("解码统计: 0 bits, 0 chars")
        
    def update_decode_stats(self):
        """更新解码统计"""
        if hasattr(self.matplotlib_widget, 'decoder'):
            stats = self.matplotlib_widget.decoder.get_stats()
            stats_text = (
                f"前置: {stats['prelude_bits']} , 已接收{stats['total_chars']} chars, "
                f"缓冲: {stats['buffer_bits']} bits, 状态: {stats['state']}"
            )
            self.decode_stats_label.setText(stats_text)
        
    def toggle_listening(self):
        """切换监听状态"""
        if not self.is_listening:
            self.start_listening()
        else:
            self.stop_listening()
            
    async def start_listening_async(self):
        """异步启动UDP监听"""
        try:
            address = self.address_input.text().strip()
            port = int(self.port_input.text().strip())
            
            loop = asyncio.get_running_loop()
            self.udp_transport, protocol = await loop.create_datagram_endpoint(
                lambda: UDPServerProtocol(self.matplotlib_widget.wave_data),
                local_addr=(address, port)
            )
            
            self.status_label.setText(f"状态: 监听中 ({address}:{port})")
            print(f"UDP服务器启动, 监听 {address}:{port}")
            
        except Exception as e:
            self.status_label.setText(f"状态: 启动失败 - {str(e)}")
            print(f"UDP服务器启动失败: {e}")
            self.is_listening = False
            self.listen_button.setText("开始监听")
            self.address_input.setEnabled(True)
            self.port_input.setEnabled(True)
            
    def start_listening(self):
        """开始监听"""
        try:
            int(self.port_input.text().strip())  # 验证端口号格式
        except ValueError:
            self.status_label.setText("状态: 端口号必须是数字")
            return
            
        self.is_listening = True
        self.listen_button.setText("停止监听")
        self.address_input.setEnabled(False)
        self.port_input.setEnabled(False)
        self.save_button.setEnabled(True)
        
        # 清空数据队列
        self.matplotlib_widget.wave_data.clear()
        
        # 启动绘图和统计更新
        self.matplotlib_widget.start_plotting()
        self.stats_timer.start()
        
        # 异步启动UDP服务器
        loop = asyncio.get_event_loop()
        loop.create_task(self.start_listening_async())

    def stop_listening(self):
        """停止监听"""
        self.is_listening = False
        self.listen_button.setText("开始监听")
        self.address_input.setEnabled(True)
        self.port_input.setEnabled(True)
        
        # 停止UDP服务器
        if self.udp_transport:
            self.udp_transport.close()
            self.udp_transport = None
            
        # 停止绘图和统计更新
        self.matplotlib_widget.stop_plotting()
        self.matplotlib_widget.wave_data.clear()
        self.stats_timer.stop()
        
        self.status_label.setText("状态: 已停止")
        
    def update_stats(self):
        """更新数据统计"""
        data_size = len(self.matplotlib_widget.signals)
        self.data_label.setText(f"接收数据: {data_size} 采样")
        
        # 更新解码统计
        self.update_decode_stats()
        
    def save_audio(self):
        """保存音频数据"""
        if len(self.matplotlib_widget.signals) > 0:
            try:
                signal_data = np.array(self.matplotlib_widget.signals)

                # 保存为WAV文件
                with wave.open("received_audio.wav", "wb") as wf:
                    wf.setnchannels(1)  # 单声道
                    wf.setsampwidth(2)  # 采样宽度为2字节
                    wf.setframerate(self.matplotlib_widget.freq)  # 设置采样率
                    wf.writeframes(signal_data.tobytes())  # 写入数据
                
                self.status_label.setText("状态: 音频已保存为 received_audio.wav")
                print("音频已保存为 received_audio.wav")
                
            except Exception as e:
                self.status_label.setText(f"状态: 保存失败 - {str(e)}")
                print(f"保存音频失败: {e}")
        else:
            self.status_label.setText("状态: 没有足够的数据可保存")


async def main():
    """异步主函数"""
    app = QApplication(sys.argv)
    
    # 设置异步事件循环
    loop = qasync.QEventLoop(app)
    asyncio.set_event_loop(loop)
    
    window = MainWindow()
    window.show()
    
    try:
        with loop:
            await loop.run_forever()
    except KeyboardInterrupt:
        print("程序被用户中断")
    finally:
        # 确保清理资源
        if window.udp_transport:
            window.udp_transport.close()