(transform_type='default', image_size=32, args=None)
| 3 | from data.augmentations.randaugment import RandAugment |
| 4 | |
| 5 | def get_transform(transform_type='default', image_size=32, args=None): |
| 6 | |
| 7 | if transform_type == 'imagenet': |
| 8 | |
| 9 | mean = (0.485, 0.456, 0.406) |
| 10 | std = (0.229, 0.224, 0.225) |
| 11 | interpolation = args.interpolation |
| 12 | crop_pct = args.crop_pct |
| 13 | |
| 14 | train_transform = transforms.Compose([ |
| 15 | transforms.Resize(int(image_size / crop_pct), interpolation), |
| 16 | transforms.RandomCrop(image_size), |
| 17 | transforms.RandomHorizontalFlip(p=0.5), |
| 18 | transforms.ColorJitter(), |
| 19 | transforms.ToTensor(), |
| 20 | transforms.Normalize( |
| 21 | mean=torch.tensor(mean), |
| 22 | std=torch.tensor(std)) |
| 23 | ]) |
| 24 | |
| 25 | test_transform = transforms.Compose([ |
| 26 | transforms.Resize(int(image_size / crop_pct), interpolation), |
| 27 | transforms.CenterCrop(image_size), |
| 28 | transforms.ToTensor(), |
| 29 | transforms.Normalize( |
| 30 | mean=torch.tensor(mean), |
| 31 | std=torch.tensor(std)) |
| 32 | ]) |
| 33 | |
| 34 | elif transform_type == 'pytorch-cifar': |
| 35 | |
| 36 | mean = (0.4914, 0.4822, 0.4465) |
| 37 | std = (0.2023, 0.1994, 0.2010) |
| 38 | |
| 39 | train_transform = transforms.Compose([ |
| 40 | transforms.RandomCrop(image_size, padding=4), |
| 41 | transforms.RandomHorizontalFlip(), |
| 42 | transforms.ToTensor(), |
| 43 | transforms.Normalize(mean=mean, std=std), |
| 44 | ]) |
| 45 | |
| 46 | test_transform = transforms.Compose([ |
| 47 | transforms.Resize((image_size, image_size)), |
| 48 | transforms.ToTensor(), |
| 49 | transforms.Normalize(mean=mean, std=std), |
| 50 | ]) |
| 51 | |
| 52 | elif transform_type == 'herbarium_default': |
| 53 | |
| 54 | train_transform = transforms.Compose([ |
| 55 | transforms.Resize((image_size, image_size)), |
| 56 | transforms.RandomResizedCrop(image_size, scale=(args.resize_lower_bound, 1)), |
| 57 | transforms.RandomHorizontalFlip(), |
| 58 | transforms.ToTensor(), |
| 59 | ]) |
| 60 | |
| 61 | test_transform = transforms.Compose([ |
| 62 | transforms.Resize((image_size, image_size)), |
no test coverage detected