# AUTOGENERATED! DO NOT EDIT! File to edit: dataset-type.ipynb (unless otherwise specified).

__all__ = ['DatasetType', 'infer_dataset_type']


# Cell
from enum import Enum
from os.path import dirname, basename


# Cell
class DatasetType(Enum):
    """
    Currently supported Machine Learning Types.
    """
    IMAGE_CLASSIFICATION = 'image_classification'
    IMAGE_SEGMENTATION = 'image_segmentation'
    IMAGE_GENERATION = 'image_generation'
    IMAGE_OBJECT_DETECTION = 'image_object_detection'

    def __str__(self):
        return self.value


# Cell
def infer_dataset_type(category_file_path):
    """
    Try to infer the dataset type from the category file path.
    `category_file_path`: the path of the categories file
    return: the infered dataset type
    raises: `ValueError` if unsupported type infered
    """
    dataset_type_from_path = None
    try:
        path = dirname(category_file_path)
        dataset_type_from_path = basename(dirname(path))
        return DatasetType(dataset_type_from_path)
    except ValueError:
        raise ValueError("Error, unsupported dataset type: {}".format(dataset_type_from_path))
