-
Notifications
You must be signed in to change notification settings - Fork 0
/
custom_dataset.py
29 lines (22 loc) · 948 Bytes
/
custom_dataset.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
# -*- coding: utf-8 -*-
"""Добро пожаловать в Colaboratory!
Automatically generated by Colab.
Original file is located at
https://colab.research.google.com/notebooks/intro.ipynb
"""
class CustomDataset(Dataset):
def __init__(self, data, is_test=False):
self.data = data
self.is_test = is_test
def get_channels(self, Np, Nr):
path = f"/kaggle/input/sleeps-converted/converted/User-{int(Np)}-Night-{int(Nr)}.npy"
return np.load(path)
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
channels = torch.tensor(self.get_channels(*self.data.iloc[idx, [0, -7]]))
item = torch.tensor(self.data.iloc[idx].drop(["AI", "HI", "AHI"]).values)
if self.is_test:
return (item, channels)
return (item, channels[:, :5000000]), torch.tensor(self.data.iloc[idx,[-5, -4, -3]])
#min value or zero value