Extracting random fixed sized feature windows from variable input sizes in Audio with a Tensorflow dataset pipeline

I have list of strong labeled audio files of variable length, i.e., with sound events with known start and end, with which I am training a Sound Event Detection model.

For this purpose I previously used a Data Generator in which I

1. extracted random fixed-size feature windows from the audio files,
2. transformed the audio slides to melspectrograms,
3. and then created a multi hot vector from the list of sound events, so that the model received as input (batch_size, time_steps, freq_bins, n_channels) and (batch_size, time_steps, n_classes) as label.

For performance reasons, I am now trying to create a tf.data pipeline in which I am fist loading the variable sized audio files and labels:

def paths_and_labels_to_dataset(file_path):
    """Constructs a dataset of audios and labels."""
    # 1) Read X (audio wave) from file_path
    path_ds  = tf.data.Dataset.from_tensor_slices(file_path)
    audio_ds = path_ds.map(lambda x: path_to_audio(x), num_parallel_calls=AUTO)

    # 2) Convert sound event labels of shape (-1, 3) to ragged tensors
    labels = [[[0,     1.5, 0], # onset, offset, class_id
               [9.85,  10,  1]]]

    rt        = tf.ragged.constant(labels)
    label_ds  = tf.data.Dataset.from_tensor_slices(rt)
    
    return tf.data.Dataset.zip((audio_ds, label_ds))

def path_to_audio(path):
    """Reads and decodes an audio file."""
    audio    = tf.io.read_file(path)
    audio, _ = tf.audio.decode_wav(audio, 1)    
    return audio

and then cache them. For there on, I am extracting random audio windows of a fixed size with:

def limit_seconds_audio(x_audio, y):
    len_audio = len(x_audio)
    effective_length = SAMPLE_RATE * SECONDS
    
    # Shorter audio: Pad the end with zeros
    if len_audio <= effective_length:
        x_audio = tf.concat([x_audio, tf.expand_dims(tf.zeros([effective_length - len_audio]), 1)], axis=0)
        
    # Longer audio: Random crop a window size of SECONDS * SAMPLE_RATE
    else:
        start   = tf.random.uniform([1], maxval = len_audio - effective_length, dtype = tf.int32)[0]
        x_audio = x_audio[start : start+effective_length]
    
    return x_audio, y

After this, x_audio is transformed to melspectograms.

Now my issue is, that I need to transform the sound event label list of the shape:

labels = [[[0,     1.5, 0], # onset seconds, offset seconds, class_id
           [9.85,  10,  1]]]

to (time_steps, n_classes) for the given audio window from limit_seconds_audio.

In my DataGenerator, I simply would slice the sound event label list for a given audio with get_event_slice with the start and stop seconds:

def get_event_slice(signal, events_for_audiofile, start, stop, sample_rate = 4_000):
    signal     = signal[int(start):int(stop)]
    signal_len = len(signal) / sample_rate # in seconds
    events     = [item for item in events_for_audiofile if item[0] > start and item[0] <= stop]
    new_events = []
    
    for ev in events:
        new_events.append([
            (ev[0] - start),
            (ev[1] - start),
            ev[2]
        ])
    
    # truncate to the signal length
    if new_events[-1][1] / sample_rate > signal_len:
        new_events[-1][1] = len(signal)
 
    return signal, new_events

and then compute the multi hot vector by convert_labels_to_timesteps, where timesteps corresponds to the width dimension of the melspectogram:

def out_frames(sec, timesteps = 622, max_seconds = 10):
    return sec * (timesteps / max_seconds)

def convert_labels_to_timesteps(x_audio, y, timesteps = 622, num_classes = 8):
    # y is of shape (-1, 3), see above
    
    label       = np.zeros((timesteps, num_classes))
    
    # seconds to frames
    frame_start = np.floor(out_frames(y[:, 0])).astype(int)
    frame_end   = np.ceil(out_frames(y[:, 1])).astype(int)
    
    # iterate each class
    se_class    = y[:, 2].astype(int)

    for ind, val in enumerate(se_class):
        label[frame_start[ind]:frame_end[ind], val] = 1 # 1 for active

    return x_audio, label

resulting in (time_steps, n_classes), which I have not been able to do within a tf.data pipeline.

Any help or hint to 1) slice the ragged tensors for the respective feature window and then 2) transform then into the shape (time_steps, n_classes) is greatly appreciated.