| 77 | |
| 78 | |
| 79 | def build_transform(is_train, config): |
| 80 | resize_im = config.DATA.IMG_SIZE > 32 |
| 81 | if is_train: |
| 82 | |
| 83 | transform = create_transform( |
| 84 | input_size=config.DATA.IMG_SIZE, |
| 85 | is_training=True, |
| 86 | color_jitter=config.AUG.COLOR_JITTER if config.AUG.COLOR_JITTER > 0 else None, |
| 87 | auto_augment=config.AUG.AUTO_AUGMENT if config.AUG.AUTO_AUGMENT != 'none' else None, |
| 88 | re_prob=config.AUG.REPROB, |
| 89 | re_mode=config.AUG.REMODE, |
| 90 | re_count=config.AUG.RECOUNT, |
| 91 | interpolation=config.DATA.INTERPOLATION, |
| 92 | ) |
| 93 | if not resize_im: |
| 94 | # replace RandomResizedCropAndInterpolation with |
| 95 | # RandomCrop |
| 96 | transform.transforms[0] = transforms.RandomCrop(config.DATA.IMG_SIZE, padding=4) |
| 97 | return transform |
| 98 | |
| 99 | t = [] |
| 100 | if resize_im: |
| 101 | if config.TEST.CROP: |
| 102 | size = int((256 / 224) * config.DATA.IMG_SIZE) |
| 103 | t.append( |
| 104 | transforms.Resize((size, size), interpolation=transforms.InterpolationMode.BICUBIC), |
| 105 | # to maintain same ratio w.r.t. 224 images |
| 106 | ) |
| 107 | t.append(transforms.CenterCrop(config.DATA.IMG_SIZE)) |
| 108 | else: |
| 109 | t.append( |
| 110 | transforms.Resize((config.DATA.IMG_SIZE, config.DATA.IMG_SIZE), |
| 111 | interpolation=transforms.InterpolationMode.BICUBIC) |
| 112 | ) |
| 113 | |
| 114 | t.append(transforms.ToTensor()) |
| 115 | t.append(transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)) |
| 116 | return transforms.Compose(t) |