app.py

from plotly.subplots import make_subplots
import plotly.graph_objects as go
import gradio as gr
import numpy as np
import itertools
import librosa
import os

example_dir = "Examples"
example_files = [os.path.join(example_dir, f) for f in os.listdir(example_dir) if f.endswith(('.wav', '.mp3', '.ogg'))]
example_pairs = [list(pair) for pair in itertools.combinations(example_files, 2)][:25]  # Limit to 5 pairs

# GENERAL HELPER FUNCTIONS
def getaudiodata(audio:gr.Audio)->tuple[int,np.ndarray]:
    # Extract audio data and sample rate
    sr, audiodata = audio

    # Ensure audiodata is a numpy array
    if not isinstance(audiodata, np.ndarray):
        audiodata = np.array(audiodata)

    # Check if audio is mono or stereo
    if len(audiodata.shape) > 1:
        # If stereo, convert to mono by averaging channels
        audiodata = np.mean(audiodata, axis=1)

    audiodata = np.astype(audiodata, np.float16)

    return sr, audiodata


# HELPER FUNCTIONS FOR SINGLE AUDIO ANALYSIS
def getBeats(audiodata:np.ndarray, sr:int):
    # Compute onset envelope
    onset_env = librosa.onset.onset_strength(y=audiodata, sr=sr)
    
    # Detect beats
    tempo, beats = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
    
    # Convert beat frames to time
    beattimes = librosa.frames_to_time(beats, sr=sr)
    
    return tempo[0], beattimes

def plotCombined(audiodata, sr):
    # Create subplots
    fig = make_subplots(rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.1,
                        subplot_titles=('Audio Waveform', 'Spectrogram'))

    # Waveform plot
    time = (np.arange(0, len(audiodata)) / sr)*2
    fig.add_trace(
        go.Scatter(x=time, y=audiodata, mode='lines', name='Waveform', line=dict(color='blue', width=1)),
        row=1, col=1
    )

    # Spectrogram plot
    D = librosa.stft(audiodata)
    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
    times = librosa.times_like(S_db)
    freqs = librosa.fft_frequencies(sr=sr)

    fig.add_trace(
        go.Heatmap(z=S_db, x=times, y=freqs, colorscale='Viridis', 
                   zmin=S_db.min(), zmax=S_db.max(), colorbar=dict(title='Magnitude (dB)')),
        row=2, col=1
    )


    # Update layout
    fig.update_layout(
        height=800, width=900,
        title_text="Audio Analysis",
    )

    fig.update_xaxes(title_text="Time (s)", row=2, col=1)
    fig.update_yaxes(title_text="Amplitude", row=1, col=1)
    fig.update_yaxes(title_text="Frequency (Hz)", type="log", row=2, col=1)

    return fig

def plotbeatshist(tempo, beattimes):
    # Calculate beat durations
    beat_durations = np.diff(beattimes)
    
    # Create histogram
    fig = go.Figure()
    fig.add_trace(go.Histogram(
        x=beat_durations,
        nbinsx=60,  # You can adjust the number of bins as needed
        name='Beat Durations'
    ))
    
    # Add vertical line for average beat duration
    avg_duration = 60 / tempo  # Convert tempo (BPM) to seconds
    fig.add_vline(x=avg_duration, line_dash="dash", line_color="red",
                  annotation_text=f"Average: {avg_duration:.2f}s", 
                  annotation_position="top right")
    
    # Update layout
    fig.update_layout(
        title_text='Histogram of Beat Durations',
        xaxis_title_text='Beat Duration (seconds)',
        yaxis_title_text='Count',
        bargap=0.05,  # gap between bars
    )
    
    return fig

def analyze_single(audio:gr.Audio):
    # Extract audio data and sample rate
    sr, audiodata = getaudiodata(audio)


    # Now you have:
    # - audiodata: a 1D numpy array containing the audio samples
    # - sr: the sample rate of the audio

    # Your analysis code goes here
    # For example, you could print basic information:
    print(f"Audio length: {len(audiodata) / sr:.2f} seconds")
    print(f"Sample rate: {sr} Hz")

    zcr = librosa.feature.zero_crossing_rate(audiodata)[0]
    print(f"Mean Zero Crossing Rate: {np.mean(zcr):.4f}")

    # Calculate RMS Energy
    rms = librosa.feature.rms(y=audiodata)[0]
    print(f"Mean RMS Energy: {np.mean(rms):.4f}")

    tempo, beattimes = getBeats(audiodata, sr)
    spectogram_wave = plotCombined(audiodata, sr)
    beats_histogram = plotbeatshist(tempo, beattimes)

    # Return your analysis results
    results = f"""
    - Audio length: {len(audiodata) / sr:.2f} seconds
    - Sample rate: {sr} Hz
    - Mean Zero Crossing Rate: {np.mean(zcr):.4f}
    - Mean RMS Energy: {np.mean(rms):.4f}
    - Tempo: {tempo:.4f}
    - Beats: {beattimes}
    - Beat durations: {np.diff(beattimes)}
    - Mean Beat Duration: {np.mean(np.diff(beattimes)):.4f}
    """
    return results, spectogram_wave, beats_histogram
#-----------------------------------------------
#-----------------------------------------------

# HELPER FUNCTIONS FOR DUAL AUDIO ANALYSIS
def analyze_double(audio1:gr.Audio, audio2:gr.Audio):
    
    sr1, audiodata1 = getaudiodata(audio1)
    sr2, audiodata2 = getaudiodata(audio2)


    combinedfig = plotCombineddouble(audiodata1, sr1, audiodata2, sr2)
    return combinedfig

def plotCombineddouble(audiodata1, sr1, audiodata2, sr2):
    # Create subplots
    fig = make_subplots(rows=2, cols=2, shared_xaxes=True, vertical_spacing=0.1,
                        subplot_titles=['Audio 1 Waveform','Audio 2 Audio Waveform', 'Audio 1 Spectrogram', 'Audio 2 Spectrogram'])

    # Waveform plot
    time = (np.arange(0, len(audiodata1)) / sr1)*2
    fig.add_trace(
        go.Scatter(x=time, y=audiodata1, mode='lines', line=dict(color='blue', width=1), showlegend=False),
        row=1, col=1
    )

    # Spectrogram plot
    D = librosa.stft(audiodata1)
    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
    times = librosa.times_like(S_db)
    freqs = librosa.fft_frequencies(sr=sr1)

    fig.add_trace(
        go.Heatmap(z=S_db, x=times, y=freqs, colorscale='Viridis', 
                   zmin=S_db.min(), zmax=S_db.max(), showlegend=False),#, colorbar=dict(title='Magnitude (dB)')),
        row=2, col=1
    )

    # Waveform plot
    time = (np.arange(0, len(audiodata2)) / sr2)*2
    fig.add_trace(
        go.Scatter(x=time, y=audiodata2, mode='lines', line=dict(color='blue', width=1), showlegend=False),
        row=1, col=2
    )

    # Spectrogram plot
    D = librosa.stft(audiodata2)
    S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
    times = librosa.times_like(S_db)
    freqs = librosa.fft_frequencies(sr=sr2)

    fig.add_trace(
        go.Heatmap(z=S_db, x=times, y=freqs, colorscale='Viridis', 
                   zmin=S_db.min(), zmax=S_db.max(), showlegend=False),#, colorbar=dict(title='Magnitude (dB)')),
        row=2, col=2
    )





    # Update layout
    fig.update_layout(
        height=800, width=1200,
        title_text="Audio Analysis",
        showlegend=False
    )

    fig.update_xaxes(title_text="Time (s)", row=2, col=1)
    fig.update_yaxes(title_text="Amplitude", row=1, col=1)
    fig.update_yaxes(title_text="Frequency (Hz)", type="log", row=2, col=1)

    fig.update_xaxes(title_text="Time (s)", row=2, col=2)
    fig.update_yaxes(title_text="Amplitude", row=1, col=2)
    fig.update_yaxes(title_text="Frequency (Hz)", type="log", row=2, col=2)

    return fig




with gr.Blocks() as app:

    gr.Markdown("# Heartbeat")
    gr.Markdown("This App helps to analyze and extract Information from Heartbeats")
    gr.Markdown("""
    - Beat (mean) (average heartbeat duration)
    - S1, S2 (mean) (average S1,S2 duration)
    - mean - herzschlag (synthesised) - Bild (Wave & Spectogram)
    - FFT & Mel Spectogram
    - Plot of Wave & Spectogram (Beats annotated)
    """)

    with gr.Tab("Single Audio"):

        audiofile = gr.Audio(
            label="Audio of a Heartbeat",
            sources="upload")
        
        analyzebtn = gr.Button("analyze")

        results = gr.Markdown()
        spectogram_wave = gr.Plot()
        beats_histogram = gr.Plot()

        analyzebtn.click(analyze_single, audiofile, [results, spectogram_wave, beats_histogram])

        gr.Examples(
            examples=example_files,
            inputs=audiofile,
            outputs=[results, spectogram_wave],
            fn=analyze_single,
            cache_examples=False
        )

        gr.Markdown("""### Open TODO's
        - Create Histogram for Beat durations
        - classify Beat's into S1 and S2
        - synthesise the mean Beat S1 & S2""")

    with gr.Tab("Two Audios"):

        with gr.Row():

            audioone = gr.Audio(
                label="Audio of a Heartbeat",
                sources="upload")
            audiotwo = gr.Audio(
                label="Audio of a Heartbeat",
                sources="upload")
            
        analyzebtn2 = gr.Button("analyze & compare")

        with gr.Accordion("Results",open=False):
            results2 = gr.Markdown()
            spectogram_wave2 = gr.Plot()

        analyzebtn2.click(analyze_double, inputs=[audioone,audiotwo], outputs=spectogram_wave2)

        # Add gr.Examples for the Two Audios tab
        gr.Examples(
            examples=example_pairs,  # Create pairs of the same file for demonstration
            inputs=[audioone, audiotwo],
            outputs=spectogram_wave2,
            fn=analyze_double,
            cache_examples=False
        )

if __name__ == "__main__":
    app.launch()