| 503 | |
| 504 | |
| 505 | class SMDSegLoader(Dataset): |
| 506 | def __init__(self, root_path, win_size, step=100, flag="train"): |
| 507 | self.flag = flag |
| 508 | self.step = step |
| 509 | self.win_size = win_size |
| 510 | self.scaler = StandardScaler() |
| 511 | data = np.load(os.path.join(root_path, "SMD_train.npy")) |
| 512 | self.scaler.fit(data) |
| 513 | data = self.scaler.transform(data) |
| 514 | test_data = np.load(os.path.join(root_path, "SMD_test.npy")) |
| 515 | self.test = self.scaler.transform(test_data) |
| 516 | self.train = data |
| 517 | data_len = len(self.train) |
| 518 | self.val = self.train[(int)(data_len * 0.8):] |
| 519 | self.test_labels = np.load(os.path.join(root_path, "SMD_test_label.npy")) |
| 520 | |
| 521 | def __len__(self): |
| 522 | if self.flag == "train": |
| 523 | return (self.train.shape[0] - self.win_size) // self.step + 1 |
| 524 | elif (self.flag == 'val'): |
| 525 | return (self.val.shape[0] - self.win_size) // self.step + 1 |
| 526 | elif (self.flag == 'test'): |
| 527 | return (self.test.shape[0] - self.win_size) // self.step + 1 |
| 528 | else: |
| 529 | return (self.test.shape[0] - self.win_size) // self.win_size + 1 |
| 530 | |
| 531 | def __getitem__(self, index): |
| 532 | index = index * self.step |
| 533 | if self.flag == "train": |
| 534 | return np.float32(self.train[index:index + self.win_size]), np.float32(self.test_labels[0:self.win_size]) |
| 535 | elif (self.flag == 'val'): |
| 536 | return np.float32(self.val[index:index + self.win_size]), np.float32(self.test_labels[0:self.win_size]) |
| 537 | elif (self.flag == 'test'): |
| 538 | return np.float32(self.test[index:index + self.win_size]), np.float32( |
| 539 | self.test_labels[index:index + self.win_size]) |
| 540 | else: |
| 541 | return np.float32(self.test[ |
| 542 | index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32( |
| 543 | self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) |
| 544 | |
| 545 | |
| 546 | class SWATSegLoader(Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected