Feature/data/btad

.github/workflows/nightly.yml

@@ -20,7 +20,8 @@ jobs:
         run: pip install tox
       - name: Coverage
         run: |
-          export ANOMALIB_DATASET_PATH=/media/data1/datasets/MVTec
+          export ANOMALIB_DATASET_PATH=/media/data1/datasets/
+          export CUDA_VISIBLE_DEVICES=2
           tox -e nightly
       - name: Upload coverage result
         uses: actions/upload-artifact@v2

.github/workflows/pre_merge.yml

@@ -22,7 +22,7 @@ jobs:
         run: tox -e black,isort,flake8,pylint,mypy,pydocstyle
       - name: Coverage
         run: |
-          export ANOMALIB_DATASET_PATH=/media/data1/datasets/MVTec
+          export ANOMALIB_DATASET_PATH=/media/data1/datasets/
           export CUDA_VISIBLE_DEVICES=3
           tox -e pre_merge
       - name: Upload coverage result

.gitignore

@@ -1,7 +1,7 @@
 # Project related
 datasets
 !anomalib/datasets
-!tests/datasets
+!tests/pre_merge/datasets
 results
 !anomalib/core/results
 

README.md

@@ -153,6 +153,7 @@ python tools/inference.py \
 ___
 
 ## Datasets
+The `development` branch supports MVTec and BeanTech datasets.
 
 ### [MVTec Dataset](https://www.mvtec.com/company/research/datasets/mvtec-ad)
 

anomalib/data/__init__.py

@@ -19,6 +19,7 @@
 from omegaconf import DictConfig, ListConfig
 from pytorch_lightning import LightningDataModule
 
+from .btech import BTechDataModule
 from .inference import InferenceDataset
 from .mvtec import MVTecDataModule
 
@@ -34,7 +35,7 @@ def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule
     """
     datamodule: LightningDataModule
 
-    if config.dataset.format.lower() == "mvtec":
+    if config.dataset.name.lower() == "mvtec":
         datamodule = MVTecDataModule(
             # TODO: Remove config values. IAAALD-211
             root=config.dataset.path,
@@ -45,8 +46,23 @@ def get_datamodule(config: Union[DictConfig, ListConfig]) -> LightningDataModule
             num_workers=config.dataset.num_workers,
             seed=config.project.seed,
         )
+    elif config.dataset.name.lower() == "btech":
+        datamodule = BTechDataModule(
+            # 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]),
+            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,
+        )
     else:
-        raise ValueError("Unknown dataset!")
+        raise ValueError(
+            "Unknown dataset! \n"
+            "If you use a custom dataset make sure you initialize it in"
+            "`get_datamodule` in `anomalib.data.__init__.py"
+        )
 
     return datamodule