-
Notifications
You must be signed in to change notification settings - Fork 0
/
dataloader.py
139 lines (93 loc) · 3.76 KB
/
dataloader.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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# -*- coding: utf-8 -*-
"""
Created on Fri Sep 8 18:57:06 2023
@author: yexin
"""
import os
import numpy as np
import torch
from torch.utils.data import Dataset
import pickle
from tqdm import tqdm
import torch.nn.functional as F
import random
class data_train_10(Dataset):
def __init__(self, cap):
self.cap = cap
# Assuming each item in cap can give at least 20 readings
self.total_sequences = sum([len(item) - 9 for item in cap])
def __len__(self):
return self.total_sequences
def __getitem__(self, idx):
cumulative_lengths = np.cumsum([len(item) - 9 for item in self.cap])
num = np.searchsorted(cumulative_lengths, idx + 1)
# If it's the first sequence, start from 0, else determine the starting point based on idx
if num == 0:
start_idx = idx
else:
start_idx = idx - cumulative_lengths[num - 1]
capacitance = self.cap[num][start_idx: start_idx + 10]
return torch.tensor(capacitance), torch.tensor(num), torch.tensor(start_idx)
class data_test_10(Dataset):
def __init__(self, cap):
self.cap = cap
# Assuming each item in cap can give at least 20 readings
self.total_sequences = sum([len(item) - 9 for item in cap])
def __len__(self):
return self.total_sequences
def __getitem__(self, idx):
cumulative_lengths = np.cumsum([len(item) - 9 for item in self.cap])
num = np.searchsorted(cumulative_lengths, idx + 1)
# If it's the first sequence, start from 0, else determine the starting point based on idx
if num == 0:
start_idx = idx
else:
start_idx = idx - cumulative_lengths[num - 1]
capacitance = self.cap[num][start_idx: start_idx + 10]
return torch.tensor(capacitance), torch.tensor(num), torch.tensor(start_idx)
class data_train_bc(Dataset):
def __init__(self, cap, b):
self.cap = cap
self.b = b
# Assuming each item in cap can give at least 20 readings
self.total_sequences = sum([len(item) - 9 for item in cap])
def __len__(self):
return self.total_sequences
def __getitem__(self, idx):
cumulative_lengths = np.cumsum([len(item) - 9 for item in self.cap])
num = np.searchsorted(cumulative_lengths, idx + 1)
# If it's the first sequence, start from 0, else determine the starting point based on idx
if num == 0:
start_idx = idx
else:
start_idx = idx - cumulative_lengths[num - 1]
capacitance = self.cap[num][start_idx: start_idx + 10]
bh = self.b[num][start_idx: start_idx + 10]
if np.mean(bh) > 0.01:
a = 1
else:
a = 0
return torch.tensor(capacitance), torch.tensor(num), torch.tensor(a)
class data_test_bc(Dataset):
def __init__(self, cap, b):
self.cap = cap
self.b = b
# Assuming each item in cap can give at least 20 readings
self.total_sequences = sum([len(item) - 9 for item in cap])
def __len__(self):
return self.total_sequences
def __getitem__(self, idx):
cumulative_lengths = np.cumsum([len(item) - 9 for item in self.cap])
num = np.searchsorted(cumulative_lengths, idx + 1)
# If it's the first sequence, start from 0, else determine the starting point based on idx
if num == 0:
start_idx = idx
else:
start_idx = idx - cumulative_lengths[num - 1]
capacitance = self.cap[num][start_idx: start_idx + 10]
bh = self.b[num][start_idx: start_idx + 10]
if np.mean(bh) > 0.01:
a = 1
else:
a = 0
return torch.tensor(capacitance), torch.tensor(num), torch.tensor(a)