import gradio as gr
import os
import soundfile as sf
from pydub import AudioSegment
import numpy as np
from transformers import Wav2Vec2Processor
import torch

def record_audio(audio):
    # 检查传入的 audio 变量
    if audio is None:
        return "没有录音文件，请录音后再提交。"
    
    print(f"传入的音频文件路径: {audio}")  # 打印调试信息
    
    audio_file = audio  # 这里 audio 应该是文件路径
    
    # 读取音频文件
    try:
        audio_data, sample_rate = sf.read(audio_file)
    except Exception as e:
        return f"读取音频文件失败: {str(e)}"
    
    # 转换采样率为 16kHz
    if sample_rate != 16000:
        audio_segment = AudioSegment(
            audio_data.tobytes(), 
            frame_rate=sample_rate,
            sample_width=audio_data.dtype.itemsize, 
            channels=1
        )
        audio_segment = audio_segment.set_frame_rate(16000)
        audio_data = np.array(audio_segment.get_array_of_samples())
        sample_rate = 16000
    
    # 保存音频文件为 voice.wav
    filename = "voice.wav"
    sf.write(filename, audio_data, sample_rate)
    
    return filename  # 返回保存的文件路径

def voice_sentiment(voice):
    MODEL_ID = "jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn"
    processor = Wav2Vec2Processor.from_pretrained(MODEL_ID)

    model_path = 'sentiment_rnn_model.pth'  # 已训练好的模型路径
    input_size = 200000  # 根据最大特征长度设定
    hidden_size = 128
    num_layers = 2

    # 加载模型
    model = load_model(model_path, input_size, hidden_size, num_layers)

    # 读取音频文件进行情感预测
    audio_data, sample_rate = sf.read(voice)

    # 进行情感预测
    max_length = input_size
    predicted_value, predicted_class = predict(voice, processor, model, max_length)

    # 输出结果
    if predicted_value is not None:
        return f'Predicted value: {predicted_value}, Predicted class: {"Positive" if predicted_class == 1 else "Negative"}'
    else:
        return "预测失败，未能处理音频文件。"

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("### 录音与情感分析")
    
    with gr.Row("音频情感分析"):
        with gr.Column():
            record_button = gr.Audio(label="录音", type="filepath")
            submit_button = gr.Button("提交录音")
        voice_output = gr.Textbox(label="情感数据")

    # 设置提交按钮的回调
    submit_button.click(record_audio, inputs=record_button, outputs=voice_output)
    voice_output.change(voice_sentiment, inputs=voice_output, outputs=voice_output)

# 启动 Gradio 应用
demo.launch()