new beat detection and more features

README.md

@@ -10,4 +10,4 @@ pinned: false
 license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,15 +1,17 @@
 from plotly.subplots import make_subplots
 from scipy.signal import find_peaks, butter, filtfilt
 import plotly.graph_objects as go
+from io import StringIO
 import pandas as pd
 import gradio as gr
 import numpy as np
 import itertools
+import tempfile
 import librosa
 import random
 import os
 
-from utils import getaudiodata, getBeats, plotBeattimes
+from utils import getaudiodata, getBeats, plotBeattimes, find_s1s2
 
 example_dir = "Examples"
 example_files = [os.path.join(example_dir, f) for f in os.listdir(example_dir) if f.endswith(('.wav', '.mp3', '.ogg'))]
@@ -185,34 +187,67 @@ def analyze_single(audio:gr.Audio):
 def getBeatsv2(audio:gr.Audio):
 
     sr, audiodata = getaudiodata(audio)
-    _, beattimes = getBeats(audiodata, sr)
+    _, beattimes, audiodata = getBeats(audiodata, sr)
 
-    fig = plotBeattimes(beattimes, audiodata, sr)
     beattimes_table = pd.DataFrame(data={"Beattimes":beattimes})
 
-    return fig, beattimes_table
+    feature_array = find_s1s2(beattimes_table)
+
+    featuredf = pd.DataFrame(
+        data=feature_array,
+        columns=[
+        "Beattimes",
+        "S1 to S2",
+        "S2 to S1",
+        "Label (S1=0/S2=1)"]
+    )
+
+    # Create boolean masks for each label
+    mask_ones = feature_array[:, 3] == 1
+    mask_zeros = feature_array[:, 3] == 0
+    
+    # Extract time/positions using the masks
+    times_label_one = feature_array[mask_ones, 0]
+    times_label_zero = feature_array[mask_zeros, 0]
+
+    fig = plotBeattimes(times_label_one, audiodata, sr, times_label_zero)
+
+
+    
+
+
+    return fig, featuredf, (sr, audiodata)
+
+def updateBeatsv2(beattimes_table:gr.Dataframe, audio:gr.Audio, uploadeddf:gr.File=None)-> go.Figure:
 
-def updateBeatsv2(beattimes_table:gr.Dataframe, audio:gr.Audio)-> go.Figure:
     sr, audiodata = getaudiodata(audio)
-    return plotBeattimes(beattimes_table["Beattimes"], audiodata, sr)
 
 
+    if uploadeddf != None:
+        beattimes_table = pd.read_csv(uploadeddf)
+
+    s1_times = beattimes_table[beattimes_table["Label (S1=0/S2=1)"] == 0]["Beattimes"].to_numpy()
+    s2_times = beattimes_table[beattimes_table["Label (S1=0/S2=1)"] == 1]["Beattimes"].to_numpy()
+
+    fig = plotBeattimes(s1_times, audiodata, sr, s2_times)
+
+    return fig, beattimes_table
+
+def download_df (df: pd.DataFrame):
 
+    temp_dir = tempfile.gettempdir()
+    temp_path = os.path.join(temp_dir, "feature_data.csv")
 
+    df.to_csv(temp_path, index=False)
 
+    return temp_path
 
 
 with gr.Blocks() as app:
 
     gr.Markdown("# Heartbeat")
     gr.Markdown("This App helps to analyze and extract Information from Heartbeat Audios")
-    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)
-    """)
+
 
     audiofile = gr.Audio(
                 type="filepath",
@@ -220,50 +255,44 @@ with gr.Blocks() as app:
                 sources="upload")
 
 
-    with gr.Tab("Single Audio V2"):
+    with gr.Tab("Preprocessing"):
         
         getBeatsbtn = gr.Button("get Beats")
+        cleanedaudio = gr.Audio(label="Cleaned Audio",show_download_button=True)
 
         beats_wave_plot = gr.Plot()
-        beattimes_table = gr.Dataframe(
-            col_count=1,
-            type='pandas',
-            interactive=True)
-        
-        updateBeatsbtn = gr.Button("update Beats")
-
 
+        beattimes_df = gr.Dataframe(
+            value = pd.DataFrame({"Beattimes":[], "S1 to S2":[], "S2 to S1":[], "Label (S1=0/S2=1)":[]}),
+            label="Beattimes")
+        
+        with gr.Row():
 
-        getBeatsbtn.click(getBeatsv2, inputs=audiofile, outputs=[beats_wave_plot, beattimes_table])
-        updateBeatsbtn.click(updateBeatsv2, inputs=[beattimes_table, audiofile], outputs=[beats_wave_plot])
-
-        gr.Examples(
-            examples=example_files,
-            inputs=audiofile,
-            fn=getBeatsv2,
-            cache_examples=False
+            csv_download = gr.DownloadButton()
+        
+            updateBeatsbtn = gr.Button("update Beats")
+        
+        uploadDF = gr.File(
+            file_count="single",
+            file_types=[".csv"],
+            label="upload a csv",
+            height=25
         )
 
-
-    with gr.Tab("Single Audio V1"):
-
-        analyzebtn = gr.Button("analyze")
-
-        results = gr.Markdown()
-        spectogram_wave = gr.Plot()
-        avg_beat_plot = gr.Plot()
-
-        analyzebtn.click(analyze_single, audiofile, [results, spectogram_wave, avg_beat_plot])
+        csv_download.click(download_df, inputs=[beattimes_df], outputs=[csv_download])
+        getBeatsbtn.click(getBeatsv2, inputs=audiofile, outputs=[beats_wave_plot, beattimes_df, cleanedaudio])
+        updateBeatsbtn.click(updateBeatsv2, inputs=[beattimes_df, audiofile, uploadDF], outputs=[beats_wave_plot, beattimes_df])
 
         gr.Examples(
             examples=example_files,
             inputs=audiofile,
-            outputs=[results, spectogram_wave],
-            fn=analyze_single,
+            fn=getBeatsv2,
             cache_examples=False
         )
 
+    with gr.Tab("Analysis"):
 
+        gr.Markdown("🚨 Please make sure to first run the 'Preprocessing'")
 
 
 app.launch()

utils.py

@@ -1,76 +1,534 @@
 import librosa
 import numpy as np
 import plotly.graph_objects as go
-from scipy.signal import find_peaks
+from scipy.signal import savgol_filter, find_peaks
+from scipy.signal import butter, filtfilt, medfilt, find_peaks, hilbert
+from scipy.ndimage import gaussian_filter1d
+from sklearn.cluster import KMeans
+from sklearn.preprocessing import StandardScaler
+from io import StringIO
+import soundfile as sf
+import pywt
+import pandas as pd
 
 
 # GENERAL HELPER FUNCTIONS
-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):
-        audiodata = np.array(audiodata)
+def denoise_audio(audiodata: np.ndarray, sr: int) -> tuple[np.ndarray, int]:
+    """
+    Enhanced denoising of audio signals optimized for heart sounds.
+    Uses a combination of bandpass filtering, adaptive wavelet denoising,
+    and improved spectral subtraction.
+    
+    Parameters:
+    -----------
+    audiodata : np.ndarray
+        Input audio signal (1D numpy array)
+    sr : int
+        Sampling rate in Hz
+        
+    Returns:
+    --------
+    tuple[np.ndarray, int]
+        Tuple containing (denoised_signal, sampling_rate)
+    """
+    # Input validation and conversion
+    if not isinstance(audiodata, np.ndarray) or audiodata.ndim != 1:
+        raise ValueError("audiodata must be a 1D numpy array")
+    if not isinstance(sr, int) or sr <= 0:
+        raise ValueError("sr must be a positive integer")
+        
+    # Convert to float32 and normalize
+    audio = audiodata.astype(np.float32)
+    audio = audio / np.max(np.abs(audio))
+    
+    # 1. Enhanced Bandpass Filter
+    # Optimize frequency range for heart sounds (20-200 Hz)
+    nyquist = sr / 2
+    low, high = 20 / nyquist, 200 / nyquist
+    order = 4  # Filter order
+    b, a = butter(order, [low, high], btype='band')
+    filtered = filtfilt(b, a, audio)
+    
+    # 2. Adaptive Wavelet Denoising
+    def apply_wavelet_denoising(sig):
+        # Use sym4 wavelet (good for biomedical signals)
+        wavelet = 'sym4'
+        level = min(6, pywt.dwt_max_level(len(sig), pywt.Wavelet(wavelet).dec_len))
+        
+        # Decompose signal
+        coeffs = pywt.wavedec(sig, wavelet, level=level)
+        
+        # Adaptive thresholding based on level
+        for i in range(1, len(coeffs)):
+            # Calculate level-dependent threshold
+            sigma = np.median(np.abs(coeffs[i])) / 0.6745
+            threshold = sigma * np.sqrt(2 * np.log(len(coeffs[i])))
+            # Adjust threshold based on decomposition level
+            level_factor = 1 - (i / len(coeffs))  # Higher levels get lower thresholds
+            coeffs[i] = pywt.threshold(coeffs[i], threshold * level_factor, mode='soft')
+        
+        return pywt.waverec(coeffs, wavelet)
+    
+    # Apply wavelet denoising
+    denoised = apply_wavelet_denoising(filtered)
+    
+    # Ensure consistent length
+    if len(denoised) != len(audio):
+        denoised = librosa.util.fix_length(denoised, len(audio))
+    
+    # 3. Improved Spectral Subtraction
+    def spectral_subtract(sig):
+        # Parameters
+        frame_length = int(sr * 0.04)  # 40ms frames
+        hop_length = frame_length // 2
+        
+        # Compute STFT
+        D = librosa.stft(sig, n_fft=frame_length, hop_length=hop_length)
+        mag, phase = np.abs(D), np.angle(D)
+        
+        # Estimate noise spectrum from low-energy frames
+        frame_energy = np.sum(mag**2, axis=0)
+        noise_threshold = np.percentile(frame_energy, 15)
+        noise_frames = mag[:, frame_energy < noise_threshold]
+        
+        if noise_frames.size > 0:
+            noise_spectrum = np.median(noise_frames, axis=1)
+            
+            # Oversubtraction factor (frequency-dependent)
+            freq_bins = np.fft.rfftfreq(frame_length, 1/sr)
+            alpha = 1.0 + 0.01 * (freq_bins / nyquist)
+            alpha = alpha[:len(noise_spectrum)].reshape(-1, 1)
+            
+            # Spectral subtraction with flooring
+            mag_clean = np.maximum(mag - alpha * noise_spectrum.reshape(-1, 1), 0.01 * mag)
+            
+            # Reconstruct signal
+            D_clean = mag_clean * np.exp(1j * phase)
+            return librosa.istft(D_clean, hop_length=hop_length)
+        
+        return sig
+    
+    # Apply spectral subtraction
+    final = spectral_subtract(denoised)
+    
+    # Final normalization
+    final = final / np.max(np.abs(final))
+    
+    return final, sr
 
-    # Check if audio is mono or stereo
+def getaudiodata(filepath: str, target_sr: int = 16000) -> tuple[int, np.ndarray]:
+    """
+    Load and process audio data with consistent output properties.
+    
+    Parameters:
+    -----------
+    filepath : str
+        Path to the audio file
+    target_sr : int
+        Target sampling rate (default: 16000 Hz)
+        
+    Returns:
+    --------
+    tuple[int, np.ndarray]
+        Sampling rate and processed audio data with consistent properties:
+        - dtype: float32
+        - shape: (N,) mono audio
+        - amplitude range: [-0.95, 0.95]
+        - no NaN or Inf values
+        - C-contiguous memory layout
+    """
+    # Load audio with specified sampling rate
+    audiodata, sr = librosa.load(filepath, sr=target_sr)
+    
+    # Ensure numpy array
+    audiodata = np.asarray(audiodata)
+    
+    # Convert to mono if 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)
-
+    
+    # Handle any NaN or Inf values
+    audiodata = np.nan_to_num(audiodata, nan=0.0, posinf=0.0, neginf=0.0)
+    
+    # Normalize to prevent clipping while maintaining relative amplitudes
+    max_abs = np.max(np.abs(audiodata))
+    if max_abs > 0:  # Avoid division by zero
+        audiodata = audiodata * (0.95 / max_abs)
+    
+    # Ensure float32 dtype and memory contiguous
+    audiodata = np.ascontiguousarray(audiodata, dtype=np.float32)
+    
     return sr, audiodata
 
-def getBeats(audiodata:np.ndarray, sr:int):
-    # Convert audio data to float32
-    audiodata = audiodata.astype(np.float32)
+def getBeats(audiodata: np.ndarray, sr: int, method='envelope') -> tuple[float, np.ndarray, np.ndarray]:
+    """
+    Advanced heartbeat detection optimized for peak detection with improved sensitivity.
     
-    # Normalize the audio data
-    audiodata = audiodata / np.max(np.abs(audiodata))
+    Parameters:
+    -----------
+    audiodata : np.ndarray
+        Audio time series
+    sr : int
+        Sampling rate
+    method : str
+        Detection method: 'onset', 'envelope', 'fusion' (default)
+        
+    Returns:
+    --------
+    tempo : float
+        Estimated heart rate in BPM
+    peak_times : np.ndarray
+        Times of detected heartbeat peaks
+    cleaned_audio : np.ndarray
+        Cleaned audio signal
+    """
+    # Denoise and normalize
+    audiodata, sr = denoise_audio(audiodata, sr)
+    cleaned_audio = audiodata / np.max(np.abs(audiodata))
     
-    # Set the threshold for peak detection (adjust this value as needed)
-    threshold = 0.5  # 50% of the maximum amplitude
+    def get_envelope_peaks():
+        """Detect peaks using enhanced envelope method with better sensitivity"""
+        # Calculate envelope using appropriate frame sizes
+        hop_length = int(sr * 0.01)  # 10ms hop
+        frame_length = int(sr * 0.04)  # 40ms window
+        
+        # Calculate RMS energy
+        rms = librosa.feature.rms(
+            y=cleaned_audio,
+            frame_length=frame_length,
+            hop_length=hop_length
+        )[0]
+        
+        # Smooth the envelope (less aggressive smoothing)
+        rms_smooth = savgol_filter(rms, 7, 3)
+        
+        # Find peaks with more lenient thresholds
+        peaks, properties = find_peaks(
+            rms_smooth,
+            distance=int(0.2 * (sr / hop_length)),  # Minimum 0.2s between peaks (300 BPM max)
+            height=np.mean(rms_smooth) + 0.1 * np.std(rms_smooth),  # Lower height threshold
+            prominence=np.mean(rms_smooth) * 0.1,  # Lower prominence threshold
+            width=(int(0.01 * (sr / hop_length)), int(0.2 * (sr / hop_length)))  # 10-200ms width
+        )
+        
+        # Refine peak locations using original signal
+        refined_peaks = []
+        window_size = int(0.05 * sr)  # 50ms window for refinement
+        
+        for peak in peaks:
+            # Convert envelope peak to sample domain
+            sample_idx = peak * hop_length
+            
+            # Define window boundaries
+            start = max(0, sample_idx - window_size//2)
+            end = min(len(cleaned_audio), sample_idx + window_size//2)
+            
+            # Find the maximum amplitude within the window
+            window = np.abs(cleaned_audio[int(start):int(end)])
+            max_idx = np.argmax(window)
+            refined_peaks.append(start + max_idx)
+        
+        return np.array(refined_peaks), rms_smooth
+
+    def get_onset_peaks():
+        """Enhanced onset detection with better sensitivity"""
+        # Multi-band onset detection with adjusted parameters
+        onset_env = librosa.onset.onset_strength(
+            y=cleaned_audio, 
+            sr=sr,
+            hop_length=256,  # Smaller hop length for better temporal resolution
+            aggregate=np.median,
+            n_mels=128
+        )
+        
+        # More lenient thresholding
+        threshold = np.mean(onset_env) + 0.3 * np.std(onset_env)
+        
+        # Get onset positions
+        onset_frames = librosa.onset.onset_detect(
+            onset_envelope=onset_env,
+            sr=sr,
+            hop_length=256,
+            backtrack=True,
+            threshold=threshold,
+            pre_max=20,  # 20 frames before peak
+            post_max=20,  # 20 frames after peak
+            pre_avg=25,   # 25 frames before for mean
+            post_avg=25,  # 25 frames after for mean
+            wait=10       # Wait 10 frames before detecting next onset
+        )
+        
+        # Refine onset positions to peaks
+        refined_peaks = []
+        window_size = int(0.05 * sr)  # 50ms window
+        
+        for frame in onset_frames:
+            # Convert frame to sample index
+            sample_idx = frame * 256  # Using hop_length=256
+            
+            # Define window boundaries
+            start = max(0, sample_idx - window_size//2)
+            end = min(len(cleaned_audio), sample_idx + window_size//2)
+            
+            # Find the maximum amplitude within the window
+            window = np.abs(cleaned_audio[int(start):int(end)])
+            max_idx = np.argmax(window)
+            refined_peaks.append(start + max_idx)
+        
+        return np.array(refined_peaks), onset_env
     
-    # Find peaks above the threshold
-    peaks, _ = find_peaks(np.abs(audiodata), height=threshold, distance=int(sr * 0.3))
+    # Apply selected method
+    if method == 'envelope':
+        peaks, _ = get_envelope_peaks()
+
+    elif method == 'onset':
+        peaks, _ = get_onset_peaks()
+
     
-    # Convert peak indices to times
-    peak_times = (peaks / sr)*2
+    else:  # fusion method
+        # Get peaks from both methods
+        env_peaks, _ = get_envelope_peaks()
+        onset_peaks, _ = get_onset_peaks()
+        
+        # Merge nearby peaks (within 50ms)
+        all_peaks = np.sort(np.concatenate([env_peaks, onset_peaks]))
+        merged_peaks = []
+        last_peak = -np.inf
+        
+        for peak in all_peaks:
+            if (peak - last_peak) / sr > 0.05:  # 50ms minimum separation
+                merged_peaks.append(peak)
+                last_peak = peak
+        
+        peaks = np.array(merged_peaks)
+
+    # Convert peaks to times
+    peak_times = peaks / sr
     
-    # Calculate tempo (beats per minute)
+    # Calculate tempo using peak times
     if len(peak_times) > 1:
-        avg_interval = np.mean(np.diff(peak_times))
-        tempo = 60 / avg_interval
+        # Use weighted average of intervals
+        intervals = np.diff(peak_times)
+        tempos = 60 / intervals  # Convert intervals to BPM
+        
+        # Remove physiologically impossible tempos
+        valid_tempos = tempos[(tempos >= 30) & (tempos <= 300)]
+        
+        if len(valid_tempos) > 0:
+            tempo = np.median(valid_tempos)  # Use median for robustness
+        else:
+            tempo = 0
     else:
         tempo = 0
-    
-    return [tempo], peak_times
 
-def plotBeattimes(beattimes:np.ndarray, audiodata:np.ndarray, sr:int)->go.Figure:
+    return tempo, peak_times, cleaned_audio
 
-        # Time array for the full audio
-    time = (np.arange(0, len(audiodata)) / sr) * 2
+# def plotBeattimes(beattimes: np.ndarray, audiodata: np.ndarray, sr: int) -> go.Figure:
+#     """
+#     Plot audio waveform with beat markers.
+    
+#     Parameters:
+#     -----------
+#     beattimes : np.ndarray
+#         Array of beat times in seconds
+#     audiodata : np.ndarray
+#         Audio time series data
+#     sr : int
+#         Sampling rate
+        
+#     Returns:
+#     --------
+#     go.Figure
+#         Plotly figure with waveform and beat markers
+#     """
+#     # Calculate correct time array for the full audio
+#     time = np.arange(len(audiodata)) / sr
+    
+#     # Convert beat times to sample indices
+#     beat_indices = np.round(beattimes * sr).astype(int)
+    
+#     # Ensure indices are within bounds
+#     beat_indices = beat_indices[beat_indices < len(audiodata)]
+    
+#     # Get actual amplitudes at beat positions
+#     beat_amplitudes = audiodata[beat_indices]
+    
+#     # Create the figure
+#     fig = go.Figure()
+    
+#     # Add waveform
+#     fig.add_trace(
+#         go.Scatter(
+#             x=time,
+#             y=audiodata,
+#             mode='lines',
+#             name='Waveform',
+#             line=dict(color='blue', width=1)
+#         )
+#     )
+    
+#     # Add beat markers
+#     fig.add_trace(
+#         go.Scatter(
+#             x=beattimes[beat_indices < len(audiodata)],  # Use filtered beat times
+#             y=beat_amplitudes,
+#             mode='markers',
+#             name='Beats',
+#             marker=dict(
+#                 color='red',
+#                 size=8,
+#                 symbol='circle',
+#                 line=dict(color='darkred', width=1)
+#             )
+#         )
+#     )
+    
+#     # Update layout
+#     fig.update_layout(
+#         title="Audio Waveform with Beat Detection",
+#         xaxis_title="Time (seconds)",
+#         yaxis_title="Amplitude",
+#         showlegend=False,
+#         hovermode='closest',
+#         plot_bgcolor='white',
+#         legend=dict(
+#             yanchor="top",
+#             y=0.99,
+#             xanchor="left",
+#             x=0.01
+#         )
+#     )
+    
+#     # Add vertical lines at beat positions (optional)
+#     for beat_time in beattimes[beat_indices < len(audiodata)]:
+#         fig.add_vline(
+#             x=beat_time,
+#             line=dict(color="rgba(255, 0, 0, 0.2)", width=1),
+#             layer="below"
+#         )
+    
+#     return fig
 
-    # CREATE BEATTIMES PLOT    
-    # Waveform plot
-    fig = go.Figure(
-        go.Scatter(x=time, y=audiodata, mode='lines', name='Waveform', line=dict(color='blue', width=1))
+def plotBeattimes(beattimes: np.ndarray, 
+                audiodata: np.ndarray, 
+                sr: int, 
+                beattimes2: np.ndarray = None) -> go.Figure:
+    """
+    Plot audio waveform with beat markers for one or two sets of beat times.
+    
+    Parameters:
+    -----------
+    beattimes : np.ndarray
+        Primary array of beat times in seconds (S1 beats if beattimes2 is provided)
+    audiodata : np.ndarray
+        Audio time series data
+    sr : int
+        Sampling rate
+    beattimes2 : np.ndarray, optional
+        Secondary array of beat times in seconds (S2 beats)
+        
+    Returns:
+    --------
+    go.Figure
+        Plotly figure with waveform and beat markers
+    """
+    # Calculate time array for the full audio
+    time = np.arange(len(audiodata)) / sr
+    
+    # Create the figure
+    fig = go.Figure()
+    
+    # Add waveform
+    fig.add_trace(
+        go.Scatter(
+            x=time,
+            y=audiodata,
+            mode='lines',
+            name='Waveform',
+            line=dict(color='blue', width=1)
+        )
     )
-    # Add beat markers
-    beat_amplitudes = np.interp(beattimes, time, audiodata)
+    
+    # Process and plot primary beat times
+    beat_indices = np.round(beattimes * sr).astype(int)
+    beat_indices = beat_indices[beat_indices < len(audiodata)]
+    beat_amplitudes = audiodata[beat_indices]
+    
+    # Define beat name based on whether secondary beats are provided
+    beat_name = "Beats S1" if beattimes2 is not None else "Beats"
+    
+    # Add primary beat markers
     fig.add_trace(
-        go.Scatter(x=beattimes, y=beat_amplitudes, mode='markers', name='Beats',
-                   marker=dict(color='red', size=8, symbol='circle'))
+        go.Scatter(
+            x=beattimes[beat_indices < len(audiodata)],
+            y=beat_amplitudes,
+            mode='markers',
+            name=beat_name,
+            marker=dict(
+                color='red',
+                size=8,
+                symbol='circle',
+                line=dict(color='darkred', width=1)
+            )
+        )
     )
-
+    
+    # Add primary beat vertical lines
+    for beat_time in beattimes[beat_indices < len(audiodata)]:
+        fig.add_vline(
+            x=beat_time,
+            line=dict(color="rgba(255, 0, 0, 0.2)", width=1),
+            layer="below"
+        )
+    
+    # Process and plot secondary beat times if provided
+    if beattimes2 is not None:
+        beat_indices2 = np.round(beattimes2 * sr).astype(int)
+        beat_indices2 = beat_indices2[beat_indices2 < len(audiodata)]
+        beat_amplitudes2 = audiodata[beat_indices2]
+        
+        # Add secondary beat markers
+        fig.add_trace(
+            go.Scatter(
+                x=beattimes2[beat_indices2 < len(audiodata)],
+                y=beat_amplitudes2,
+                mode='markers',
+                name="Beats S2",
+                marker=dict(
+                    color='green',
+                    size=8,
+                    symbol='circle',
+                    line=dict(color='darkgreen', width=1)
+                )
+            )
+        )
+        
+        # Add secondary beat vertical lines
+        for beat_time in beattimes2[beat_indices2 < len(audiodata)]:
+            fig.add_vline(
+                x=beat_time,
+                line=dict(color="rgba(0, 255, 0, 0.2)", width=1),
+                layer="below"
+            )
+    
+    # Update layout
     fig.update_layout(
-        showlegend=False
+        title="Audio Waveform with Beat Detection",
+        xaxis_title="Time (seconds)",
+        yaxis_title="Amplitude",
+        showlegend=True,  # Changed to True to show beat types
+        hovermode='closest',
+        plot_bgcolor='white',
+        legend=dict(
+            yanchor="top",
+            y=0.99,
+            xanchor="left",
+            x=0.01
+        )
     )
-
+    
     return fig
 
-
 def iterate_beat_segments(beat_times, sr, audio):
     """
     Iterate over audio segments between beats.
@@ -96,8 +554,6 @@ def iterate_beat_segments(beat_times, sr, audio):
 
         segment_metrics = segment_analysis(segment, sr)
 
-
-
 def segment_analysis(segment, sr):
     """
     Analyze an audio segment and compute various metrics.
@@ -138,4 +594,41 @@ def segment_analysis(segment, sr):
         duration,
         s1_to_s2_duration,
         s2_to_s1_duration
-    ]
+    ]
+
+def find_s1s2(df:pd.DataFrame):
+
+
+    times = df['Beattimes'].to_numpy()
+    n_peaks = len(times)
+    
+    # Initialize the feature array
+    feature_array = np.zeros((n_peaks, 4))
+
+    # Fill in the peak times (first column)
+    feature_array[:, 0] = times
+
+    # Calculate and fill distances to previous peaks (second column)
+    feature_array[1:, 1] = np.diff(times)  # For all except first peak
+    feature_array[0, 1] = feature_array[1, 1]  # First peak uses same as second
+    
+    # Calculate and fill distances to next peaks (third column)
+    feature_array[:-1, 2] = np.diff(times)  # For all except last peak
+    feature_array[-1, 2] = feature_array[-2, 2]  # Last peak uses same as second-to-last
+    
+    # Extract features (distances to prev and next peaks)
+    X = feature_array[:, 1:3]
+    
+    # Scale features
+    scaler = StandardScaler()
+    X_scaled = scaler.fit_transform(X)
+    
+    # Apply K-means clustering
+    kmeans = KMeans(n_clusters=2, random_state=42)
+    labels = kmeans.fit_predict(X_scaled)
+    
+    # Update the labels in the feature array
+    feature_array[:, 3] = labels
+    
+    return feature_array
+