| import os | |
| import pandas as pd | |
| import numpy as np | |
| from torch.utils.data import Dataset | |
| class XrdDataset(Dataset): | |
| def __init__(self, data_dir, annotations_file): | |
| self.labels = pd.read_csv(annotations_file) | |
| self.data_dir = data_dir | |
| def __len__(self): | |
| return len(self.labels) | |
| def __getitem__(self, idx): | |
| dataid = str(self.labels.iloc[idx, 0]) | |
| data_path = os.path.join(self.data_dir, dataid + '.csv') | |
| data_csv = pd.read_csv(data_path) | |
| data = data_csv.values.astype(np.float32).T | |
| label = self.labels.iloc[idx, 1] | |
| return data, label | |
| class mixDataset_cls_dynamic(Dataset): | |
| def __init__(self, data_dir, anno_struc, mode): | |
| self.data_dir = data_dir | |
| self.codIdList = pd.read_csv(anno_struc).values[:, 0].astype(np.int32) | |
| self.mode = mode | |
| def __len__(self): | |
| return 1000000 | |
| def __getitem__(self, idx): | |
| choice1, choice2 = np.random.randint(0, 23073, 2) | |
| if self.mode == 'train': | |
| rand1, rand2 = np.random.randint(1, 25, 2) | |
| else: | |
| rand1, rand2 = np.random.randint(1, 7, 2) | |
| data_path1 = os.path.join(self.data_dir, '{}_{}.csv'.format(self.codIdList[choice1], rand1)) | |
| data_path2 = os.path.join(self.data_dir, '{}_{}.csv'.format(self.codIdList[choice2], rand2)) | |
| data1 = pd.read_csv(data_path1).values.astype(np.float32).T | |
| data2 = pd.read_csv(data_path2).values.astype(np.float32).T | |
| ratio1 = np.random.randint(20, 81) | |
| ratio2 = 100 - ratio1 | |
| label = np.zeros(23073).astype(np.float32) | |
| label[choice1] = 0.4 | |
| label[choice2] = 0.4 | |
| return data1, data2, ratio1, ratio2, label | |