def get_datamodule(config: DictConfig | ListConfig | dict) -> AnomalibDataModule:
    """Get Anomaly Datamodule from config.

    Args:
        config: Configuration for the anomaly model. Can be either:
            - DictConfig from OmegaConf
            - ListConfig from OmegaConf
            - Python dictionary

    Returns:
        PyTorch Lightning DataModule configured according to the input.

    Raises:
        UnknownDatamoduleError: If the specified datamodule cannot be found.

    Example:
        >>> from omegaconf import DictConfig
        >>> config = DictConfig({
        ...     "data": {
        ...         "class_path": "MVTecAD",
        ...         "init_args": {"root": "./datasets/MVTec"}
        ...     }
        ... })
        >>> datamodule = get_datamodule(config)
    """
    logger.info("Loading the datamodule and dataset class from the config.")

    # Getting the datamodule class from the config.
    if isinstance(config, dict):
        config = DictConfig(config)
    _config = config.data if "data" in config else config

    # All the sub data modules are imported to anomalib.data. So need to import the module dynamically using paths.
    module = importlib.import_module("anomalib.data")
    data_class_name = _config.class_path.split(".")[-1]
    # check if the data_class exists in the module
    if not hasattr(module, data_class_name):
        logger.error(
            f"Dataclass '{data_class_name}' not found in module '{module.__name__}'. "
            f"Available classes are {AnomalibDataModule.__subclasses__()}",
        )
        error_str = f"Dataclass '{data_class_name}' not found in module '{module.__name__}'."
        raise UnknownDatamoduleError(error_str)
    dataclass = getattr(module, data_class_name)

    init_args = {**_config.get("init_args", {})}  # get dict
    if "image_size" in init_args:
        init_args["image_size"] = to_tuple(init_args["image_size"])
    return dataclass(**init_args)
Get Anomaly Datamodule from config.

    Args:
        config: Configuration for the anomaly model. Can be either:
            - DictConfig from OmegaConf
            - ListConfig from OmegaConf
            - Python dictionary

    Returns:
        PyTorch Lightning DataModule configured according to the input.

    Raises:
        UnknownDatamoduleError: If the specified datamodule cannot be found.

    Example:
        >>> from omegaconf import DictConfig
        >>> config = DictConfig({
        ...     "data": {
        ...         "class_path": "MVTecAD",
        ...         "init_args": {"root": "./datasets/MVTec"}
        ...     }
        ... })
        >>> datamodule = get_datamodule(config)
    
src/anomalib/data/__init__.py