Function get_num_classes

ML/src/python/neuralforge/data/datasets.py:304–321

(dataset_name)

        raise ValueError(f"Unknown dataset: {name}")

def get_num_classes(dataset_name):
    dataset_name = dataset_name.lower()
    if dataset_name in ['cifar10', 'mnist', 'fashion_mnist', 'fashionmnist', 'stl10']:
        return 10
    elif dataset_name == 'cifar100':
        return 100
    elif dataset_name in ['tiny_imagenet', 'tinyimagenet']:
        return 200
    elif dataset_name == 'imagenet':
        return 1000
    elif dataset_name == 'food101':
        return 101
    elif dataset_name == 'caltech256':
        return 257
    elif dataset_name in ['oxford_pets', 'oxfordpets']:
        return 37
    else:
        return 10


def get_class_names(dataset_name):