def _setup(self, _stage: str | None = None) -> None:
        """Set up the datasets and perform train/validation/test split.

        Args:
            _stage: str | None: Optional argument for compatibility with pytorch
                lightning. Defaults to None.
        """
        self.train_data = MVTecAD2Dataset(
            root=self.root,
            category=self.category,
            split=Split.TRAIN,
            augmentations=self.train_augmentations,
        )

        # MVTec AD 2 has a dedicated validation set
        self.val_data = MVTecAD2Dataset(
            root=self.root,
            category=self.category,
            split=Split.VAL,
            augmentations=self.val_augmentations,
        )

        # Create datasets for all test types
        self.test_public_data = MVTecAD2Dataset(
            root=self.root,
            category=self.category,
            split=Split.TEST,
            test_type=TestType.PUBLIC,
            augmentations=self.test_augmentations,
        )

        self.test_private_data = MVTecAD2Dataset(
            root=self.root,
            category=self.category,
            split=Split.TEST,
            test_type=TestType.PRIVATE,
            augmentations=self.test_augmentations,
        )

        self.test_private_mixed_data = MVTecAD2Dataset(
            root=self.root,
            category=self.category,
            split=Split.TEST,
            test_type=TestType.PRIVATE_MIXED,
            augmentations=self.test_augmentations,
        )

        # Always set test_data to public test set for standard evaluation
        self.test_data = self.test_public_data
Set up the datasets and perform train/validation/test split.

        Args:
            _stage: str | None: Optional argument for compatibility with pytorch
                lightning. Defaults to None.
        
src/anomalib/data/datamodules/image/mvtecad2.py