| 544 | |
| 545 | |
| 546 | class SWATSegLoader(Dataset): |
| 547 | def __init__(self, root_path, win_size, step=1, flag="train"): |
| 548 | self.flag = flag |
| 549 | self.step = step |
| 550 | self.win_size = win_size |
| 551 | self.scaler = StandardScaler() |
| 552 | |
| 553 | train_data = pd.read_csv(os.path.join(root_path, 'swat_train2.csv')) |
| 554 | test_data = pd.read_csv(os.path.join(root_path, 'swat2.csv')) |
| 555 | labels = test_data.values[:, -1:] |
| 556 | train_data = train_data.values[:, :-1] |
| 557 | test_data = test_data.values[:, :-1] |
| 558 | |
| 559 | self.scaler.fit(train_data) |
| 560 | train_data = self.scaler.transform(train_data) |
| 561 | test_data = self.scaler.transform(test_data) |
| 562 | self.train = train_data |
| 563 | self.test = test_data |
| 564 | self.val = test_data |
| 565 | self.test_labels = labels |
| 566 | print("test:", self.test.shape) |
| 567 | print("train:", self.train.shape) |
| 568 | |
| 569 | def __len__(self): |
| 570 | """ |
| 571 | Number of images in the object dataset. |
| 572 | """ |
| 573 | if self.flag == "train": |
| 574 | return (self.train.shape[0] - self.win_size) // self.step + 1 |
| 575 | elif (self.flag == 'val'): |
| 576 | return (self.val.shape[0] - self.win_size) // self.step + 1 |
| 577 | elif (self.flag == 'test'): |
| 578 | return (self.test.shape[0] - self.win_size) // self.step + 1 |
| 579 | else: |
| 580 | return (self.test.shape[0] - self.win_size) // self.win_size + 1 |
| 581 | |
| 582 | def __getitem__(self, index): |
| 583 | index = index * self.step |
| 584 | if self.flag == "train": |
| 585 | return np.float32(self.train[index:index + self.win_size]), np.float32(self.test_labels[0:self.win_size]) |
| 586 | elif (self.flag == 'val'): |
| 587 | return np.float32(self.val[index:index + self.win_size]), np.float32(self.test_labels[0:self.win_size]) |
| 588 | elif (self.flag == 'test'): |
| 589 | return np.float32(self.test[index:index + self.win_size]), np.float32( |
| 590 | self.test_labels[index:index + self.win_size]) |
| 591 | else: |
| 592 | return np.float32(self.test[ |
| 593 | index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]), np.float32( |
| 594 | self.test_labels[index // self.step * self.win_size:index // self.step * self.win_size + self.win_size]) |
| 595 | |
| 596 | |
| 597 | class UEAloader(Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected