import gradio as gr
import numpy as np
import plotly.graph_objects as go
import librosa

def plot_stft(audio_file):
    # Load the audio file
    y, sr = librosa.load(audio_file)
    
    # Compute the STFT
    D = librosa.stft(y)
    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
    
    # Generate time and frequency axes
    times = librosa.times_like(S_db)
    freqs = librosa.fft_frequencies(sr=sr)
    
    # Create Plotly figure
    fig = go.Figure(data=go.Heatmap(
        z=S_db,
        x=times,
        y=freqs,
        colorscale='Viridis'
    ))
    
    # Update layout for better visualization
    fig.update_layout(
        title="STFT (Short-Time Fourier Transform)",
        xaxis_title="Time (s)",
        yaxis_title="Frequency (Hz)",
        yaxis_type="log"
    )
    
    # Return the HTML representation of the plot
    return fig.to_html()

# Gradio interface
demo = gr.Interface(fn=plot_stft, 
                    inputs=gr.Audio(type="filepath"), 
                    outputs="html")

demo.launch()