| import torchvision.transforms as tvs_trans | |
| from openood.utils.config import Config | |
| from .transform import Convert, interpolation_modes, normalization_dict | |
| class CiderPreprocessor(): | |
| def __init__(self, config: Config): | |
| self.pre_size = config.dataset.pre_size | |
| self.image_size = config.dataset.image_size | |
| self.interpolation = interpolation_modes[config.dataset.interpolation] | |
| normalization_type = config.dataset.normalization_type | |
| if normalization_type in normalization_dict.keys(): | |
| self.mean = normalization_dict[normalization_type][0] | |
| self.std = normalization_dict[normalization_type][1] | |
| else: | |
| self.mean = [0.5, 0.5, 0.5] | |
| self.std = [0.5, 0.5, 0.5] | |
| if 'imagenet' in config.dataset.name: | |
| self.transform = tvs_trans.Compose([ | |
| tvs_trans.RandomResizedCrop(size=self.image_size, | |
| scale=(0.4, 1.), | |
| interpolation=self.interpolation), | |
| tvs_trans.RandomHorizontalFlip(), | |
| tvs_trans.RandomApply( | |
| [tvs_trans.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8), | |
| tvs_trans.RandomGrayscale(p=0.2), | |
| tvs_trans.ToTensor(), | |
| tvs_trans.Normalize(mean=self.mean, std=self.std), | |
| ]) | |
| else: | |
| self.transform = tvs_trans.Compose([ | |
| Convert('RGB'), | |
| tvs_trans.RandomResizedCrop(size=self.image_size, | |
| scale=(0.2, 1.), | |
| interpolation=self.interpolation), | |
| tvs_trans.RandomHorizontalFlip(), | |
| tvs_trans.RandomApply( | |
| [tvs_trans.ColorJitter(0.4, 0.4, 0.4, 0.1)], p=0.8), | |
| tvs_trans.RandomGrayscale(p=0.2), | |
| tvs_trans.ToTensor(), | |
| tvs_trans.Normalize(mean=self.mean, std=self.std), | |
| ]) | |
| self.transform = TwoCropTransform(self.transform) | |
| def setup(self, **kwargs): | |
| pass | |
| def __call__(self, image): | |
| return self.transform(image) | |
| class TwoCropTransform: | |
| """Create two crops of the same image.""" | |
| def __init__(self, transform): | |
| self.transform = transform | |
| def __call__(self, x): | |
| return [self.transform(x), self.transform(x)] | |