infer_gui.py

import _thread
import argparse
import functools
import os
import platform
import time
import tkinter.messagebox
from tkinter import *
from tkinter.filedialog import askopenfilename

import numpy as np
import soundcard
import soundfile
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, AutoModelForCausalLM
from zhconv import convert

from utils.utils import print_arguments, add_arguments

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

parser = argparse.ArgumentParser(description=__doc__)
add_arg = functools.partial(add_arguments, argparser=parser)
add_arg("model_path",  type=str,  default="models/whisper-tiny-finetune/", help="合并模型的路径，或者是huggingface上模型的名称")
add_arg("language",    type=str,  default="chinese", help="设置语言，如果为None则预测的是多语言")
add_arg("use_gpu",     type=bool, default=True,      help="是否使用gpu进行预测")
add_arg("num_beams",   type=int,  default=1,         help="解码搜索大小")
add_arg("batch_size",  type=int,  default=16,        help="预测batch_size大小")
add_arg("use_compile", type=bool, default=False,     help="是否使用Pytorch2.0的编译器")
add_arg("task",        type=str,  default="transcribe", choices=['transcribe', 'translate'], help="模型的任务")
add_arg("assistant_model_path",  type=str, default=None, help="助手模型，可以提高推理速度，例如openai/whisper-tiny")
add_arg("local_files_only",      type=bool, default=True, help="是否只在本地加载模型，不尝试下载")
add_arg("use_flash_attention_2", type=bool, default=False, help="是否使用FlashAttention2加速")
add_arg("use_bettertransformer", type=bool, default=False, help="是否使用BetterTransformer加速")
args = parser.parse_args()
print_arguments(args)


class SpeechRecognitionApp:
    def __init__(self, window: Tk, args):
        self.window = window
        self.wav_path = None
        self.predicting = False
        self.playing = False
        self.recording = False
        # 录音参数
        self.frames = []
        self.sample_rate = 16000
        self.interval_time = 0.5
        self.block_size = int(self.sample_rate * self.interval_time)
        # 最大录音时长
        self.max_record = 600
        # 录音保存的路径
        self.output_path = 'dataset/record'
        # 指定窗口标题
        self.window.title("夜雨飘零语音识别")
        # 固定窗口大小
        self.window.geometry('870x500')
        self.window.resizable(False, False)
        # 识别短语音按钮
        self.short_button = Button(self.window, text="选择文件", width=20, command=self.predict_audio_thread)
        self.short_button.place(x=10, y=10)
        # 录音按钮
        self.record_button = Button(self.window, text="录音识别", width=20, command=self.record_audio_thread)
        self.record_button.place(x=170, y=10)
        # 播放音频按钮
        self.play_button = Button(self.window, text="播放音频", width=20, command=self.play_audio_thread)
        self.play_button.place(x=330, y=10)
        # 输出结果文本框
        self.result_label = Label(self.window, text="输出日志：")
        self.result_label.place(x=10, y=70)
        self.result_text = Text(self.window, width=120, height=30)
        self.result_text.place(x=10, y=100)
        # 转阿拉伯数字控件
        self.check_frame = Frame(self.window)
        self.joint_text_check_var = BooleanVar()
        self.joint_text_check = Checkbutton(self.check_frame, text='拼接文本', variable=self.joint_text_check_var)
        self.joint_text_check.grid(column=0, row=0)
        self.to_simple_check_var = BooleanVar()
        self.to_simple_check = Checkbutton(self.check_frame, text='繁体转简体', variable=self.to_simple_check_var)
        self.to_simple_check.grid(column=1, row=0)
        self.to_simple_check.select()
        self.task_check_var = BooleanVar()
        self.task_check = Checkbutton(self.check_frame, text='音频转录', variable=self.task_check_var)
        self.task_check.grid(column=2, row=0)
        self.task_check.select()
        self.check_frame.grid(row=1)
        self.check_frame.place(x=600, y=10)

        # 设置设备
        device = "cuda:0" if torch.cuda.is_available() and args.use_gpu else "cpu"
        torch_dtype = torch.float16 if torch.cuda.is_available() and args.use_gpu else torch.float32

        # 获取Whisper的特征提取器、编码器和解码器
        processor = AutoProcessor.from_pretrained(args.model_path)

        # 获取模型
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            args.model_path, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True,
            use_flash_attention_2=args.use_flash_attention_2
        )
        if args.use_bettertransformer and not args.use_flash_attention_2:
            model = model.to_bettertransformer()
        # 使用Pytorch2.0的编译器
        if args.use_compile:
            if torch.__version__ >= "2" and platform.system().lower() != 'windows':
                model = torch.compile(model)
        model.to(device)

        # 获取助手模型
        generate_kwargs_pipeline = None
        if args.assistant_model_path is not None:
            assistant_model = AutoModelForCausalLM.from_pretrained(
                args.assistant_model_path, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
            )
            assistant_model.to(device)
            generate_kwargs_pipeline = {"assistant_model": assistant_model}

        # 获取管道
        self.infer_pipe = pipeline("automatic-speech-recognition",
                                   model=model,
                                   tokenizer=processor.tokenizer,
                                   feature_extractor=processor.feature_extractor,
                                   max_new_tokens=128,
                                   chunk_length_s=30,
                                   batch_size=2,
                                   torch_dtype=torch_dtype,
                                   generate_kwargs=generate_kwargs_pipeline,
                                   device=device)

        # 预热
        _ = self.infer_pipe("dataset/test.wav")

    # 预测短语音线程
    def predict_audio_thread(self):
        if not self.predicting:
            self.wav_path = askopenfilename(filetypes=[("音频文件", "*.wav"), ("音频文件", "*.mp3")],
                                            initialdir='./dataset')
            if self.wav_path == '': return
            self.result_text.delete('1.0', 'end')
            self.result_text.insert(END, "已选择音频文件：%s\n" % self.wav_path)
            self.result_text.insert(END, "正在识别中...\n")
            _thread.start_new_thread(self.predict_audio, (self.wav_path,))
        else:
            tkinter.messagebox.showwarning('警告', '正在预测，请等待上一轮预测结束！')

    # 预测短语音
    def predict_audio(self, wav_path):
        self.predicting = True
        self.result_text.delete('1.0', 'end')
        try:
            task = "transcribe" if self.task_check_var.get() else "translate"
            # 推理参数
            generate_kwargs = {"task": task, "num_beams": args.num_beams}
            if args.language is not None:
                generate_kwargs["language"] = args.language
            # 推理
            result = self.infer_pipe(wav_path, return_timestamps=True, generate_kwargs=generate_kwargs)
            # 判断是否要分段输出
            if self.joint_text_check_var.get():
                text = result['text']
                # 繁体转简体
                if self.to_simple_check_var.get():
                    text = convert(text, 'zh-cn')
                self.result_text.delete('1.0', 'end')
                self.result_text.insert(END, f"{text}\n")
            else:
                for chunk in result["chunks"]:
                    text = chunk['text']
                    # 繁体转简体
                    if self.to_simple_check_var.get():
                        text = convert(text, 'zh-cn')
                    self.result_text.insert(END, f"[{chunk['timestamp'][0]} - {chunk['timestamp'][1]}]：{text}\n")
            self.predicting = False
        except Exception as e:
            print(e)
            self.predicting = False

    # 录音识别线程
    def record_audio_thread(self):
        if not self.playing and not self.recording:
            self.result_text.delete('1.0', 'end')
            self.recording = True
            _thread.start_new_thread(self.record_audio, ())
        else:
            if self.playing:
                tkinter.messagebox.showwarning('警告', '正在播放音频，无法录音！')
            else:
                # 停止播放
                self.recording = False

    # 播放音频线程
    def play_audio_thread(self):
        if self.wav_path is None or self.wav_path == '':
            tkinter.messagebox.showwarning('警告', '音频路径为空！')
        else:
            if not self.playing and not self.recording:
                _thread.start_new_thread(self.play_audio, ())
            else:
                if self.recording:
                    tkinter.messagebox.showwarning('警告', '正在录音，无法播放音频！')
                else:
                    # 停止播放
                    self.playing = False

    def record_audio(self):
        self.frames = []
        self.record_button.configure(text='停止录音')
        self.result_text.insert(END, "正在录音...\n")
        # 打开默认的输入设备
        input_device = soundcard.default_microphone()
        recorder = input_device.recorder(samplerate=self.sample_rate, channels=1, blocksize=self.block_size)
        with recorder:
            while True:
                if len(self.frames) * self.interval_time > self.max_record: break
                # 开始录制并获取数据
                data = recorder.record(numframes=self.block_size)
                data = data.squeeze()
                self.frames.append(data)
                self.result_text.delete('1.0', 'end')
                self.result_text.insert(END, f"已经录音{len(self.frames) * self.interval_time}秒\n")
                if not self.recording: break
        # 拼接录音数据
        data = np.concatenate(self.frames)
        # 保存音频数据
        os.makedirs(self.output_path, exist_ok=True)
        self.wav_path = os.path.join(self.output_path, '%s.wav' % str(int(time.time())))
        soundfile.write(self.wav_path, data=data, samplerate=self.sample_rate)
        self.recording = False
        self.record_button.configure(text='录音识别')
        self.result_text.delete('1.0', 'end')
        _thread.start_new_thread(self.predict_audio, (self.wav_path,))

    # 播放音频
    def play_audio(self):
        self.play_button.configure(text='停止播放')
        self.playing = True
        default_speaker = soundcard.default_speaker()
        data, sr = soundfile.read(self.wav_path)
        with default_speaker.player(samplerate=sr) as player:
            for i in range(0, data.shape[0], sr):
                if not self.playing: break
                d = data[i:i + sr]
                player.play(d / np.max(np.abs(d)))
        self.playing = False
        self.play_button.configure(text='播放音频')


tk = Tk()
myapp = SpeechRecognitionApp(tk, args)

if __name__ == '__main__':
    tk.mainloop()