from .mydataset import *

# Load the dataset
T2I_DATASET_REGISTRY = {
    'geneval': GenEvalDataset,
    'dpg': DPGDataset
}

def build_dataset(args):
    if args.task_type == 't2i':
        dataset = T2I_DATASET_REGISTRY[args.dataset_type](args)
    else:
        raise NotImplementedError(f"Dataset type {args.dataset_type} not implemented")
    
    return dataset
        