# Copyright (c) 2022 Horizon Robotics. (authors: Binbin Zhang)
#               2022 Chengdong Liang (liangchengdong@mail.nwpu.edu.cn)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import gradio as gr
import numpy as np
import wenetruntime as wenet
import librosa


wenet.set_log_level(2)
decoder_cn = wenet.Decoder(lang='chs')


def recognition(audio):
    if audio is None:
        return "Input Error! Please enter one audio!"
    sr, y = audio
    # NOTE: model supports 16k sample_rate
    if sr != 16000:
        y = librosa.resample(y, sr, 16000)
    ans = decoder_cn.decode(y.tobytes(), True)
    if ans is None:
        return "ERROR! No text output! Please try again!"
    # NOTE: ans (json)
    # {
    #    'nbest' : [{"sentence" : ""}], 'type' : 'final_result
    # }
    ans = json.loads(ans)
    txt = ans['nbest'][0]['sentence']
    return txt


# input
inputs = gr.inputs.Audio(source="microphone", type="numpy", label='Input audio')

output = gr.outputs.Textbox(label="Output Text")

examples = [
   ['examples/BAC009S0767W0127.wav'],
   ['examples/BAC009S0767W0424.wav'],
   ['examples/BAC009S0767W0488.wav'],
]

text = "Speech Recognition in WeNet | 基于 WeNet 的语音识别"

# description
description = ("Wenet Demo ! This is a Mandarin streaming speech recognition !")

article = (
    "<p style='text-align: center'>"
    "<a href='https://github.com/wenet-e2e/wenet' target='_blank'>Github: Learn more about WeNet</a>"
    "</p>")

interface = gr.Interface(
    fn=recognition,
    inputs=inputs,
    outputs=output,
    title=text,
    description=description,
    article=article,
    examples=examples,
    theme='huggingface',
)

interface.launch(enable_queue=True)