Overview
The PatchCore class is the core implementation of the PatchCore algorithm for industrial anomaly detection. It provides methods for loading models, training on normal data, and predicting anomalies in images.
Class Definition
class PatchCore(torch.nn.Module):
def __init__(self, device)
torch.nn.Module.
The device (CPU or CUDA) on which to run the model.
Methods
load
def load(
self,
backbone,
layers_to_extract_from,
device,
input_shape,
pretrain_embed_dimension,
target_embed_dimension,
patchsize=3,
patchstride=1,
anomaly_score_num_nn=1,
featuresampler=patchcore.sampler.IdentitySampler(),
nn_method=patchcore.common.FaissNN(False, 4),
**kwargs
)
Parameters
Pre-trained neural network backbone for feature extraction (e.g., ResNet, WideResNet).
List of layer names from which to extract features from the backbone.
Device on which to load the model.
Shape of input images (e.g., (3, 224, 224) for RGB images).
Dimension to which features are projected before aggregation.
Target dimension for final feature embeddings.
Size of patches to extract from feature maps.
Stride for patch extraction from feature maps.
Number of nearest neighbors to use for anomaly scoring.
featuresampler
Sampler
default:"IdentitySampler()"
Sampler to use for subsampling features during training. Can be IdentitySampler, GreedyCoresetSampler, or RandomSampler.
nn_method
FaissNN
default:"FaissNN(False, 4)"
Nearest neighbor search method for anomaly scoring.
Example
import torch
import patchcore.backbones
from patchcore.patchcore import PatchCore
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = PatchCore(device)
# Load WideResNet50 backbone
backbone = patchcore.backbones.load("wideresnet50")
model.load(
backbone=backbone,
layers_to_extract_from=["layer2", "layer3"],
device=device,
input_shape=(3, 224, 224),
pretrain_embed_dimension=1024,
target_embed_dimension=1024,
patchsize=3,
anomaly_score_num_nn=1
)
fit
def fit(self, training_data)
training_data
DataLoader or iterable
required
DataLoader or iterable containing normal (non-anomalous) training images. Images should be dictionaries with an "image" key or raw tensors.
Example
from torch.utils.data import DataLoader
from torchvision import transforms
# Assuming you have a dataset of normal images
train_loader = DataLoader(train_dataset, batch_size=32)
# Train the model
model.fit(train_loader)
The fit method only requires normal (non-anomalous) images. PatchCore learns what “normal” looks like and detects deviations from this during inference.
predict
Computes anomaly scores and segmentation masks for input images.
data
torch.Tensor or DataLoader
required
Input image tensor or DataLoader. If a tensor, shape should be [B, C, H, W]. If a DataLoader, images can be dictionaries with "image", "is_anomaly", and "mask" keys.
Returns
List of anomaly scores for each image. Higher scores indicate higher likelihood of anomaly.
List of anomaly segmentation masks for each image. Values are normalized between 0 and 1.
Ground truth labels (only returned when using DataLoader with labeled data).
Ground truth masks (only returned when using DataLoader with labeled data).
Example
import torch
# Single batch prediction
test_images = torch.randn(4, 3, 224, 224).to(device)
scores, masks = model.predict(test_images)
for i, (score, mask) in enumerate(zip(scores, masks)):
print(f"Image {i}: Anomaly Score = {score:.4f}")
# mask is a numpy array with shape [H, W]
# DataLoader prediction with ground truth
test_loader = DataLoader(test_dataset, batch_size=32)
scores, masks, labels_gt, masks_gt = model.predict(test_loader)
embed
Extracts feature embeddings from input images without computing anomaly scores.
data
torch.Tensor or DataLoader
required
Input images as tensor [B, C, H, W] or DataLoader.
Returns
features
np.ndarray or list[np.ndarray]
Feature embeddings extracted from the images. Shape depends on the number of patches and target embedding dimension.
save_to_path
def save_to_path(self, save_path: str, prepend: str = "") -> None
Directory path where the model will be saved.
Optional prefix to prepend to saved filenames.
Example
# Save the trained model
model.save_to_path("./models/patchcore_trained")
# Save with a prefix
model.save_to_path("./models", prepend="bottle_")
load_from_path
def load_from_path(
self,
load_path: str,
device: torch.device,
nn_method: patchcore.common.FaissNN,
prepend: str = ""
) -> None
Directory path from which to load the model.
Device on which to load the model.
Nearest neighbor search method to use.
Optional prefix used when saving the model.
Example
import torch
from patchcore.patchcore import PatchCore
import patchcore.common
device = torch.device("cuda")
model = PatchCore(device)
nn_method = patchcore.common.FaissNN(False, 4)
model.load_from_path(
"./models/patchcore_trained",
device=device,
nn_method=nn_method
)
# Now you can use the model for prediction
scores, masks = model.predict(test_images)
PatchMaker Class
Internal utility class for patch extraction and scoring.
class PatchMaker:
def __init__(self, patchsize, stride=None)
Size of patches to extract.
Stride for patch extraction. If None, defaults to patchsize.
Methods
patchify
def patchify(self, features, return_spatial_info=False)
Input feature tensor with shape [B, C, H, W].
Whether to return spatial information about patches.
unpatch_scores
def unpatch_scores(self, x, batchsize)
score
Aggregates patch scores to produce image-level scores by taking the maximum across dimensions.