STFPM

This is the implementation of the STFPM paper.

Model Type: Segmentation

Description

STFPM algorithm which consists of a pre-trained teacher network and a student network with identical architecture. The student network learns the distribution of anomaly-free images by matching the features with the counterpart features in the teacher network. Multi-scale feature matching is used to enhance robustness. This hierarchical feature matching enables the student network to receive a mixture of multi-level knowledge from the feature pyramid thus allowing for anomaly detection of various sizes.

During inference, the feature pyramids of teacher and student networks are compared. Larger difference indicates a higher probability of anomaly occurrence.

Architecture

Usage

$ python tools/train.py --model stfpm

PyTorch model for the STFPM model implementation.

class anomalib.models.stfpm.torch_model.STFPMModel(layers: list[str], input_size: tuple[int, int], backbone: str = 'resnet18')

Bases: Module

STFPM: Student-Teacher Feature Pyramid Matching for Unsupervised Anomaly Detection.

Parameters:

• layers (list[str]) – Layers used for feature extraction

• input_size (tuple[int, int]) – Input size for the model.

• backbone (str, optional) – Pre-trained model backbone. Defaults to “resnet18”.

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

forward(images: Tensor) → Tensor | dict[str, Tensor] | tuple[dict[str, Tensor]]

Forward-pass images into the network.

During the training mode the model extracts the features from the teacher and student networks. During the evaluation mode, it returns the predicted anomaly map.

Parameters:

images (Tensor) – Batch of images.

Returns:

Teacher and student features when in training mode, otherwise the predicted anomaly maps.

training: bool

STFPM: Student-Teacher Feature Pyramid Matching for Unsupervised Anomaly Detection.

https://arxiv.org/abs/2103.04257

class anomalib.models.stfpm.lightning_model.StfpmLightning(hparams: DictConfig | ListConfig)

Bases: Stfpm

PL Lightning Module for the STFPM algorithm.

Parameters:

hparams (DictConfig | ListConfig) – Model params

configure_callbacks() → list[EarlyStopping]

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

Configures optimizers.

Note

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

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

Returns:

SGD optimizer

Return type:

Optimizer

Anomaly Map Generator for the STFPM model implementation.

class anomalib.models.stfpm.anomaly_map.AnomalyMapGenerator(image_size: ListConfig | tuple)

Bases: Module

Generate Anomaly Heatmap.

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

compute_anomaly_map(teacher_features: dict[str, Tensor], student_features: dict[str, Tensor]) → torch.Tensor

Compute the overall anomaly map via element-wise production the interpolated anomaly maps.

Parameters:

• teacher_features (dict[str, Tensor]) – Teacher features

• student_features (dict[str, Tensor]) – Student features

Returns:

Final anomaly map

compute_layer_map(teacher_features: Tensor, student_features: Tensor) → Tensor

Compute the layer map based on cosine similarity.

Parameters:

• teacher_features (Tensor) – Teacher features

• student_features (Tensor) – Student features

Returns:

Anomaly score based on cosine similarity.

forward(**kwargs: dict[str, Tensor]) → torch.Tensor

Returns anomaly map.

Expects teach_features and student_features keywords to be passed explicitly.

Example

>>> anomaly_map_generator = AnomalyMapGenerator(image_size=tuple(hparams.model.input_size))
>>> output = self.anomaly_map_generator(
        teacher_features=teacher_features,
        student_features=student_features
    )

Raises:

ValueError – teach_features and student_features keys are not found

Returns:

anomaly map

Return type:

torch.Tensor

training: bool