def make_pipeline(imageset_path, output_dir): """returns an augmentation pipeline for a given image set""" p = Pipeline(imageset_path, output_dir) p.random_distortion(probability=0.7, grid_width=4, grid_height=4, magnitude=8) p.flip_left_right(probability=0.5) p.flip_top_bottom(probability=0.5) p.zoom(probability=0.3, min_factor=1.1, max_factor=1.4) p.rotate(probability=0.5, max_left_rotation=10, max_right_rotation=10) return p
class Dataset(data.Dataset): def __init__(self, train=True, root='../data'): if train: self.dataset = MNIST(root=root, train=True) else: self.dataset = MNIST(root=root, train=False) self.transform = transforms.Compose( [transforms.ToTensor()] ) self.p = Pipeline() self.p.random_distortion(probability=0.5, grid_width=7, grid_height=7, magnitude=1) def __len__(self,): return len(self.dataset) def __getitem__(self, i): img = self.dataset[i][0] lab = self.dataset[i][1] # img = self._sample([img], self.p)[0] img = self.transform(img) # img = 2 * (img - 0.5) lab = torch.eye(10)[lab] return img, lab def _sample(self, image, p): for op in p.operations: r = round(random.uniform(0, 1), 1) if r <= op.probability: image = op.perform_operation(image) return image