SuperSimpleNet

SuperSimpleNet#

Architecture#

SuperSimpleNet Architecture

SuperSimpleNet: Unifying Unsupervised and Supervised Learning for Fast and Reliable Surface Defect Detection.

This module implements the SuperSimpleNet model for surface defect / anomaly detection. SuperSimpleNet is a simple yet strong discriminative model consisting of a pretrained feature extractor with upscaling, feature adaptor, train-time feature-level synthetic anomaly generation module, and segmentation-detection module.

Using the adapted features, the model predicts an anomaly map via the segmentation head and an anomaly score using the classification head. It delivers strong performance while maintaining fast inference.

Example

>>> from anomalib.data import MVTec
>>> from anomalib.models import Supersimplenet
>>> from anomalib.engine import Engine

>>> datamodule = MVTec()
>>> model = Supersimplenet()
>>> engine = Engine()

>>> engine.fit(model, datamodule=datamodule)  
>>> predictions = engine.predict(model, datamodule=datamodule)
Paper:

Title: SuperSimpleNet: Unifying Unsupervised and Supervised Learning for Fast and Reliable Surface Defect Detection. URL: https://arxiv.org/pdf/2408.03143

Notes

This implementation supports both unsupervised and supervised setting, but Anomalib currently supports only unsupervised learning.

See also

anomalib.models.image.supersimplenet.torch_model.SupersimplenetModel:

PyTorch implementation of the SuperSimpleNet model.

class anomalib.models.image.supersimplenet.lightning_model.Supersimplenet(perlin_threshold=0.2, backbone='wide_resnet50_2', layers=['layer2', 'layer3'], supervised=False, pre_processor=True, post_processor=True, evaluator=True, visualizer=True)#

Bases: AnomalibModule

PL Lightning Module for the SuperSimpleNet model.

Parameters:

  • perlin_threshold (float) – threshold value for Perlin noise thresholding during anomaly generation.

  • backbone (str) – backbone name

  • layers (list[str]) – backbone layers utilised

  • supervised (bool) – whether the model will be trained in supervised mode. False by default (unsupervised).

  • pre_processor (PreProcessor | bool, optional) – Pre-processor instance or flag to use default. Defaults to True.

  • post_processor (PostProcessor | bool, optional) – Post-processor instance or flag to use default. Defaults to True.

  • evaluator (Evaluator | bool, optional) – Evaluator instance or flag to use default. Defaults to True.

  • visualizer (Visualizer | bool, optional) – Visualizer instance or flag to use default. Defaults to True.

configure_optimizers()#

Configure AdamW optimizer and MultiStepLR scheduler.

Return type:

Union[Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRScheduler, ReduceLROnPlateau, LRSchedulerConfig]]], OptimizerConfig, OptimizerLRSchedulerConfig, Sequence[OptimizerConfig], Sequence[OptimizerLRSchedulerConfig], None]

classmethod configure_pre_processor(image_size=None)#

Configure the default pre-processor for SuperSimpleNet.

Pre-processor resizes images and normalizes using ImageNet statistics.

Parameters:

image_size (tuple[int, int] | None, optional) – Target size for resizing. Defaults to (256, 256).

Returns:

Configured SuperSimpleNet pre-processor

Return type:

PreProcessor

property learning_type: LearningType#

Return the learning type of the model.

This is subject to change in the future when support for supervised training is introduced.

Returns:

Learning type of the model.

Return type:

LearningType

property trainer_arguments: dict[str, Any]#

Return SuperSimpleNet trainer arguments.

training_step(batch, *args, **kwargs)#

Perform the training step input and return the loss.

Parameters:

  • (batch (batch) – dict[str, str | torch.Tensor]): Input batch

  • args – Additional arguments.

  • kwargs – Additional keyword arguments.

Returns:

Dictionary containing the loss value.

Return type:

STEP_OUTPUT

validation_step(batch, *args, **kwargs)#

Perform the validation step and return the anomaly map and anomaly score.

Parameters:

  • batch (dict[str, str | torch.Tensor]) – Input batch

  • args – Additional arguments.

  • kwargs – Additional keyword arguments.

Returns:

batch dictionary containing anomaly-maps.

Return type:

STEP_OUTPUT | None

PyTorch model for the SuperSimpleNet model implementation.

See also

anomalib.models.image.supersimplenet.lightning_model.Supersimplenet:

SuperSimpleNet Lightning model.

class anomalib.models.image.supersimplenet.torch_model.AnomalyMapGenerator(sigma)#

Bases: Module

Final anomaly map generator, responsible for upscaling and smoothing.

Parameters:

sigma (float)

forward(out_map, final_size)#

Upscale and smooth anomaly map to get final anomaly map of same size as input image.

Parameters:

  • out_map (torch.Tensor) – output anomaly map from seg. head.

  • final_size (tuple[int, int]) – size (h, w) of final anomaly map.

Returns:

final anomaly map.

Return type:

torch.Tensor

class anomalib.models.image.supersimplenet.torch_model.FeatureAdapter(channel_dim)#

Bases: Module

Feature adapter used to adapt raw features for the task of anomaly detection.

Parameters:

channel_dim (int) – channel dimension of features.

forward(features)#

Adapt features.

Parameters:

features (torch.Tensor) – input features

Return type:

Tensor

Returns:

(torch.Tensor) adapted features

class anomalib.models.image.supersimplenet.torch_model.FeatureExtractor(backbone, layers, patch_size=3)#

Bases: Module

Feature extractor module.

Parameters:

  • backbone (str) – backbone name.

  • layers (list[str]) – list of layers used for extraction.

forward(input_tensor)#

Extract features from input tensor.

Parameters:

input_tensor (Tensor) – input tensor (images)

Returns:

extracted feature map.

Return type:

(torch.Tensor)

get_channels_dim()#

Get feature channel dimension.

Returns:

feature channel dimension.

Return type:

(int)

class anomalib.models.image.supersimplenet.torch_model.SegmentationDetectionModule(channel_dim, stop_grad=False)#

Bases: Module

SegmentationDetection module responsible for prediction of anomaly map and score.

Parameters:

  • channel_dim (int) – channel dimension of features.

  • stop_grad (bool) – whether to stop gradient from class. head to seg. head.

forward(features)#

Predict anomaly map and anomaly score.

Parameters:

features (Tensor) – adapted features.

Return type:

tuple[Tensor, Tensor]

Returns:

predicted anomaly map and score.

get_params()#

Get segmentation and classification head parameters.

Return type:

tuple[list[Parameter], list[Parameter]]

Returns:

seg. head parameters and class. head parameters.

class anomalib.models.image.supersimplenet.torch_model.SupersimplenetModel(perlin_threshold=0.2, backbone='wide_resnet50_2', layers=['layer2', 'layer3'], stop_grad=True)#

Bases: Module

SuperSimpleNet Pytorch model.

It consists of feature extractor, feature adaptor, anomaly generation mechanism and segmentation-detection module.

Parameters:

  • perlin_threshold (float) – threshold value for Perlin noise thresholding during anomaly generation.

  • backbone (str) – backbone name

  • layers (list[str]) – backbone layers utilised

  • stop_grad (bool) – whether to stop gradient from class. to seg. head.

static downsample_mask(masks, feat_h, feat_w)#

Downsample the masks according to the feature dimensions.

Primarily used in supervised setting.

Parameters:

  • masks (torch.Tensor) – input GT masks

  • feat_h (int) – feature height.

  • feat_w (int) – feature width.

Returns:

downsampled masks.

Return type:

(torch.Tensor)

forward(images, masks=None, labels=None)#

SuperSimpleNet forward pass.

Extract and process features, adapt them, generate anomalies (train only) and predict anomaly map and score.

Parameters:
Returns:

anomaly map and score training: anomaly map, score and GT masks and labels

Return type:

inference

anomalib.models.image.supersimplenet.torch_model.init_weights(module)#

Init weight of the model.

Parameters:

module (nn.Module) – torch module.

Return type:

None

Anomaly generator for the SuperSimplenet model implementation.

class anomalib.models.image.supersimplenet.anomaly_generator.AnomalyGenerator(noise_mean, noise_std, threshold)#

Bases: Module

Anomaly generator for the SuperSimpleNet model.

Parameters:

  • noise_mean (float) – Mean of the Gaussian noise distribution.

  • noise_std (float) – Standard deviation of the Gaussian noise distribution.

  • threshold (float) – Threshold used to binarize Perlin noise.

forward(features, mask, labels)#

Generate anomaly on features using thresholded perlin noise and Gaussian noise.

Also update GT masks and labels with new anomaly information.

Parameters:
Return type:

tuple[Tensor, Tensor, Tensor]

Returns:

perturbed features, updated GT masks and labels.

generate_perlin(batches, height, width)#

Generate 2d perlin noise masks with dims [b, 1, h, w].

Parameters:

  • batches (int) – number of batches (different masks)

  • height (int) – height of features

  • width (int) – width of features

Return type:

Tensor

Returns:

tensor with b perlin binarized masks

static next_power_2(num)#

Get the next power of 2 for given number.

Parameters:

num (int) – value of interest

Return type:

int

Returns:

next power of 2 value for given number