PatchCore

This is the implementation of the PatchCore paper.

Model Type: Segmentation

Description

The PatchCore algorithm is based on the idea that an image can be classified as anomalous as soon as a single patch is anomalous. The input image is tiled. These tiles act as patches which are fed into the neural network. It consists of a single pre-trained network which is used to extract “mid” level features patches. The “mid” level here refers to the feature extraction layer of the neural network model. Lower level features are generally too broad and higher level features are specific to the dataset the model is trained on. The features extracted during training phase are stored in a memory bank of neighbourhood aware patch level features.

During inference this memory bank is coreset subsampled. Coreset subsampling generates a subset which best approximates the structure of the available set and allows for approximate solution finding. This subset helps reduce the search cost associated with nearest neighbour search. The anomaly score is taken as the maximum distance between the test patch in the test patch collection to each respective nearest neighbour.

Architecture

Usage

$ python tools/train.py --model patchcore

PyTorch model for the PatchCore model implementation.

class anomalib.models.patchcore.torch_model.PatchcoreModel(input_size: tuple[int, int], layers: list[str], backbone: str = 'wide_resnet50_2', pre_trained: bool = True, num_neighbors: int = 9)

Bases: DynamicBufferModule, Module

Patchcore Module.

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

compute_anomaly_score(patch_scores: Tensor, locations: Tensor, embedding: Tensor) → Tensor

Compute Image-Level Anomaly Score.

Parameters:

• patch_scores (Tensor) – Patch-level anomaly scores

• locations – Memory bank locations of the nearest neighbor for each patch location

• embedding – The feature embeddings that generated the patch scores

Returns:

Image-level anomaly scores

Return type:

Tensor

forward(input_tensor: Tensor) → Tensor | tuple[Tensor, Tensor]

Return Embedding during training, or a tuple of anomaly map and anomaly score during testing.

Steps performed: 1. Get features from a CNN. 2. Generate embedding based on the features. 3. Compute anomaly map in test mode.

Parameters:

input_tensor (Tensor) – Input tensor

Returns:

Embedding for training,

anomaly map and anomaly score for testing.

Return type:

Tensor | tuple[Tensor, Tensor]

generate_embedding(features: dict[str, Tensor]) → Tensor

Generate embedding from hierarchical feature map.

Parameters:

• features – Hierarchical feature map from a CNN (ResNet18 or WideResnet)

• features – dict[str:Tensor]:

Returns:

Embedding vector

nearest_neighbors(embedding: Tensor, n_neighbors: int) → tuple[Tensor, Tensor]

Nearest Neighbours using brute force method and euclidean norm.

Parameters:

• embedding (Tensor) – Features to compare the distance with the memory bank.

• n_neighbors (int) – Number of neighbors to look at

Returns:

Patch scores. Tensor: Locations of the nearest neighbor(s).

Return type:

Tensor

static reshape_embedding(embedding: Tensor) → Tensor

Reshape Embedding.

Reshapes Embedding to the following format: [Batch, Embedding, Patch, Patch] to [Batch*Patch*Patch, Embedding]

Parameters:

embedding (Tensor) – Embedding tensor extracted from CNN features.

Returns:

Reshaped embedding tensor.

Return type:

Tensor

subsample_embedding(embedding: Tensor, sampling_ratio: float) → None

Subsample embedding based on coreset sampling and store to memory.

Parameters:

• embedding (np.ndarray) – Embedding tensor from the CNN

• sampling_ratio (float) – Coreset sampling ratio

training: bool

Towards Total Recall in Industrial Anomaly Detection.

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

class anomalib.models.patchcore.lightning_model.Patchcore(input_size: tuple[int, int], backbone: str, layers: list[str], pre_trained: bool = True, coreset_sampling_ratio: float = 0.1, num_neighbors: int = 9)

Bases: AnomalyModule

PatchcoreLightning Module to train PatchCore algorithm.

Parameters:

• input_size (tuple[int, int]) – Size of the model input.

• backbone (str) – Backbone CNN network

• layers (list[str]) – Layers to extract features from the backbone CNN

• pre_trained (bool, optional) – Boolean to check whether to use a pre_trained backbone.

• coreset_sampling_ratio (float, optional) – Coreset sampling ratio to subsample embedding. Defaults to 0.1.

• num_neighbors (int, optional) – Number of nearest neighbors. Defaults to 9.

configure_optimizers() → None

Configure optimizers.

Returns:

Do not set optimizers by returning None.

Return type:

None

on_validation_start() → None

Apply subsampling to the embedding collected from the training set.

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

Generate feature embedding of the batch.

Parameters:

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

Returns:

Embedding Vector

Return type:

dict[str, np.ndarray]

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

Get batch of anomaly maps from input image batch.

Parameters:

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

Returns:

Image filenames, test images, GT and predicted label/masks

Return type:

dict[str, Any]

class anomalib.models.patchcore.lightning_model.PatchcoreLightning(hparams)

Bases: Patchcore

PatchcoreLightning Module to train PatchCore algorithm.

Parameters:

hparams (DictConfig | ListConfig) – Model params

Anomaly Map Generator for the PatchCore model implementation.

class anomalib.models.patchcore.anomaly_map.AnomalyMapGenerator(input_size: ListConfig | tuple, sigma: int = 4)

Bases: Module

Generate Anomaly Heatmap.

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

compute_anomaly_map(patch_scores: Tensor) → Tensor

Pixel Level Anomaly Heatmap.

Parameters:

patch_scores (Tensor) – Patch-level anomaly scores

Returns:

Map of the pixel-level anomaly scores

Return type:

Tensor

forward(patch_scores: Tensor) → Tensor

Returns anomaly_map and anomaly_score.

Parameters:

patch_scores (Tensor) – Patch-level anomaly scores

Example >>> anomaly_map_generator = AnomalyMapGenerator(input_size=input_size) >>> map = anomaly_map_generator(patch_scores=patch_scores)

Returns:

anomaly_map

Return type:

Tensor

training: bool