Spaces:
Running
Running
| from typing import Dict, Type | |
| from .ab_dataset import ABDataset | |
| static_dataset_registery = {} | |
| # (NIPS'20) Measuring Robustness to Natural Distribution Shifts in Image Classification | |
| POSSIBLE_SHIFT_TYPES = ['Consistency Shifts', 'Dataset Shifts', 'Adversarially Filtered Shifts', | |
| 'Image Corruptions', 'Geometric Transformations', 'Style Transfer', 'Adversarial Examples'] | |
| def dataset_register(name, classes, task_type, object_type, class_aliases=[], shift_type: Dict[str, str]=None): | |
| assert shift_type is None or len(shift_type.keys()) == 1 | |
| if shift_type is not None: | |
| assert list(shift_type.values())[0] in POSSIBLE_SHIFT_TYPES | |
| class _Register: | |
| def __init__(self, func: Type[ABDataset]): | |
| self.func = func | |
| static_dataset_registery[name] = (self, classes, task_type, object_type, class_aliases, shift_type) | |
| def __call__(self, *args, **kwargs): | |
| res = self.func(*args, **kwargs) | |
| res.name = name | |
| res.classes = classes | |
| res.class_aliases = class_aliases | |
| res.shift_type = shift_type | |
| res.task_type = task_type | |
| res.object_type = object_type | |
| res.build() | |
| return res | |
| return _Register | |