Anomalib Adding A Model

Adding a New Model

Models live under src/anomalib/models/image/<model_name>/ (or .../video/<model_name>/ for video models). Every model is a LightningModule subclass of AnomalibModule (src/anomalib/models/components/base/anomalib_module.py) that wraps a plain torch.nn.Module.

Reference implementation

Read src/anomalib/models/image/padim/ first — it is the smallest complete example:

• torch_model.py — PadimModel(nn.Module): pure PyTorch forward pass. In training mode it returns intermediate embeddings; in eval mode it returns an InferenceBatch(pred_score=..., anomaly_map=...).
• anomaly_map.py — AnomalyMapGenerator(nn.Module): post-processing (score/map computation) kept out of torch_model.py for clarity. Not all models need a separate file for this.
• lightning_model.py — Padim(MemoryBankMixin, AnomalibModule): the Lightning-facing wrapper. This is the class users construct (Padim()), pass to Engine, and reference from CLI/config as anomalib.models.Padim.
• __init__.py — exports the Lightning class only: from .lightning_model import Padim.
• README.md — usage + benchmark notes (see model-doc-sync/model-sample-image-export skills for docs work).

Only mix in MemoryBankMixin (src/anomalib/models/components/base/memory_bank_module.py) if the model accumulates a memory bank / feature bank across training (as PaDiM and PatchCore do). Most models just subclass AnomalibModule directly.

Required members on your AnomalibModule subclass

import torch
from anomalib import LearningType
from anomalib.data import Batch
from anomalib.metrics import Evaluator
from anomalib.models.components import AnomalibModule
from anomalib.post_processing import PostProcessor
from anomalib.pre_processing import PreProcessor
from anomalib.visualization import Visualizer

class MyModel(AnomalibModule):
    def __init__(self, some_param: int = 1, pre_processor: PreProcessor | bool = True,
                 post_processor: PostProcessor | bool = True,
                 evaluator: Evaluator | bool = True,
                 visualizer: Visualizer | bool = True) -> None:
        super().__init__(pre_processor=pre_processor, post_processor=post_processor,
                          evaluator=evaluator, visualizer=visualizer)
        self.model = MyTorchModel(some_param=some_param)

    @property
    def trainer_arguments(self) -> dict:
        """Default Trainer overrides for this model, e.g. {"max_epochs": 1} for training-free models."""
        return {"max_epochs": 1, "num_sanity_val_steps": 0}

    @property
    def learning_type(self) -> LearningType:
        return LearningType.ONE_CLASS

    def training_step(self, batch: Batch, *args, **kwargs) -> torch.Tensor:
        # Training-free models (e.g. Padim) still return a dummy loss for Lightning.
        _ = self.model(batch.image)
        return torch.tensor(0.0, requires_grad=True, device=self.device)

    def validation_step(self, batch: Batch, *args, **kwargs) -> Batch:
        predictions = self.model(batch.image)
        return batch.update(**predictions._asdict())

    def configure_optimizers(self) -> None:
        # Return None for training-free / statistical models (Padim does this).
        return None

AnomalibModule provides working defaults you can override only when the model needs something different:

• configure_pre_processor(image_size=None) — resize + ImageNet normalization.
• configure_post_processor() — thresholding/normalization for ONE_CLASS models.
• configure_evaluator() — image/pixel AUROC and F1 metrics.
• configure_visualizer() — default ImageVisualizer.

pre_processor / post_processor / evaluator / visualizer constructor args each accept an instance, True (use the configured default), or False (disable).

Registration — how the model becomes discoverable

1. Add the export in src/anomalib/models/image/__init__.py (or video/__init__.py):

   from .my_model import MyModel

   and add "MyModel" to __all__ and to the Available Models docstring list at the top of the file.

2. Also add the import in src/anomalib/models/__init__.py, which explicitly imports all image (and video) models and defines the top-level __all__. Without this, anomalib.models.MyModel won't resolve.

3. That's the only registration needed. list_models() and get_model() (src/anomalib/models/__init__.py) discover models by walking AnomalibModule.__subclasses__() and matching on cls.__name__ (case-insensitive) — there is no separate name-string registry to update.

4. get_model() also accepts a dict/DictConfig/Namespace with class_path + init_args, restricted to modules listed in ALLOWED_MODULES (same file) — you don't need to touch that set for models already under anomalib.models.

5. Once exported, the model is usable as anomalib.models.MyModel, from get_model("MyModel"), and from the CLI: anomalib train --model MyModel --data anomalib.data.MVTecAD.