config_load.py

"""
Temporary file for testing loading from config without modifying armory.utils
"""

from armory.datasets import load, preprocessing, generator, filtering


def load_dataset(
    name=None,
    version=None,
    config=None,
    batch_size=1,
    num_batches=None,
    epochs=1,
    split="test",
    framework="numpy",
    preprocessor_name="DEFAULT",
    preprocessor_kwargs=None,
    shuffle_files=False,
    label_key="label",  # TODO: make this smarter or more flexible
    index=None,
    class_ids=None,
    drop_remainder=False,
):
    # All are keyword elements by design
    if name is None:
        raise ValueError("name must be specified, not None")
    info, ds_dict = load.load(
        name, version=version, config=config, shuffle_files=shuffle_files
    )

    if class_ids is None:
        element_filter = None
    else:
        if isinstance(class_ids, int):
            class_ids = [class_ids]
        if not isinstance(class_ids, list):
            raise ValueError(
                f"class_ids must be a list, int, or None, not {type(class_ids)}"
            )
        element_filter = filtering.get_filter_by_class(class_ids, label_key=label_key)
    if index is None:
        index_filter = None
    elif isinstance(index, list):
        index_filter = filtering.get_enum_filter_by_index(index)
    elif isinstance(index, str):
        index_filter = filtering.get_enum_filter_by_slice(index)
    else:
        raise ValueError(f"index must be a list, str, or None, not {type(index)}")

    if preprocessor_name is None:
        preprocessor = None
    elif preprocessor_name == "DEFAULT":
        if preprocessing.has(name):
            preprocessor = preprocessing.get(name)
        else:
            preprocessor = preprocessing.infer_from_dataset_info(info, split)
    else:
        preprocessor = preprocessing.get(preprocessor_name)

    if preprocessor is not None and preprocessor_kwargs is not None:
        preprocessing_fn = lambda x: preprocessor(  # noqa: E731
            x, **preprocessor_kwargs
        )
    else:
        preprocessing_fn = preprocessor

    shuffle_elements = shuffle_files

    return generator.ArmoryDataGenerator(
        info,
        ds_dict,
        split=split,
        batch_size=batch_size,
        framework=framework,
        epochs=epochs,
        drop_remainder=drop_remainder,
        num_batches=num_batches,
        index_filter=index_filter,
        element_filter=element_filter,
        element_map=preprocessing_fn,
        shuffle_elements=shuffle_elements,
        key_map=None,
    )