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 patterns are similar, but base classes and split behavior differ by modality.) ## 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. - For image and depth datamodules, use `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. - It also already implements `_create_test_split()` / `_create_val_split()` — don't hand-roll `random_split()` + `concatenate_datasets()` in `_setup()`; just assign `self.train_data` / `self.test_data` from directory splits and let the base class do the rest (see `MPDD`/`BTech` `_setup()` for the minimal pattern). The two methods trigger differently, though: - `_create_test_split()` samples normal images from `train_data` only when `test_data` lacks normal samples, `test_split_mode` is not `NONE`, and `test_split_ratio` is set. - `_create_val_split()` is driven purely by `val_split_mode`, not by missing samples: for `FROM_TRAIN`/`FROM_TEST`/`SAME_AS_TEST`/`SYNTHETIC` it auto-derives `val_data`; for `FROM_DIR` it does nothing, so you must assign `self.val_data` yourself in `_setup()` if you support `FROM_DIR`. - 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`. - Validate dataset-specific params (e.g. `category`, `modality`) against an explicit allowlist and raise a clear `ValueError` listing valid options — see `AutoVI`/`RealIAD`'s `category not in CATEGORIES` check. Without this, a typo just surfaces as a generic "found 0 images" error. - Video datamodules use `AnomalibVideoDataModule` (`src/anomalib/data/datamodules/base/video.py`): its `_create_test_split()` is a no-op, and it rejects `SYNTHETIC` validation. Follow that base class rather than applying the image/depth split guidance above. ## 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. ## Module docstring: License and Reference For modules implementing a published dataset or its datamodule, include `License:` and `Reference:` sections in the module docstring (see `bmad.py` for an example). Verify the license against the dataset's actual source (the Hugging Face dataset card or the dataset repo's own `LICENSE`), not the paper's code repository — the two are often different (e.g. code under BSD/MIT while the data itself is CC BY). Include the arXiv ID and/or DOI link in `Reference:`. Generic format loaders or adapters that do not represent a published dataset should not invent a dataset license or paper reference. ## 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 ``` ### `prepare_data()` — downloading the dataset Check the dataset's actual license and hosting terms before choosing a pattern. A HEAD request can be a preliminary probe without transferring the archive: `curl -s -o /dev/null -w '%{http_code}' -I -L <direct-file-url>`. Some hosts reject HEAD, and a redirect to a login page may still return `200`; a `401`/`403` alone does not prove the dataset is gated. Confirm access from the source's terms or with a minimal ranged GET to the download endpoint that `prepare_data()` will use. - **Open dataset** (direct download link, no auth): use `DownloadInfo` + `download_and_extract` from `anomalib.data.utils` — see `BMAD`/`MVTecAD` `prepare_data()`. Define a module-level `DOWNLOAD_INFO` with `name`, `url`, and `hashsum`. - **Gated dataset** (e.g. Hugging Face dataset requiring click-through access): use `huggingface_hub` conditionally — if the package or token (`HF_TOKEN` env var or cached `hf auth login`) is unavailable, raise `FileNotFoundError` with manual-download instructions. Pass an immutable commit SHA as `revision` to `hf_hub_download` (see `HF_REVISION` in `Kaputt`). Translate only expected authentication, network, or remote-availability errors into the manual-download fallback. Do not blanket-catch download errors: integrity, archive-validation, extraction, and local filesystem failures should remain visible. ## 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`. 3. Add a CLI config at `examples/configs/data/my_dataset.yaml` (see `examples/configs/data/bmad.yaml` for the `class_path`/`init_args` format). 4. Add a docs page under the matching modality's reference folder — `docs/source/markdown/guides/reference/data/datamodules/{image,video,depth}/my_dataset.md` (copy an existing page from that same folder, e.g. `image/bmad.md` — it's just an `automodule` stub). Add a grid card + toctree entry to that modality's `index.md` and a card under the matching modality in `docs/source/markdown/guides/reference/data/datamodules/index.md`. 5. Add a `CHANGELOG.md` entry under `## [Unreleased]`. ## 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"