from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import gradio as gr

import os
import shutil


def greet(audio):
    print(audio)
    if not os.path.exists("/tmp/test_wavs"):
        os.makedirs("/tmp/test_wavs")
    shutil.copy(audio, "/tmp/test_wavs/01_000001.wav")
    print("-----------------------------------------------")
    print(os.listdir("/tmp/test_wavs/"))
    print("-----------------------------------------------")
    
    inference_pipeline = pipeline(task=Tasks.auto_speech_recognition, model='damo/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch')
    rec_result = inference_pipeline(audio_in=audio)
    semantic_cls = pipeline(Tasks.text_classification, 'damo/nlp_structbert_sentiment-classification_chinese-base')
    cls = semantic_cls(input=rec_result['text'])
    cls_dict = {cls['labels'][0]: cls['scores'][0], cls['labels'][1]: cls['scores'][1]}
    print(cls_dict)
    return rec_result['text'], cls_dict


with gr.Blocks() as demo:
    gr.Markdown("# 语言情感判断")

    with gr.Tab("测试音频"):
        test = gr.Audio(type='filepath', value='test.wav')
        button = gr.Button("情感分析", variant="primary")
        output_txt = gr.Textbox(label='文本')
        output_label = gr.Label(label='情感')
        button.click(greet,
                     inputs=test,
                     outputs=[output_txt, output_label])

    with gr.Tab("麦克风"):
        in_audio = gr.Audio(source='microphone', type='filepath')
        button = gr.Button("情感分析", variant="primary")
        output_txt = gr.Textbox(label='文本')
        output_label = gr.Label(label='情感')
        button.click(greet,
                     inputs=in_audio,
                     outputs=[output_txt, output_label])



    with gr.Tab("上传音频"):
        in_audio = gr.Audio(type='filepath')
        button = gr.Button("情感分析", variant="primary")
        output_txt = gr.Textbox(label='文本')
        output_label = gr.Label(label='情感')
        button.click(greet,
                     inputs=in_audio,
                     outputs=[output_txt, output_label])

gr.close_all()
demo.queue(concurrency_count=5)
demo.launch()