jhauret/spectrogram

1import gradio as gr2import numpy as np3import plotly.graph_objects as go4import scipy.signal as ssig5import librosa6import plotly.io as pio7 8def plot_stft(audio_file):9    # Load audio file10    audio, sampling_rate = librosa.load(audio_file)11 12    # Compute STFT13    freq, frames, stft = ssig.stft(audio,14                                   sampling_rate,15                                   window='hann',16                                   nperseg=512,17                                   noverlap=412,18                                   nfft=1024,19                                   return_onesided=True,20                                   boundary='zeros',21                                   padded=True,22                                   axis=-1)23 24    # Create spectrogram heatmap25    spectrogram = go.Heatmap(z=librosa.amplitude_to_db(np.abs(stft), ref=np.max),26                             x=frames,27                             y=freq,28                             colorscale='Viridis')29 30    # Create Plotly figure31    fig = go.Figure(spectrogram)32 33    # Customize layout34    fig.update_layout(35        font=dict(family='Latin Modern Roman', size=18),36        xaxis=dict(title='Time (seconds)',37                   titlefont=dict(family='Latin Modern Roman', size=18)),38        yaxis=dict(title='Frequency (Hz)',39                   titlefont=dict(family='Latin Modern Roman', size=18)),40        margin=dict(l=0, r=0, t=0, b=0),41    )42 43    fig.update_traces(colorbar_thickness=8, selector=dict(type='heatmap'))44    fig.update_traces(showscale=True, showlegend=False, visible=True)45    fig.update_xaxes(visible=True, showgrid=False)46    fig.update_yaxes(visible=True, showgrid=False)47 48    # Save the figure as an image49    image_path = 'stft_plot.png'50    fig.write_image(image_path)51 52    return image_path53 54# Gradio interface55demo = gr.Interface(fn=plot_stft,56                    inputs=gr.Audio(type="filepath"),57                    outputs="image")58 59demo.launch()60