def setup(self, stage: str | None = None) -> None:
        """Set up train, validation and test data.

        This method handles the data splitting logic based on the configured
        modes.

        Args:
            stage (str | None): Current stage (fit/validate/test/predict).
                Defaults to ``None``.
        """
        has_subset = any(hasattr(self, subset) for subset in ["train_data", "val_data", "test_data"])
        if not has_subset or not self._is_setup:
            self._setup(stage)
            self._create_test_split()
            self._create_val_split()
            if isinstance(stage, TrainerFn):
                # only set flag if called from trainer
                self._is_setup = True

        self._update_augmentations()
Set up train, validation and test data.

        This method handles the data splitting logic based on the configured
        modes.

        Args:
            stage (str | None): Current stage (fit/validate/test/predict).
                Defaults to ``None``.
        
src/anomalib/data/datamodules/base/image.py