from .cifar10 import CIFAR10
from .cifar100 import CIFAR100
from .stl10 import STL10
from .tiny_in import TinyImageNet
from .imagenet import ImageNet


DS_LIST = ["cifar10", "cifar100", "stl10", "tinyimagenet", "imagenet"]


def get_ds(name):
    assert name in DS_LIST
    if name == "cifar10":
        return CIFAR10
    elif name == "cifar100":
        return CIFAR100
    elif name == "stl10":
        return STL10
    elif name == "tinyimagenet":
        return TinyImageNet
    elif name == "imagenet":
        return ImageNet