Spaces:
Runtime error
Runtime error
| """Image preprocessing transforms shared by trainers and inference. | |
| Three callables are exported: | |
| * :data:`IMAGENET_MEAN` / :data:`IMAGENET_STD` — normalisation statistics that | |
| match the ImageNet-pretrained ResNet50 weights we fine-tune from. | |
| * :func:`eval_transform` — deterministic resize+centre-crop+normalise. Used | |
| for validation, test, feature extraction, and inference. | |
| * :func:`train_transform` — augmentation pipeline with optional RandAugment. | |
| Both transforms accept ``image_size`` (typically 224 for ResNet50). The train | |
| transform optionally takes RandAugment parameters; setting ``num_ops=0`` falls | |
| back to a simpler ColorJitter for backwards compatibility. | |
| """ | |
| from __future__ import annotations | |
| from typing import Callable | |
| from torchvision import transforms | |
| IMAGENET_MEAN: tuple[float, float, float] = (0.485, 0.456, 0.406) | |
| IMAGENET_STD: tuple[float, float, float] = (0.229, 0.224, 0.225) | |
| def eval_transform(image_size: int = 224) -> Callable: | |
| """Return the deterministic preprocessing pipeline. | |
| Resize to 1.15× the target size, centre-crop down to ``image_size``, then | |
| convert to a normalised float tensor. The 1.15× margin matches what we use | |
| during training so train/eval distributions agree. | |
| """ | |
| return transforms.Compose([ | |
| transforms.Resize(int(image_size * 1.15)), | |
| transforms.CenterCrop(image_size), | |
| transforms.ToTensor(), | |
| transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), | |
| ]) | |
| def train_transform( | |
| image_size: int = 224, | |
| randaug_num_ops: int = 0, | |
| randaug_magnitude: int = 9, | |
| ) -> Callable: | |
| """Return a training-time augmentation pipeline. | |
| With ``randaug_num_ops > 0`` the pipeline uses :class:`torchvision.transforms.RandAugment` | |
| with the given operation count and magnitude (typical values: ``num_ops=2, | |
| magnitude=9``). With ``num_ops=0`` it falls back to a hand-tuned ColorJitter | |
| that matches our pre-RandAugment baseline so old behaviour is recoverable. | |
| """ | |
| pipeline = [ | |
| transforms.Resize(int(image_size * 1.15)), | |
| transforms.RandomResizedCrop(image_size, scale=(0.7, 1.0)), | |
| transforms.RandomHorizontalFlip(), | |
| ] | |
| if randaug_num_ops > 0: | |
| pipeline.append(transforms.RandAugment( | |
| num_ops=randaug_num_ops, magnitude=randaug_magnitude, | |
| )) | |
| else: | |
| pipeline.append(transforms.ColorJitter( | |
| brightness=0.2, contrast=0.2, saturation=0.15, | |
| )) | |
| pipeline += [ | |
| transforms.ToTensor(), | |
| transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD), | |
| ] | |
| return transforms.Compose(pipeline) | |