Aluode's picture
Upload folder using huggingface_hub
3bb804c verified
raw
history blame
2 kB
from typing import Any, Generic, Optional, TypeVar, Union
import numpy as np
from numpy.typing import NDArray
_DigitalDtype = TypeVar("_DigitalDtype", bound=Union[np.int16, np.int32])
class LazyLoader(Generic[_DigitalDtype]):
"""
Class to load data for a single signal from a buffer of EDF data records (array or memmap).
Parameters
----------
buffer : numpy.ndarray or numpy.memmap
Buffer of EDF data records.
start_sample : int
Offset of the signal in the data records (in samples).
end_sample : int
End of the signal in the data records (in samples).
"""
def __init__(
self,
buffer: Union[NDArray[_DigitalDtype], np.memmap[Any, np.dtype[_DigitalDtype]]],
start_sample: int,
end_sample: int,
) -> None:
self.buffer = buffer
self.start_sample = start_sample
self.end_sample = end_sample
def load(
self, start_record: Optional[int] = None, end_record: Optional[int] = None
) -> NDArray[_DigitalDtype]:
"""
Load signal data from the buffer.
Parameters
----------
start_record : int, optional
The first EDF data record to load samples from. If None, load from the beginning of the buffer.
end_record : int, optional
The last EDF data record to load samples from. If None, load until the end of the buffer.
Returns
-------
numpy.ndarray
Signal data (digital).
"""
if start_record is None:
start_record = 0
if end_record is None:
end_record = self.buffer.shape[0]
if (
end_record < start_record
or start_record < 0
or end_record > self.buffer.shape[0]
):
raise ValueError("Invalid slice: Slice exceeds EDF duration")
return self.buffer[
start_record:end_record, self.start_sample : self.end_sample
].flatten()