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
| class Mydata(Dataset): def __init__(self, file_path, img_file, mode = 'train', train_weight = 0.8): super().__init__() self.img_file = img_file data = pd.read_csv(file_path) self.mode = mode if mode == 'train': self.len = int(len(data) * train_weight) self.features = np.asarray(data.iloc[0:self.len,0]) self.labels = np.asarray(data.iloc[0:self.len,1]) elif mode == 'valid': self.pos = int(len(data) * (train_weight)) self.features = np.asarray(data.iloc[self.pos:,0]) self.labels = np.asarray(data.iloc[self.pos:,1]) self.len = self.features.shape[0] else: self.len = len(data) self.features = np.asarray(data.iloc[:,0]) def __getitem__(self, index): img = Image.open(self.img_file + self.features[index]) if self.mode == 'train': trans = transforms.Compose([ transforms.Resize((224,224)), transforms.RandomHorizontalFlip(p=0.5), transforms.ToTensor() ]) img = trans(img) label = class_to_num[self.labels[index]] return img, label elif self.mode == 'valid': trans = transforms.ToTensor() img = trans(img) label = class_to_num[self.labels[index]] return img, label else: trans = transforms.ToTensor() img = trans(img) return img def __len__(self): return self.len
|