New datamodules design

anomalib/data/base/datamodule.py

@@ -27,14 +27,23 @@ class AnomalibDataModule(LightningDataModule, ABC):
         train_batch_size (int): Batch size used by the train dataloader.
         test_batch_size (int): Batch size used by the val and test dataloaders.
         num_workers (int): Number of workers used by the train, val and test dataloaders.
+        seed (Optional[int], optional): Seed used during random subset splitting.
     """
 
-    def __init__(self, train_batch_size: int, eval_batch_size: int, num_workers: int, val_split_mode: ValSplitMode):
+    def __init__(
+        self,
+        train_batch_size: int,
+        eval_batch_size: int,
+        num_workers: int,
+        val_split_mode: ValSplitMode,
+        seed: Optional[int] = None,
+    ):
         super().__init__()
         self.train_batch_size = train_batch_size
         self.eval_batch_size = eval_batch_size
         self.num_workers = num_workers
         self.val_split_mode = val_split_mode
+        self.seed = seed
 
         self.train_data: Optional[AnomalibDataset] = None
         self.val_data: Optional[AnomalibDataset] = None
@@ -55,26 +64,36 @@ def setup(self, stage: Optional[str] = None):
     def _setup(self, _stage: Optional[str] = None) -> None:
         """Set up the datasets and perform dynamic subset splitting.
 
-        May be overridden in subclass for custom splitting behaviour.
+        This method yay be overridden in subclass for custom splitting behaviour.
+
+        Note: The stage argument is not used here. This is because, for a given instance of an AnomalibDataModule
+        subclass, all three subsets are created at the first call of setup(). This is to accommodate the subset
+        splitting behaviour of anomaly tasks, where the validation set is usually extracted from the test set, and
+        the test set must therefore be created as early as the `fit` stage.
         """
         assert self.train_data is not None
         assert self.test_data is not None
 
         self.train_data.setup()
         self.test_data.setup()
         if self.val_split_mode == ValSplitMode.FROM_TEST:
-            self.val_data, self.test_data = random_split(self.test_data, [0.5, 0.5], label_aware=True)
+            self.val_data, self.test_data = random_split(self.test_data, [0.5, 0.5], label_aware=True, seed=self.seed)
         elif self.val_split_mode == ValSplitMode.SAME_AS_TEST:
             self.val_data = self.test_data
-        else:
+        elif self.val_split_mode != ValSplitMode.NONE:
             raise ValueError(f"Unknown validation split mode: {self.val_split_mode}")
 
     @property
     def is_setup(self):
         """Checks if setup() has been called."""
-        if self.train_data is None or self.val_data is None or self.test_data is None:
-            return False
-        return self.train_data.is_setup and self.val_data.is_setup and self.test_data.is_setup
+        # at least one of [train_data, val_data, test_data] should be setup
+        if self.train_data is not None and self.train_data.is_setup:
+            return True
+        if self.val_data is not None and self.val_data.is_setup:
+            return True
+        if self.test_data is not None and self.test_data.is_setup:
+            return True
+        return False
 
     def train_dataloader(self) -> TRAIN_DATALOADERS:
         """Get train dataloader."""

anomalib/data/btech.py

@@ -124,7 +124,6 @@ def __init__(
             pre_process: List of pre_processing object containing albumentation compose.
             split: 'train', 'val' or 'test'
             task: ``classification`` or ``segmentation``
-            seed: seed used for the random subset splitting
             create_validation_set: Create a validation subset in addition to the train and test subsets
 
         Examples:
@@ -182,6 +181,7 @@ def __init__(
         transform_config_train: Optional[Union[str, A.Compose]] = None,
         transform_config_eval: Optional[Union[str, A.Compose]] = None,
         val_split_mode: ValSplitMode = ValSplitMode.SAME_AS_TEST,
+        seed: Optional[int] = None,
     ) -> None:
         """Instantiate BTech Lightning Data Module.
 
@@ -195,8 +195,8 @@ def __init__(
             task: ``classification`` or ``segmentation``
             transform_config_train: Config for pre-processing during training.
             transform_config_val: Config for pre-processing during validation.
-            seed: seed used for the random subset splitting
             create_validation_set: Create a validation subset in addition to the train and test subsets
+            seed (Optional[int], optional): Seed used during random subset splitting.
 
         Examples:
             >>> from anomalib.data import BTech
@@ -224,7 +224,7 @@ def __init__(
             >>> data["image"].shape, data["mask"].shape
             (torch.Size([32, 3, 256, 256]), torch.Size([32, 256, 256]))
         """
-        super().__init__(train_batch_size, eval_batch_size, num_workers, val_split_mode)
+        super().__init__(train_batch_size, eval_batch_size, num_workers, val_split_mode, seed)
 
         self.root = Path(root)
         self.category = Path(category)

anomalib/data/folder.py

@@ -231,6 +231,7 @@ class Folder(AnomalibDataModule):
             during validation.
             Defaults to None.
         val_split_mode (ValSplitMode): Setting that determines how the validation subset is obtained.
+        seed (Optional[int], optional): Seed used during random subset splitting.
     """
 
     def __init__(
@@ -251,12 +252,14 @@ def __init__(
         transform_config_train: Optional[Union[str, A.Compose]] = None,
         transform_config_eval: Optional[Union[str, A.Compose]] = None,
         val_split_mode: ValSplitMode = ValSplitMode.FROM_TEST,
+        seed: Optional[int] = None,
     ):
         super().__init__(
             train_batch_size=train_batch_size,
             eval_batch_size=eval_batch_size,
             num_workers=num_workers,
             val_split_mode=val_split_mode,
+            seed=seed,
         )
 
         self.split_ratio = split_ratio
@@ -298,7 +301,7 @@ def _setup(self, _stage: Optional[str] = None):
 
         # add some normal images to the test set
         if not self.test_data.has_normal:
-            self.train_data, normal_test_data = random_split(self.train_data, self.split_ratio)
+            self.train_data, normal_test_data = random_split(self.train_data, self.split_ratio, seed=self.seed)
             self.test_data += normal_test_data
 
         super()._setup()

anomalib/data/mvtec.py

@@ -162,12 +162,14 @@ def __init__(
         transform_config_train: Optional[Union[str, A.Compose]] = None,
         transform_config_eval: Optional[Union[str, A.Compose]] = None,
         val_split_mode: ValSplitMode = ValSplitMode.SAME_AS_TEST,
+        seed: Optional[int] = None,
     ):
         super().__init__(
             train_batch_size=train_batch_size,
             eval_batch_size=eval_batch_size,
             num_workers=num_workers,
             val_split_mode=val_split_mode,
+            seed=seed,
         )
 
         self.root = Path(root)

anomalib/data/utils/split.py

@@ -35,6 +35,7 @@ class Split(str, Enum):
 class ValSplitMode(str, Enum):
     """Splitting mode used to obtain validation subset."""
 
+    NONE = "none"
     SAME_AS_TEST = "same_as_test"
     FROM_TEST = "from_test"
 

tools/test.py

@@ -5,7 +5,7 @@
 
 from argparse import ArgumentParser, Namespace
 
-from pytorch_lightning import Trainer
+from pytorch_lightning import Trainer, seed_everything
 
 from anomalib.config import get_configurable_parameters
 from anomalib.data import get_datamodule
@@ -40,6 +40,9 @@ def test():
         weight_file=args.weight_file,
     )
 
+    if config.project.seed:
+        seed_everything(config.project.seed)
+
     datamodule = get_datamodule(config)
     model = get_model(config)