Add Custom Dataset Training Support

README.md

@@ -102,6 +102,33 @@ where the currently available models are:
 - [DFKDE](anomalib/models/dfkde)
 - [GANomaly](anomalib/models/ganomaly)
 
+### Custom Dataset
+It is also possible to train on a custom dataset. To do so, `data` section in `config.yaml` is to be modified as follows:
+```yaml
+dataset:
+  name: custom
+  path: <path/to/custom/dataset>
+  normal: normal # name of the folder containing normal images.
+  abnormal: abnormal # name of the folder containing abnormal images.
+  task: segmentation # classification or segmentation
+  mask: <path/to/mask/annotations> #optional
+  extensions: null
+  split_ratio: 0.2
+  seed: 0
+  image_size: 256
+  train_batch_size: 32
+  test_batch_size: 32
+  num_workers: 8
+  transform_config: null
+  create_validation_set: true
+  tiling:
+    apply: false
+    tile_size: null
+    stride: null
+    remove_border_count: 0
+    use_random_tiling: False
+    random_tile_count: 16
+```
 ## Inference
 
 Anomalib contains several tools that can be used to perform inference with a trained model. The script in [`tools/inference`](tools/inference.py) contains an example of how the inference tools can be used to generate a prediction for an input image.

anomalib/config/config.py

@@ -177,7 +177,8 @@ def get_configurable_parameters(
     config = update_input_size_config(config)
 
     # Project Configs
-    project_path = Path(config.project.path) / config.model.name / config.dataset.name / config.dataset.category
+    category = config.dataset.category if "category" in config.dataset.keys() else ""
+    project_path = Path(config.project.path) / config.model.name / config.dataset.name / category
     (project_path / "weights").mkdir(parents=True, exist_ok=True)
     (project_path / "images").mkdir(parents=True, exist_ok=True)
     config.project.path = str(project_path)

anomalib/data/__init__.py

@@ -20,6 +20,7 @@
 from pytorch_lightning import LightningDataModule
 
 from .btech import BTechDataModule
+from .custom import CustomDataModule
 from .inference import InferenceDataset
 from .mvtec import MVTecDataModule
 
@@ -51,12 +52,29 @@ def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule
             # TODO: Remove config values. IAAALD-211
             root=config.dataset.path,
             category=config.dataset.category,
-            image_size=(config.dataset.image_size[0], config.dataset.image_size[0]),
+            image_size=(config.dataset.image_size[0], config.dataset.image_size[1]),
             train_batch_size=config.dataset.train_batch_size,
             test_batch_size=config.dataset.test_batch_size,
             num_workers=config.dataset.num_workers,
             seed=config.project.seed,
         )
+    elif config.dataset.name.lower() == "custom":
+        datamodule = CustomDataModule(
+            root=config.dataset.path,
+            normal=config.dataset.normal,
+            abnormal=config.dataset.abnormal,
+            task=config.dataset.task,
+            mask_dir=config.dataset.mask,
+            extensions=config.dataset.extensions,
+            split_ratio=config.dataset.split_ratio,
+            seed=config.dataset.seed,
+            image_size=(config.dataset.image_size[0], config.dataset.image_size[1]),
+            train_batch_size=config.dataset.train_batch_size,
+            test_batch_size=config.dataset.test_batch_size,
+            num_workers=config.dataset.num_workers,
+            transform_config=config.dataset.transform_config,
+            create_validation_set=config.dataset.create_validation_set,
+        )
     else:
         raise ValueError(
             "Unknown dataset! \n"