from transformers import pipeline
import gradio as gr

model_id = "artyomboyko/distilhubert-finetuned-gtzan"
pipe = pipeline("audio-classification", model=model_id)

def classify_audio(filepath):
    preds = pipe(filepath)
    outputs = {}
    for p in preds:
        outputs[p["label"]] = p["score"]
    return outputs


demo = gr.Blocks()

title = "Audio classification"
description = """
This demo is designed to test the music classification. It is important to remember that music classification depends very much on the quality of the recording, for example, when classifying a recording from a 
microphone with poor high frequencies, the song "Evanescence - bring me to life" may not be correctly defined as classical music. But if we transfer a recording of the same song as a file, it is correctly 
identified as metal music. In addition, a real problem is caused by compositions in which there is a mixture of different styles of music, or modern styles of examples of which were not in the training dataset.
"""

mic_classify_audio = gr.Interface(
    fn=classify_audio,
    inputs=gr.Audio(source="microphone", type="filepath"),
    outputs=gr.outputs.Label(),
    title=title,
    description=description,
)

file_classify_audio = gr.Interface(
    fn=classify_audio,
    inputs=gr.Audio(source="upload", type="filepath"),
    outputs=gr.outputs.Label(),
    #examples=[["./example.wav"]],
    title=title,
    description=description,
)

with demo:
    gr.TabbedInterface([mic_classify_audio, file_classify_audio], ["Microphone", "Audio File"])

demo.launch(debug=True)