|
from ..data_aug import pil_image_to_tensor |
|
from ..ab_dataset import ABDataset |
|
from ..dataset_split import train_val_split |
|
from ..dataset_cache import get_dataset_cache_path, read_cached_dataset_status, cache_dataset_status |
|
from .mm_image_folder import MMImageFolder |
|
import os |
|
from typing import Dict, List, Optional |
|
from torchvision.transforms import Compose |
|
|
|
from ..registery import dataset_register |
|
|
|
|
|
with open(os.path.join(os.path.dirname(__file__), 'imagenet_classes.txt'), 'r') as f: |
|
classes = [line.split(' ')[2].strip() for line in f.readlines()] |
|
assert len(classes) == 1000 |
|
|
|
@dataset_register( |
|
name='ImageNet-MM-128', |
|
classes=classes, |
|
task_type='MM Image Classification', |
|
object_type='Generic Object', |
|
class_aliases=[], |
|
shift_type=None |
|
) |
|
class ImageNet128(ABDataset): |
|
def create_dataset(self, root_dir: str, split: str, transform: Optional[Compose], |
|
classes: List[str], ignore_classes: List[str], idx_map: Optional[Dict[int, int]]): |
|
if transform is None: |
|
transform = pil_image_to_tensor(128) |
|
self.transform = transform |
|
root_dir = os.path.join(root_dir, 'train' if split != 'test' else 'val') |
|
dataset = MMImageFolder(root_dir, [c for c in classes if c not in ignore_classes], transform=transform) |
|
|
|
cache_file_path = get_dataset_cache_path(root_dir, classes, ignore_classes, idx_map) |
|
if os.path.exists(cache_file_path): |
|
dataset.samples = read_cached_dataset_status(cache_file_path, 'ImageNet-' + split) |
|
else: |
|
if len(ignore_classes) > 0: |
|
ignore_classes_idx = [classes.index(c) for c in ignore_classes] |
|
dataset.samples = [s for s in dataset.samples if s[1] not in ignore_classes_idx] |
|
if idx_map is not None: |
|
dataset.samples = [(s[0], idx_map[s[1]]) if s[1] in idx_map.keys() else s for s in dataset.samples] |
|
|
|
cache_dataset_status(dataset.samples, cache_file_path, 'ImageNet-' + split) |
|
|
|
if split != 'test': |
|
dataset = train_val_split(dataset, split) |
|
return dataset |
|
|
