124 lines
3.2 KiB
Python
124 lines
3.2 KiB
Python
import os.path as osp
|
|
|
|
from dassl.utils import listdir_nohidden
|
|
|
|
from ..build import DATASET_REGISTRY
|
|
from ..base_dataset import Datum, DatasetBase
|
|
|
|
AVAI_C_TYPES = [
|
|
"brightness",
|
|
"contrast",
|
|
"defocus_blur",
|
|
"elastic_transform",
|
|
"fog",
|
|
"frost",
|
|
"gaussian_blur",
|
|
"gaussian_noise",
|
|
"glass_blur",
|
|
"impulse_noise",
|
|
"jpeg_compression",
|
|
"motion_blur",
|
|
"pixelate",
|
|
"saturate",
|
|
"shot_noise",
|
|
"snow",
|
|
"spatter",
|
|
"speckle_noise",
|
|
"zoom_blur",
|
|
]
|
|
|
|
|
|
@DATASET_REGISTRY.register()
|
|
class CIFAR10C(DatasetBase):
|
|
"""CIFAR-10 -> CIFAR-10-C.
|
|
|
|
Dataset link: https://zenodo.org/record/2535967#.YFwtV2Qzb0o
|
|
|
|
Statistics:
|
|
- 2 domains: the normal CIFAR-10 vs. a corrupted CIFAR-10
|
|
- 10 categories
|
|
|
|
Reference:
|
|
- Hendrycks et al. Benchmarking neural network robustness
|
|
to common corruptions and perturbations. ICLR 2019.
|
|
"""
|
|
|
|
dataset_dir = ""
|
|
domains = ["cifar10", "cifar10_c"]
|
|
|
|
def __init__(self, cfg):
|
|
root = osp.abspath(osp.expanduser(cfg.DATASET.ROOT))
|
|
self.dataset_dir = root
|
|
|
|
self.check_input_domains(
|
|
cfg.DATASET.SOURCE_DOMAINS, cfg.DATASET.TARGET_DOMAINS
|
|
)
|
|
source_domain = cfg.DATASET.SOURCE_DOMAINS[0]
|
|
target_domain = cfg.DATASET.TARGET_DOMAINS[0]
|
|
assert source_domain == self.domains[0]
|
|
assert target_domain == self.domains[1]
|
|
|
|
c_type = cfg.DATASET.CIFAR_C_TYPE
|
|
c_level = cfg.DATASET.CIFAR_C_LEVEL
|
|
|
|
if not c_type:
|
|
raise ValueError(
|
|
"Please specify DATASET.CIFAR_C_TYPE in the config file"
|
|
)
|
|
|
|
assert (
|
|
c_type in AVAI_C_TYPES
|
|
), f'C_TYPE is expected to belong to {AVAI_C_TYPES}, but got "{c_type}"'
|
|
assert 1 <= c_level <= 5
|
|
|
|
train_dir = osp.join(self.dataset_dir, source_domain, "train")
|
|
test_dir = osp.join(
|
|
self.dataset_dir, target_domain, c_type, str(c_level)
|
|
)
|
|
|
|
if not osp.exists(test_dir):
|
|
raise ValueError
|
|
|
|
train = self._read_data(train_dir)
|
|
test = self._read_data(test_dir)
|
|
|
|
super().__init__(train_x=train, test=test)
|
|
|
|
def _read_data(self, data_dir):
|
|
class_names = listdir_nohidden(data_dir)
|
|
class_names.sort()
|
|
items = []
|
|
|
|
for label, class_name in enumerate(class_names):
|
|
class_dir = osp.join(data_dir, class_name)
|
|
imnames = listdir_nohidden(class_dir)
|
|
|
|
for imname in imnames:
|
|
impath = osp.join(class_dir, imname)
|
|
item = Datum(impath=impath, label=label, domain=0)
|
|
items.append(item)
|
|
|
|
return items
|
|
|
|
|
|
@DATASET_REGISTRY.register()
|
|
class CIFAR100C(CIFAR10C):
|
|
"""CIFAR-100 -> CIFAR-100-C.
|
|
|
|
Dataset link: https://zenodo.org/record/3555552#.YFxpQmQzb0o
|
|
|
|
Statistics:
|
|
- 2 domains: the normal CIFAR-100 vs. a corrupted CIFAR-100
|
|
- 10 categories
|
|
|
|
Reference:
|
|
- Hendrycks et al. Benchmarking neural network robustness
|
|
to common corruptions and perturbations. ICLR 2019.
|
|
"""
|
|
|
|
dataset_dir = ""
|
|
domains = ["cifar100", "cifar100_c"]
|
|
|
|
def __init__(self, cfg):
|
|
super().__init__(cfg)
|