ria-toolkit-oss/src/ria_toolkit_oss/annotations/qualify_slice.py

import numpy as np

from ria_toolkit_oss.datatypes import Recording


def qualify_slice_from_annotations(recording: Recording, slice_length: int):
    """
    Slice a recording into many smaller recordings,
    discarding any slices which do not have annotations that apply to those samples.
    Used together with an annotation based qualifier.

    :param recording: The recording to slice.
    :type recording: Recording
    :param slice_length: The length in samples of a slice.
    :type slice_length: int"""

    if len(recording.annotations) == 0:
        print("Warning, no annotations.")

    annotation_mask = np.zeros(len(recording.data[0]))

    for annotation in recording.annotations:
        annotation_mask[annotation.sample_start : annotation.sample_start + annotation.sample_count] = 1

    output_recordings = []

    for i in range((len(recording.data[0]) // slice_length) - 1):
        start_index = slice_length * i
        end_index = slice_length * (i + 1)

        if 1 in annotation_mask[start_index:end_index]:
            sl = recording.data[:, start_index:end_index]
            output_recordings.append(Recording(data=sl, metadata=recording.metadata))

    return output_recordings