import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import librosa
import librosa.display
import base64
from io import BytesIO

def plot_stft(audio_file):
    # Load audio file
    y, sr = librosa.load(audio_file)
    
    # Compute the Short-Time Fourier Transform (STFT)
    D = librosa.stft(y)
    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
    
    # Plot the STFT
    plt.figure(figsize=(10, 6))
    librosa.display.specshow(S_db, sr=sr, x_axis='time', y_axis='log')
    plt.colorbar(format='%+2.0f dB')
    plt.title('STFT (Short-Time Fourier Transform)')
    
    # Save the plot to a BytesIO object
    buf = BytesIO()
    plt.savefig(buf, format='png')
    plt.close()
    buf.seek(0)
    
    # Encode the image as base64
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    buf.close()
    
    # Create an HTML img tag with the base64 encoded image
    html_img = f'<img src="data:image/png;base64,{image_base64}" alt="STFT plot"/>'
    
    return html_img

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

demo.launch()