spectogram has correct length, filename displayed, average beat

app.py

@@ -15,9 +15,9 @@ random.shuffle(all_pairs)
 example_pairs = [list(pair) for pair in all_pairs[:25]]
 
 # GENERAL HELPER FUNCTIONS
-def getaudiodata(audio:gr.Audio)->tuple[int,np.ndarray]:
-    # Extract audio data and sample rate
-    sr, audiodata = audio
+def getaudiodata(filepath)->tuple[int,np.ndarray]:
+
+    audiodata, sr = librosa.load(filepath, sr=None)
 
     # Ensure audiodata is a numpy array
     if not isinstance(audiodata, np.ndarray):
@@ -70,8 +70,39 @@ def getHRV(beattimes: np.ndarray) -> np.ndarray:
     
     return instantaneous_hr
 
+def create_average_heartbeat(audiodata, sr):
+    # 1. Detect individual heartbeats
+    onset_env = librosa.onset.onset_strength(y=audiodata, sr=sr)
+    peaks, _ = find_peaks(onset_env, distance=sr//2)  # Assume at least 0.5s between beats
+
+    # 2. Extract individual heartbeats
+    beat_length = sr  # Assume 1 second for each beat
+    beats = []
+    for peak in peaks:
+        if peak + beat_length < len(audiodata):
+            beat = audiodata[peak:peak+beat_length]
+            beats.append(beat)
+
+    # 3. Align and average the beats
+    if beats:
+        avg_beat = np.mean(beats, axis=0)
+    else:
+        avg_beat = np.array([])
+
+    # 4. Create a Plotly figure of the average heartbeat
+    time = np.arange(len(avg_beat)) / sr
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=time, y=avg_beat, mode='lines', name='Average Beat'))
+    fig.update_layout(
+        title='Average Heartbeat',
+        xaxis_title='Time (s)',
+        yaxis_title='Amplitude'
+    )
+
+    return fig, avg_beat
+
 # HELPER FUNCTIONS FOR SINGLE AUDIO ANALYSIS
-def plotCombined(audiodata, sr):
+def plotCombined(audiodata, sr, filename):
     # Get beat times
     tempo, beattimes = getBeats(audiodata, sr)
     
@@ -80,7 +111,7 @@ def plotCombined(audiodata, sr):
                         subplot_titles=('Audio Waveform', 'Spectrogram', 'Heart Rate Variability'))
     
     # Time array for the full audio
-    time = (np.arange(0, len(audiodata)) / sr)*2
+    time = (np.arange(0, len(audiodata)) / sr) * 2
     
     # Waveform plot
     fig.add_trace(
@@ -105,23 +136,30 @@ def plotCombined(audiodata, sr):
     )
     
     # Spectrogram plot
-    D = librosa.stft(audiodata)
+    n_fft = 2048  # You can adjust this value
+    hop_length = n_fft // 4  # You can adjust this value
+
+    D = librosa.stft(audiodata, n_fft=n_fft, hop_length=hop_length)
     S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
-    times = librosa.times_like(S_db)
-    freqs = librosa.fft_frequencies(sr=sr)
+
+    # Calculate the correct time array for the spectrogram
+    spec_times = librosa.times_like(S_db, sr=sr, hop_length=hop_length)
+
+    freqs = librosa.fft_frequencies(sr=sr, n_fft=n_fft)
     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)')),
+        go.Heatmap(z=S_db, x=spec_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=1000,
-        title_text="Audio Analysis with Heart Rate Variability",
+        title_text=filename,
         showlegend=False
     )
     fig.update_xaxes(title_text="Time (s)", row=2, col=1)
+    fig.update_xaxes(range=[0, len(audiodata)/sr], row=2, col=1)
     fig.update_yaxes(title_text="Amplitude", row=1, col=1)
     fig.update_yaxes(title_text="HRV", row=3, col=1)
     fig.update_yaxes(title_text="Frequency (Hz)", type="log", row=2, col=1)
@@ -167,7 +205,11 @@ def plotbeatscatter(tempo, beattimes):
 
 def analyze_single(audio:gr.Audio):
     # Extract audio data and sample rate
-    sr, audiodata = getaudiodata(audio)
+    filepath = audio
+    filename = filepath.split("/")[-1]
+
+
+    sr, audiodata = getaudiodata(filepath)
 
 
     # Now you have:
@@ -187,11 +229,21 @@ def analyze_single(audio:gr.Audio):
     # print(f"Mean RMS Energy: {np.mean(rms):.4f}")
 
     tempo, beattimes = getBeats(audiodata, sr)
-    spectogram_wave = plotCombined(audiodata, sr)
+    spectogram_wave = plotCombined(audiodata, sr, filename)
     beats_histogram = plotbeatscatter(tempo[0], beattimes)
 
+    # Add the new average heartbeat analysis
+    avg_beat_plot, avg_beat = create_average_heartbeat(audiodata, sr)
+
+    # Calculate some statistics about the average beat
+    avg_beat_duration = len(avg_beat) / sr
+    avg_beat_energy = np.sum(np.square(avg_beat))
+
     # Return your analysis results
     results = f"""
+    Average Heartbeat Analysis:
+    - Duration: {avg_beat_duration:.3f} seconds
+    - Energy: {avg_beat_energy:.3f}
     - Audio length: {len(audiodata) / sr:.2f} seconds
     - Sample rate: {sr} Hz
     - Mean Zero Crossing Rate: {np.mean(zcr):.4f}
@@ -201,7 +253,7 @@ def analyze_single(audio:gr.Audio):
     - Beat durations: {np.diff(beattimes)}
     - Mean Beat Duration: {np.mean(np.diff(beattimes)):.4f}
     """
-    return results, spectogram_wave, beats_histogram
+    return results, spectogram_wave, avg_beat_plot, beats_histogram
 #-----------------------------------------------
 #-----------------------------------------------
 
@@ -297,6 +349,7 @@ with gr.Blocks() as app:
     with gr.Tab("Single Audio"):
 
         audiofile = gr.Audio(
+            type="filepath",
             label="Audio of a Heartbeat",
             sources="upload")
         
@@ -304,9 +357,10 @@ with gr.Blocks() as app:
 
         results = gr.Markdown()
         spectogram_wave = gr.Plot()
+        avg_beat_plot = gr.Plot()
         beats_histogram = gr.Plot()
 
-        analyzebtn.click(analyze_single, audiofile, [results, spectogram_wave, beats_histogram])
+        analyzebtn.click(analyze_single, audiofile, [results, spectogram_wave, avg_beat_plot, beats_histogram])
 
         gr.Examples(
             examples=example_files,
@@ -321,32 +375,34 @@ with gr.Blocks() as app:
         - classify Beat's into S1 and S2
         - synthesise the mean Beat S1 & S2""")
 
-    with gr.Tab("Two Audios"):
+    # with gr.Tab("Two Audios"):
 
-        with gr.Row():
+    #     with gr.Row():
 
-            audioone = gr.Audio(
-                label="Audio of a Heartbeat",
-                sources="upload")
-            audiotwo = gr.Audio(
-                label="Audio of a Heartbeat",
-                sources="upload")
+    #         audioone = gr.Audio(
+    #             type="filepath",
+    #             label="Audio of a Heartbeat",
+    #             sources="upload")
+    #         audiotwo = gr.Audio(
+    #             type="filepath",
+    #             label="Audio of a Heartbeat",
+    #             sources="upload")
             
-        analyzebtn2 = gr.Button("analyze & compare")
+    #     analyzebtn2 = gr.Button("analyze & compare")
 
-        with gr.Accordion("Results",open=False):
-            results2 = gr.Markdown()
-            spectogram_wave2 = gr.Plot()
+    #     with gr.Accordion("Results",open=False):
+    #         results2 = gr.Markdown()
+    #         spectogram_wave2 = gr.Plot()
 
-        analyzebtn2.click(analyze_double, inputs=[audioone,audiotwo], outputs=spectogram_wave2)
+    #     analyzebtn2.click(analyze_double, inputs=[audioone,audiotwo], outputs=spectogram_wave2)
 
-        gr.Examples(
-            examples=example_pairs,
-            inputs=[audioone, audiotwo],
-            outputs=spectogram_wave2,
-            fn=analyze_double,
-            cache_examples=False
-        )
+    #     gr.Examples(
+    #         examples=example_pairs,
+    #         inputs=[audioone, audiotwo],
+    #         outputs=spectogram_wave2,
+    #         fn=analyze_double,
+    #         cache_examples=False
+    #     )
 
 
 app.launch()