(is_train, config)
| 73 | |
| 74 | |
| 75 | def build_dataset(is_train, config): |
| 76 | transform = build_transform(is_train, config) |
| 77 | if config.DATA.DATASET == 'imagenet': |
| 78 | prefix = 'train' if is_train else 'val' |
| 79 | if config.DATA.ZIP_MODE: |
| 80 | ann_file = prefix + "_map.txt" |
| 81 | prefix = prefix + ".zip@/" |
| 82 | dataset = CachedImageFolder(config.DATA.DATA_PATH, ann_file, prefix, transform, |
| 83 | cache_mode=config.DATA.CACHE_MODE if is_train else 'part') |
| 84 | else: |
| 85 | # ToDo: test custom_image_folder |
| 86 | root = os.path.join(config.DATA.DATA_PATH, prefix) |
| 87 | dataset = CustomImageFolder(root, transform=transform) |
| 88 | nb_classes = 1000 |
| 89 | else: |
| 90 | raise NotImplementedError("We only support ImageNet Now.") |
| 91 | |
| 92 | return dataset, nb_classes |
| 93 | |
| 94 | |
| 95 | def build_transform(is_train, config): |
no test coverage detected