tensor-md 0.1.4

Location-aware tensor Mahalanobis anomaly detection

tensor-md

tensor-md detects and localizes unusual image regions from normal training images. It keeps each local observation as a tensor and scores it with a location-aware tensor Mahalanobis model using separable covariance factors.

The package is label-free during training: provide a folder of normal images, fit the detector, and then score images from another folder. Larger scores mean that a patch differs more strongly from the normal variation learned at the same spatial location.

Installation

Install the core package from PyPI:

python -m pip install tensor-md

The PyPI distribution is named tensor-md; import it in Python as tensor_md, because Python module names cannot contain hyphens.

Install the optional CNN dependencies when using the built-in feature extractors:

python -m pip install "tensor-md[cnn]"

Diagnostics and notebooks are available through the evaluation and notebooks extras.

Quick start

The following example learns directly from image patches. The training and scoring directories may contain PNG, JPEG, BMP, or TIFF images; no category name, labels, or MVTec layout is required.

from tensor_md import (
    LocationAwareTensorMahalanobisDetector,
    PatchExtractionConfig,
    load_image_patches,
    load_normal_patches,
)

config = PatchExtractionConfig(
    train_image_dir="data/normal",
    image_size=(256, 256),
    patch_size=(16, 16),
    stride=16,
)

normal = load_normal_patches(config)
images = load_image_patches("data/to_check", config)

detector = LocationAwareTensorMahalanobisDetector(
    patches_per_image=normal.patches_per_image,
)
detector.fit(normal)
scores = detector.score(images)

# Rows are images and columns are spatial patch locations.
score_maps = scores.reshape(len(images.image_paths), images.patches_per_image)

Training and scoring images must use the same configuration. Location-aware modelling is most useful when images are approximately aligned, so the same grid location usually represents the same object part or texture region.

Using CNN features

Set input_representation="cnn_features" to model intermediate CNN features instead of RGB patches. This example uses two built-in ResNet50 layers and reduces both to 128 channels before stacking them as a tensor mode:

config = PatchExtractionConfig(
    train_image_dir="data/normal",
    input_representation="cnn_features",
    cnn_backbone="ResNet50",
    cnn_layer_names=("conv3_block4_out", "conv4_block6_out"),
    cnn_dimensionality_reduction="pca",
    cnn_reduction_dimensions=128,
    cnn_feature_patch_size=(1, 1),
)

Available dimensionality-reduction settings are:

• "pca": fit one channel PCA per selected layer using normal training descriptors.
• "random": retain a reproducible random subset of channels from each layer.
• "none" or None: keep all channels.

Multiple feature maps must have the same retained channel count so they can be stacked. Without reduction, their original channel counts must already match. The loader raises a clear error if cnn_reduction_dimensions is larger than the channel count of any selected layer.

Supplying any CNN

The tensor detector is not tied to ResNet. Pass a callable through cnn_feature_extractor. It receives an NHWC float32 batch with values in [0, 1] and returns either one NHWC feature-map batch or a list of them:

def extract_features(batch):
    # Return shape: (N, H, W, C), or a list of such arrays.
    return my_model(batch)

config = PatchExtractionConfig(
    train_image_dir="data/normal",
    input_representation="cnn_features",
    cnn_feature_extractor=extract_features,
)

Keras and PyTorch models can also be wrapped with the convenience adapter:

from tensor_md import make_cnn_feature_extractor

extractor = make_cnn_feature_extractor(model, framework="pytorch")
config = PatchExtractionConfig(
    train_image_dir="data/normal",
    input_representation="cnn_features",
    cnn_feature_extractor=extractor,
)

The adapter handles the common NCHW/NHWC layout conversion. A custom extractor remains appropriate for models with unusual inputs or outputs.

Neighbourhood scoring

The neighbourhood detector pools completed Mahalanobis scores across nearby grid locations. This can make localization less sensitive to small spatial movements:

from tensor_md import NeighborhoodScoreLocationAwareTensorMahalanobisDetector

detector = NeighborhoodScoreLocationAwareTensorMahalanobisDetector(
    patches_per_image=normal.patches_per_image,
    grid_shape=(16, 16),
    score_neighbor_radius=1,
    score_neighbor_pooling="weighted_mean",
)
detector.fit(normal)
scores = detector.score(images)

The product of grid_shape must equal patches_per_image.

Score diagnostics

Diagnostics are optional and do not require anomaly labels. They compare the scores of the normal training images with the images being checked and can save arrays, distributions, heatmaps, or floating-point TIFF score maps:

artifacts = detector.fit_and_save_diagnostics(
    normal,
    images,
    "outputs/diagnostics",
    grid_shape=(16, 16),
    formats=("npy", "json", "distribution", "heatmaps", "tiff"),
)

Available formats are npy, csv, json, tiff, distribution, and heatmaps. Diagnostics are for inspection; benchmark metrics still require the benchmark's official ground-truth masks and evaluator.

Saving a fitted detector

detector.save("models/detector.pkl")

restored = LocationAwareTensorMahalanobisDetector.load(
    "models/detector.pkl"
)
scores = restored.score(images)

Model files use Python pickle and must only be loaded from a trusted source. Reuse the same image preprocessing and feature extractor when creating data for the restored detector.

MVTec AD layout

MVTec AD is optional. If the dataset is arranged as <root>/<category>/train/good and <root>/<category>/test/..., both splits can be loaded together:

from tensor_md import load_patch_datasets

config = PatchExtractionConfig(
    category="bottle",
    data_root="/path/to/mvtec",
    input_representation="cnn_features",
    cnn_backbone="ResNet50",
    cnn_layer_names=("conv3_block4_out", "conv4_block6_out"),
    cnn_dimensionality_reduction="pca",
    cnn_reduction_dimensions=128,
)
datasets = load_patch_datasets(config)

The package does not download or bundle MVTec AD.

License

MIT; see LICENSE.