Generator/__init__.py

"""
Datasets interface.
"""
from .constants import dataset_setups
from .datasets import BaseGen, BrainIDGen



dataset_options = {  
    'default': BaseGen,
    'brain_id': BrainIDGen,
}




def build_datasets(gen_args, device):
    """Helper function to build dataset for different splits ('train' or 'test')."""
    datasets = {'all': dataset_options[gen_args.dataset_option](gen_args, device)}
    return datasets