import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import librosa
import librosa.display

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 as an image
    plt.savefig('stft_plot.png')
    plt.close()
    
    return 'stft_plot.png'

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

demo.launch()