anomalib-adding-a-datamodule

# Adding a New Dataset / DataModule anomalib splits data support into two layers per source, both under `src/anomalib/data/`: - `datasets/image/<name>.py` — a `torch`-facing `AnomalibDataset` subclass (one dataset = one split). - `datamodules/image/<name>.py` — a Lightning-facing `AnomalibDataModule` subclass that owns train/val/test dataloaders and split logic. (Use `datasets/video/` and `datamodules/video/`, or `depth/`, for other modalities — the pattern is identical.) ## Base classes to implement against - `AnomalibDataset` — `src/anomalib/data/datasets/base/image.py` - `__init__(self, augmentations=None)` — call via `super().__init__(...)`. - You must build a `pandas.DataFrame` and assign it to `self.samples`. Required columns: `image_path`, `split`, `label_index` (0 for normal, 1 for abnormal); segmentation datasets also need `mask_path` (set to empty string `""` for normal samples). After building the DataFrame, set `samples.attrs["task"]` to `"classification"` or `"segmentation"`. - `collate_fn` defaults to `ImageBatch.collate`; override only for non-image batch types. - `AnomalibDataModule` — `src/anomalib/data/datamodules/base/image.py` - Only abstract method you must implement: `_setup(self, _stage=None) -> None`, where you set `self.train_data` and `self.test_data` (and `self.val_data` if you don't rely on the base class's `val_split_mode` machinery). - The base class already implements `setup()`, `train_dataloader()`, `val_dataloader()`, `test_dataloader()`, and `from_config()` (jsonargparse subclass integration) — do not override these unless the data source genuinely needs custom dataloader construction. - Constructor should accept and forward: `train_batch_size`, `eval_batch_size`, `num_workers`, `train_augmentations` / `val_augmentations` / `test_augmentations` / `augmentations`, `test_split_mode` / `test_split_ratio`, `val_split_mode` / `val_split_ratio`, `seed`. ## Reference: MVTecAD (standard benchmark-style dataset) - `src/anomalib/data/datasets/image/mvtecad.py` — `MVTecADDataset(AnomalibDataset)`; builds `self.samples` via the `make_mvtec_ad_dataset(root_category, split, extensions)` helper. - `src/anomalib/data/datamodules/image/mvtecad.py` — `MVTecAD(AnomalibDataModule)`; `_setup()` constructs `MVTecADDataset(split=Split.TRAIN, root=self.root, category=self.category)` for train/test, and `prepare_data()` downloads the dataset archive if missing. ## Reference: Folder (generic custom-folder dataset — use this as your template for ad hoc data) - `src/anomalib/data/datasets/image/folder.py` — `FolderDataset(AnomalibDataset)`, built via the `make_folder_dataset(...)` helper: collects filenames/labels from directories, builds the `samples` DataFrame, and attaches mask paths to abnormal samples when `mask_dir` is given. - `src/anomalib/data/datamodules/image/folder.py` — `Folder(AnomalibDataModule)` constructor (key args): ```python Folder( name: str, # required, becomes datamodule.name normal_dir: str | Path | Sequence[str | Path], # required root: str | Path | None = None, abnormal_dir: str | Path | Sequence[str | Path] | None = None, normal_test_dir: str | Path | Sequence[str | Path] | None = None, # separate normal images for test set mask_dir: str | Path | Sequence[str | Path] | None = None, # for segmentation masks normal_split_ratio: float = 0.2, extensions: tuple[str] | None = None, train_batch_size: int = 32, eval_batch_size: int = 32, num_workers: int = 8, train_augmentations: Transform | None = None, val_augmentations: Transform | None = None, test_augmentations: Transform | None = None, augmentations: Transform | None = None, test_split_mode: TestSplitMode = TestSplitMode.FROM_DIR, test_split_ratio: float = 0.2, val_split_mode: ValSplitMode = ValSplitMode.FROM_TEST, val_split_ratio: float = 0.5, seed: int | None = None, ) ``` Note: `test_split_ratio` (inherited from `AnomalibDataModule`) controls the fraction of training images held out for testing when `test_split_mode` triggers a synthetic split. `normal_split_ratio` is stored by `Folder` but used only by `FolderDataset` internally to split normal images between train and test sets when `normal_test_dir` is not provided and test data must come from the normal pool. Use `Folder` directly (no new code needed) whenever the data is already laid out as `root/normal_dir/*`, `root/abnormal_dir/*`, optionally `root/mask_dir/*`. Only write a brand-new dataset/datamodule pair when the data needs custom parsing logic `Folder` can't express. ## Writing a brand-new datamodule (skeleton) ```python # src/anomalib/data/datasets/image/my_dataset.py from anomalib.data.datasets.base import AnomalibDataset class MyDataset(AnomalibDataset): def __init__(self, root=None, augmentations=None, split=None): super().__init__(augmentations=augmentations) samples = make_my_dataset_samples(root=root, split=split) # build the DataFrame yourself # DataFrame must have columns: image_path, split, label_index (and mask_path for segmentation) samples.attrs["task"] = "segmentation" # or "classification" self.samples = samples ``` ```python # src/anomalib/data/datamodules/image/my_dataset.py from pathlib import Path from anomalib.data.datamodules.base.image import AnomalibDataModule from anomalib.data.datasets.image.my_dataset import MyDataset from anomalib.data.utils import Split class MyDataModule(AnomalibDataModule): def __init__( self, root: str | Path = "./datasets/MyDataset", train_batch_size: int = 32, eval_batch_size: int = 32, num_workers: int = 8, train_augmentations=None, val_augmentations=None, test_augmentations=None, augmentations=None, test_split_mode=None, test_split_ratio: float = 0.2, val_split_mode=None, val_split_ratio: float = 0.5, seed: int | None = None, ) -> None: super().__init__( train_batch_size=train_batch_size, eval_batch_size=eval_batch_size, num_workers=num_workers, train_augmentations=train_augmentations, val_augmentations=val_augmentations, test_augmentations=test_augmentations, augmentations=augmentations, test_split_mode=test_split_mode, test_split_ratio=test_split_ratio, val_split_mode=val_split_mode, val_split_ratio=val_split_ratio, seed=seed, ) self.root = Path(root) def _setup(self, _stage=None) -> None: self.train_data = MyDataset(split=Split.TRAIN, root=self.root) self.test_data = MyDataset(split=Split.TEST, root=self.root) def prepare_data(self) -> None: ... # optional: download/validate on rank-zero ``` ## Registration — how the datamodule becomes discoverable 1. Export from the image (or video/depth) package `__init__.py` — `src/anomalib/data/datamodules/image/__init__.py`: add the import and `__all__` entry there first. 2. Then add the import and `__all__` entry in `src/anomalib/data/__init__.py`, alongside the existing `datamodules.image` import block: ```python from .datamodules.image import ( ..., MyDataModule, ) ``` Once exported, it is usable as `anomalib.data.MyDataModule`, and from the CLI: `anomalib train --model Patchcore --data anomalib.data.MyDataModule --data.root ./datasets/mine`. ## Tests Add `tests/unit/data/datamodule/image/test_my_dataset.py` following the pattern in `tests/unit/data/datamodule/image/test_mvtec_ad.py`: a `datamodule` fixture that instantiates the datamodule against a generated dummy dataset, calls `prepare_data()` + `setup()`, then reuses the shared assertions in `tests/unit/data/datamodule/base/image.py` (batch shapes, split non-overlap, etc.). ### Add a dummy dataset generator (required for a new `DataFormat`) Real datasets aren't checked into the repo — tests generate synthetic data on the fly via `tests/helpers/data.py`. If your new datamodule corresponds to a new `DataFormat` value (i.e. it isn't just `Folder` under another name), you must add a matching generator method: 1. Add the format to `ImageDataFormat` (or `VideoDataFormat`) in `src/anomalib/data/datamodules/image/__init__.py` (or `.../video/__init__.py`), e.g. `MY_DATASET = "my_dataset"`. 2. Implement `_generate_dummy_my_dataset_dataset(self) -> None` on `DummyImageDatasetGenerator` (`tests/helpers/data.py`) for image datasets, or `DummyVideoDatasetGenerator` for video datasets — the method name must be `_generate_dummy_{data_format.value}_dataset`; `DummyDatasetGenerator.generate_dataset()` dispatches to it via `getattr`. Build the on-disk layout your datamodule expects using the low-level `DummyImageGenerator` (`tests/helpers/data.py::DummyImageGenerator`) for image datasets, or `DummyVideoGenerator` for video datasets: ```python def _generate_dummy_my_dataset_dataset(self) -> None: """Generate dummy MyDataset dataset in a temporary directory.""" dataset_category = "dummy" # normal train/test images for split in ("train", "test"): path = self.dataset_root / dataset_category / split / self.normal_category num_images = self.num_train if split == "train" else self.num_test for i in range(num_images): image_filename = path / f"{i:03}.png" self.image_generator.generate_image(label=LabelName.NORMAL, image_filename=image_filename) # abnormal test images + masks path = self.dataset_root / dataset_category / "test" / self.abnormal_category mask_path = self.dataset_root / dataset_category / "ground_truth" / self.abnormal_category for i in range(self.num_test): image_filename = path / f"{i:03}.png" mask_filename = mask_path / f"{i:03}_mask.png" self.image_generator.generate_image(LabelName.ABNORMAL, image_filename, mask_filename) ``` See `_generate_dummy_mvtecad_dataset` and `_generate_dummy_folder_dataset` in the same file for the two canonical layouts (category-per-split-per-class vs. flat normal/abnormal/mask dirs) — mirror whichever matches your real dataset's directory structure. 3. The session-scoped `dataset_path` fixture in `tests/conftest.py` dispatches image formats to `DummyImageDatasetGenerator` and video formats to `DummyVideoDatasetGenerator` (skipping `folder`/ `tabular`, which tests construct manually) — you only need to implement the `_generate_dummy_*` method on the appropriate generator class. 4. In your datamodule test, consume the generated data via the shared fixture: ```python @pytest.fixture() def datamodule(dataset_path: Path) -> MyDataModule: dm = MyDataModule(root=dataset_path / "my_dataset") dm.prepare_data() dm.setup() return dm ``` If your datamodule is just a thin wrapper around `Folder` (same on-disk convention, different defaults), you don't need a new `DataFormat`/generator — reuse `_generate_dummy_folder_dataset` and construct your datamodule directly against its output directory. ## Gotchas - `self.samples` must be assigned (not mutated in place before assignment) — the `samples` setter on `AnomalibDataset` validates required columns and paths. - Don't skip `samples.attrs["task"]` — post-processing and metrics branch on `"classification"` vs `"segmentation"`. - Prefer `Folder` over a new dataset class whenever the on-disk layout is a plain normal/abnormal/mask directory split — writing a new class is only needed for non-standard parsing. - Tests never touch real downloaded datasets. If you add a new `DataFormat`, you must also add a `_generate_dummy_<format>_dataset` method on the appropriate generator (`DummyImageDatasetGenerator` for image formats, `DummyVideoDatasetGenerator` for video formats) — otherwise the shared `dataset_path` fixture (`tests/conftest.py`) will raise `NotImplementedError` for that format. ## Path confinement and data security rules When writing a datamodule or dataset that parses metadata (split CSVs, JSON, Parquet, annotations, or directory trees): - **Never trust dataset paths:** Never join path strings from metadata files directly to `root` with `root / path` or `os.path.join`. Always use `resolve_path_under_root(root, path)` or `validate_path(path, base_dir=root)` from `anomalib.data.utils.path`. - **Validate categorical metadata:** If categorical fields (`category`, `split`, `label`) from CSV/JSON are used to build directory paths, validate them against explicit allowlists/enums before building paths to prevent directory traversal via metadata fields. - **Confine destination directories:** When copying or creating split trees on disk in `prepare_data()`: - Verify that `split_root` is within `root` (`is_within_directory(self.root, self.split_root)`). - Confine leaf folders before creation: `validate_path(folder, base_dir=self.root, should_exist=False).mkdir(...)`. - Confine destination file paths: `validate_path(dest_file, base_dir=self.root, should_exist=False)`. - **Add confinement tests:** Any datamodule that reads external split/metadata files must include unit tests asserting that out-of-root paths (e.g. `../` traversal) raise `ValueError` (see `tests/unit/data/utils/test_path_confinement.py`). ## Reviewer / self-check before opening a PR - [ ] `AnomalibDataset` subclass sets `self.samples` (DataFrame with required columns + `task` attr). - [ ] `AnomalibDataModule` subclass implements `_setup()` only; no unnecessary overrides of `train_dataloader`/`val_dataloader`/`test_dataloader`. - [ ] Any paths parsed from split files, annotations, or metadata are confined to `root` using `resolve_path_under_root` or `validate_path(..., base_dir=root)`. - [ ] Datamodule exported from `src/anomalib/data/__init__.py` and `__all__` updated.