Draem

This is the implementation of the DRAEM paper.

Model Type: Segmentation

Description

DRAEM is a reconstruction based algorithm that consists of a reconstructive subnetwork and a discriminative subnetwork. DRAEM is trained on simulated anomaly images, generated by augmenting normal input images from the training set with a random Perlin noise mask extracted from an unrelated source of image data. The reconstructive subnetwork is an autoencoder architecture that is trained to reconstruct the original input images from the augmented images. The reconstructive submodel is trained using a combination of L2 loss and Structural Similarity loss. The input of the discriminative subnetwork consists of the channel-wise concatenation of the (augmented) input image and the output of the reconstructive subnetwork. The output of the discriminative subnetwork is an anomaly map that contains the predicted anomaly scores for each pixel location. The discriminative subnetwork is trained using Focal Loss.

For optimal results, DRAEM requires specifying the path to a folder of image data that will be used as the source of the anomalous pixel regions in the simulated anomaly images. The path can be specified by editing the value of the model.anomaly_source_path parameter in the config.yaml file. The authors of the original paper recommend using the DTD dataset as anomaly source.

Architecture

DRAEM Architecture

Usage

$ python tools/train.py --model draem

PyTorch model for the DRAEM model implementation.

class anomalib.models.draem.torch_model.DecoderDiscriminative(base_width: int, out_channels: int = 1)[source]

Bases: Module

Decoder part of the discriminator network.

Parameters:

  • base_width (int) – Base dimensionality of the layers of the autoencoder.

  • out_channels (int) – Number of output channels.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(act1: Tensor, act2: Tensor, act3: Tensor, act4: Tensor, act5: Tensor, act6: Tensor) → Tensor[source]

Computes predicted anomaly class scores from the intermediate outputs of the encoder sub network.

Parameters:

  • act1 (Tensor) – Encoder activations of the first block of convolutional layers.

  • act2 (Tensor) – Encoder activations of the second block of convolutional layers.

  • act3 (Tensor) – Encoder activations of the third block of convolutional layers.

  • act4 (Tensor) – Encoder activations of the fourth block of convolutional layers.

  • act5 (Tensor) – Encoder activations of the fifth block of convolutional layers.

  • act6 (Tensor) – Encoder activations of the sixth block of convolutional layers.

Returns:

Predicted anomaly class scores per pixel.

training: bool
class anomalib.models.draem.torch_model.DecoderReconstructive(base_width: int, out_channels: int = 1)[source]

Bases: Module

Decoder part of the reconstructive network.

Parameters:

  • base_width (int) – Base dimensionality of the layers of the autoencoder.

  • out_channels (int) – Number of output channels.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(act5: Tensor) → Tensor[source]

Reconstruct the image from the activations of the bottleneck layer.

Parameters:

act5 (Tensor) – Activations of the bottleneck layer.

Returns:

Batch of reconstructed images.

training: bool
class anomalib.models.draem.torch_model.DiscriminativeSubNetwork(in_channels: int = 3, out_channels: int = 3, base_width: int = 64)[source]

Bases: Module

Discriminative model that predicts the anomaly mask from the original image and its reconstruction.

Parameters:

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • base_width (int) – Base dimensionality of the layers of the autoencoder.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(batch: Tensor) → Tensor[source]

Generate the predicted anomaly masks for a batch of input images.

Parameters:

batch (Tensor) – Batch of inputs consisting of the concatenation of the original images and their reconstructions.

Returns:

Activations of the output layer corresponding to the normal and anomalous class scores on the pixel level.

training: bool
class anomalib.models.draem.torch_model.DraemModel(sspcab: bool = False)[source]

Bases: Module

DRAEM PyTorch model consisting of the reconstructive and discriminative sub networks.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(batch: Tensor) → Tensor | tuple[Tensor, Tensor][source]

Compute the reconstruction and anomaly mask from an input image.

Parameters:

x (Tensor) – batch of input images

Returns:

Predicted confidence values of the anomaly mask. During training the reconstructed input images are returned as well.

training: bool
class anomalib.models.draem.torch_model.EncoderDiscriminative(in_channels: int, base_width: int)[source]

Bases: Module

Encoder part of the discriminator network.

Parameters:

  • in_channels (int) – Number of input channels.

  • base_width (int) – Base dimensionality of the layers of the autoencoder.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(batch: Tensor) → tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor][source]

Convert the inputs to the salient space by running them through the encoder network.

Parameters:

batch (Tensor) – Batch of inputs consisting of the concatenation of the original images and their reconstructions.

Returns:

Computed feature maps for each of the layers in the encoder sub network.

training: bool
class anomalib.models.draem.torch_model.EncoderReconstructive(in_channels: int, base_width: int, sspcab: bool = False)[source]

Bases: Module

Encoder part of the reconstructive network.

Parameters:

  • in_channels (int) – Number of input channels.

  • base_width (int) – Base dimensionality of the layers of the autoencoder.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(batch: Tensor) → Tensor[source]

Encode a batch of input images to the salient space.

Parameters:

batch (Tensor) – Batch of input images.

Returns:

Feature maps extracted from the bottleneck layer.

training: bool
class anomalib.models.draem.torch_model.ReconstructiveSubNetwork(in_channels: int = 3, out_channels: int = 3, base_width=128, sspcab: bool = False)[source]

Bases: Module

Autoencoder model that encodes and reconstructs the input image.

Parameters:

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • base_width (int) – Base dimensionality of the layers of the autoencoder.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(batch: Tensor) → Tensor[source]

Encode and reconstruct the input images.

Parameters:

batch (Tensor) – Batch of input images

Returns:

Batch of reconstructed images.

training: bool

DRÆM – A discriminatively trained reconstruction embedding for surface anomaly detection.

Paper https://arxiv.org/abs/2108.07610

class anomalib.models.draem.lightning_model.Draem(enable_sspcab: bool = False, sspcab_lambda: float = 0.1, anomaly_source_path: str | None = None)[source]

Bases: AnomalyModule

DRÆM: A discriminatively trained reconstruction embedding for surface anomaly detection.

Parameters:

anomaly_source_path (str | None) – Path to folder that contains the anomaly source images. Random noise will be used if left empty.

setup_sspcab() → None[source]

Prepare the model for the SSPCAB training step by adding forward hooks for the SSPCAB layer activations.

training_step(batch: dict[str, str | Tensor], *args, **kwargs) → STEP_OUTPUT[source]

Training Step of DRAEM.

Feeds the original image and the simulated anomaly image through the network and computes the training loss.

Parameters:

batch (dict[str, str | Tensor]) – Batch containing image filename, image, label and mask

Returns:

Loss dictionary

validation_step(batch: dict[str, str | Tensor], *args, **kwargs) → STEP_OUTPUT[source]

Validation step of DRAEM. The Softmax predictions of the anomalous class are used as anomaly map.

Parameters:

batch (dict[str, str | Tensor]) – Batch of input images

Returns:

Dictionary to which predicted anomaly maps have been added.

class anomalib.models.draem.lightning_model.DraemLightning(hparams: DictConfig | ListConfig)[source]

Bases: Draem

DRÆM: A discriminatively trained reconstruction embedding for surface anomaly detection.

Parameters:

hparams (DictConfig | ListConfig) – Model parameters

configure_callbacks() → list[EarlyStopping][source]

Configure model-specific callbacks.

Note

This method is used for the existing CLI. When PL CLI is introduced, configure callback method will be

deprecated, and callbacks will be configured from either config.yaml file or from CLI.

configure_optimizers() → Optimizer[source]

Configure the Adam optimizer.