anomalib.models.draem.torch_model

PyTorch model for the DRAEM model implementation.

Module Contents

Classes

DraemModel

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

ReconstructiveSubNetwork

Autoencoder model that encodes and reconstructs the input image.

DiscriminativeSubNetwork

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

EncoderDiscriminative

Encoder part of the discriminator network.

DecoderDiscriminative

Decoder part of the discriminator network.

EncoderReconstructive

Encoder part of the reconstructive network.

DecoderReconstructive

Decoder part of the reconstructive network.
class anomalib.models.draem.torch_model.DraemModel[source]

Bases: torch.nn.Module

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

forward(batch: torch.Tensor) → Union[torch.Tensor, Tuple[torch.Tensor, torch.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.

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

Bases: torch.nn.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.

forward(batch: torch.Tensor) → torch.Tensor[source]

Encode and reconstruct the input images.

Parameters

batch (Tensor) – Batch of input images

Returns

Batch of reconstructed images.

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

Bases: torch.nn.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.

forward(batch: torch.Tensor) → torch.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.

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

Bases: torch.nn.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.

forward(batch: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.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.

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

Bases: torch.nn.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.

forward(act1: torch.Tensor, act2: torch.Tensor, act3: torch.Tensor, act4: torch.Tensor, act5: torch.Tensor, act6: torch.Tensor) → torch.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.

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

Bases: torch.nn.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.

forward(batch: torch.Tensor) → torch.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.

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

Bases: torch.nn.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.

forward(act5: torch.Tensor) → torch.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.