third_party/SOLD2/sold2/dataset/dataset_util.py

"""
The interface of initializing different datasets.
"""
from .synthetic_dataset import SyntheticShapes
from .wireframe_dataset import WireframeDataset
from .holicity_dataset import HolicityDataset
from .merge_dataset import MergeDataset


def get_dataset(mode="train", dataset_cfg=None):
    """Initialize different dataset based on a configuration."""
    # Check dataset config is given
    if dataset_cfg is None:
        raise ValueError("[Error] The dataset config is required!")

    # Synthetic dataset
    if dataset_cfg["dataset_name"] == "synthetic_shape":
        dataset = SyntheticShapes(mode, dataset_cfg)

        # Get the collate_fn
        from .synthetic_dataset import synthetic_collate_fn

        collate_fn = synthetic_collate_fn

    # Wireframe dataset
    elif dataset_cfg["dataset_name"] == "wireframe":
        dataset = WireframeDataset(mode, dataset_cfg)

        # Get the collate_fn
        from .wireframe_dataset import wireframe_collate_fn

        collate_fn = wireframe_collate_fn

    # Holicity dataset
    elif dataset_cfg["dataset_name"] == "holicity":
        dataset = HolicityDataset(mode, dataset_cfg)

        # Get the collate_fn
        from .holicity_dataset import holicity_collate_fn

        collate_fn = holicity_collate_fn

    # Dataset merging several datasets in one
    elif dataset_cfg["dataset_name"] == "merge":
        dataset = MergeDataset(mode, dataset_cfg)

        # Get the collate_fn
        from .holicity_dataset import holicity_collate_fn

        collate_fn = holicity_collate_fn

    else:
        raise ValueError(
            "[Error] The dataset '%s' is not supported" % dataset_cfg["dataset_name"]
        )

    return dataset, collate_fn