File size: 1,772 Bytes
38f7d61 |
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 |
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
|