import gradio as gr
import torch
import librosa
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor

model = AutoModelForAudioClassification.from_pretrained(".")
extractor = AutoFeatureExtractor.from_pretrained(".")

maps = {
    0:"blues",
    1:"classical",
    2:"country",
    3:"disco",
    4:"hiphop",
    5:"jazz",
    6:"metal",
    7:"pop",
    8:"reggae",
    9:"rock"
}

def predict(audio):

    audio, sr = librosa.load(audio, sr=16000)

    inputs = extractor(audio, sampling_rate=16000, return_tensors="pt")

    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits, dim=-1)[0].tolist()

    result = {maps[i]: float(probs[i]) for i in range(10)}

    return result


with gr.Blocks(title="AST_Audio_Classfication") as demo:

    gr.Markdown("""
    <div style="text-align:center;">
        <h1 style="font-size:3rem;">🎵 Music Genre Detection</h1>
        <p style="font-size:1.2rem; color:#555;">
        Audio classification using <b>Audio Spectrogram Transformer (AST)</b><br>
        Built as part of the <b>IIT Madras Intro to Deep Learning & GenAI Project (2026 Term 1)</b>
        </p>
    </div>
    """)

    with gr.Row():

        with gr.Column(scale=1):

            gr.HTML("""
            <div style="
                background:white;
                padding:20px;
                border-radius:14px;
                box-shadow:0 2px 12px rgba(0,0,0,0.08);
                margin-bottom:20px;
            ">
            <h2>📌 Model Overview</h2>
            <ul style="line-height:1.6;">
                <li><b>Architecture:</b> Audio Spectrogram Transformer</li>
                <li><b>Base Model:</b> MIT AST</li>
                <li><b>Task:</b> Music Genre Classification</li>
                <li><b>Classes:</b> 10 Genres</li>
                <li><b>Framework:</b> PyTorch + Hugging Face</li>
            </ul>
            </div>
            """)

            gr.HTML("""
            <div style="
                background:white;
                padding:20px;
                border-radius:14px;
                box-shadow:0 2px 12px rgba(0,0,0,0.08);
                margin-bottom:20px;
            ">
            <h2>🎧 Genres Detected by Model</h2>
            <ul>
                <li>🎷 Jazz</li>
                <li>🎸 Rock</li>
                <li>🎤 Pop</li>
                <li>🎶 Blues</li>
                <li>🎼 Classical</li>
                <li>🤠 Country</li>
                <li>💃 Disco</li>
                <li>🔥 HipHop</li>
                <li>🎛 Metal</li>
                <li>🌴 Reggae</li>
            </ul>
            </div>
            """)

            gr.HTML("""
            <div style="
                background:white;
                padding:20px;
                border-radius:14px;
                box-shadow:0 2px 12px rgba(0,0,0,0.08);
            ">
            <h2>🎓 Project Info</h2>
            <ul style="line-height:1.6;">
                <li><b>Course:</b> Deep Learning & Generative AI</li>
                <li><b>Institution:</b> IIT Madras</li>
                <li><b>Term:</b> 2026 Term 1</li>
                <li><b>Student:</b> Ayusman Samasi</li>
                <li><b>Roll:</b> 22f3001XXX</li>
                <li><b>Email:</b> 22f3001XXX@ds.study.iitm.ac.in</li>
            </ul>
            </div>
            """)

        with gr.Column(scale=2):

            audio_input = gr.Audio(
                sources=["upload","microphone"],
                type="filepath",
                label="Upload or Record Audio"
            )

            btn = gr.Button("🎯 Detect Genre")

            output = gr.Label(
                num_top_classes=5,
                label="Prediction Probabilities"
            )

            btn.click(predict, inputs=audio_input, outputs=output)

    gr.Markdown("""
    <br>
    <p style="text-align:center; color:#777;">
    Built with ❤️ using Hugging Face Transformers & Gradio<br>
    <b>Ayusman Samasi • IIT Madras DL & GenAI Project T1 2026</b>
    </p>
    """)

demo.launch()