-
Notifications
You must be signed in to change notification settings - Fork 115
/
datasets.py
48 lines (41 loc) · 1.6 KB
/
datasets.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import h5py
import math
import pandas as pd
from tensorflow.keras.utils import Sequence
import numpy as np
class ECGSequence(Sequence):
@classmethod
def get_train_and_val(cls, path_to_hdf5, hdf5_dset, path_to_csv, batch_size=8, val_split=0.02):
n_samples = len(pd.read_csv(path_to_csv))
n_train = math.ceil(n_samples*(1-val_split))
train_seq = cls(path_to_hdf5, hdf5_dset, path_to_csv, batch_size, end_idx=n_train)
valid_seq = cls(path_to_hdf5, hdf5_dset, path_to_csv, batch_size, start_idx=n_train)
return train_seq, valid_seq
def __init__(self, path_to_hdf5, hdf5_dset, path_to_csv=None, batch_size=8,
start_idx=0, end_idx=None):
if path_to_csv is None:
self.y = None
else:
self.y = pd.read_csv(path_to_csv).values
# Get tracings
self.f = h5py.File(path_to_hdf5, "r")
self.x = self.f[hdf5_dset]
self.batch_size = batch_size
if end_idx is None:
end_idx = len(self.x)
self.start_idx = start_idx
self.end_idx = end_idx
@property
def n_classes(self):
return self.y.shape[1]
def __getitem__(self, idx):
start = self.start_idx + idx * self.batch_size
end = min(start + self.batch_size, self.end_idx)
if self.y is None:
return np.array(self.x[start:end, :, :])
else:
return np.array(self.x[start:end, :, :]), np.array(self.y[start:end])
def __len__(self):
return math.ceil((self.end_idx - self.start_idx) / self.batch_size)
def __del__(self):
self.f.close()