| """Image preprocessing pipeline for Imager-1-Tiny.""" | |
| from torchvision import transforms | |
| IMAGENET_MEAN = [0.485, 0.456, 0.406] | |
| IMAGENET_STD = [0.229, 0.224, 0.225] | |
| def get_transform(train: bool = False): | |
| """ | |
| Return the torchvision transform pipeline. | |
| Args: | |
| train: If True, returns the training pipeline with augmentation. | |
| If False, returns the deterministic inference pipeline. | |
| """ | |
| if train: | |
| return transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.RandomResizedCrop(224, scale=(0.8, 1.0)), | |
| transforms.RandomHorizontalFlip(), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), | |
| ]) | |
| return transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), | |
| ]) | |