import pickle
import glob


def get_num_classes(data_dir):
    """
    Get the number of classes.
    :param data_dir: str - data directory
    :return: int - number of classes
    """
    mode = 'training' # arbitrary
    loc = "{}/{}".format(data_dir, mode)
    with open('{}/labels.pickle'.format(data_dir), 'rb') as f:
        data = pickle.load(f)

    modes = list(data.keys())

    assert glob.glob(data_dir), "Check directory."
    assert glob.glob("{}/*.jpg".format(loc)), "Check file extension (should be 'jpg')."
    i = 0  # Arbitrarily chosen
    return len(data[modes[i]][i][-1])