DSR -- A dual subspace re-projection network for surface anomaly detection

In plain words

Automated inspection cameras look for scratches, dents and other flaws on factory parts. The flaws often look a lot like a good part. Many methods therefore learn only from good parts, then flag whatever they cannot rebuild correctly. That idea works poorly when a flaw is only slightly different from normal.

DSR first compresses each image into a small vocabulary of patterns, called a codebook. It then rebuilds the image with two decoders. One is a general decoder that knows how natural images look. The other only knows how this particular object looks when it is fine. Where the second decoder cannot reproduce what it sees, the part is suspicious.

To teach the system what subtle defects look like, the method does not paste fake scars onto photos. It changes some of the patterns in the compressed representation to other patterns from the same vocabulary that are close to the originals. The result looks like a plausible defect on the real part. The system then learns to restore normal appearance, and the difference between the input and the restored version shows where the defect is.

TL;DR

An unsupervised surface-defect detector. In the unsupervised setting it trains on defect-free images of one object class, with no defect samples and no external anomaly dataset. Use it when defects look close to normal appearance. The authors report state-of-the-art results on KSDD2 and MVTec, with official code under Apache-2.0. Stage 1 needs a pretrained codebook model, and its training code is not in the repository.

Problem

Reconstruction-based detectors assume that anomalies are reconstructed badly, and that a pixel or feature distance will expose them. Both assumptions fail when the reconstruction network generalises well or when the defect is subtle §1 §2. Discriminative methods that synthesise anomalies usually paste image-level patches taken from external datasets. Their results drop for near-in-distribution anomalies, which are hard to simulate realistically §1 §2.

The paper targets two gaps: the reliance on auxiliary anomaly datasets, and weak detection of near-in-distribution anomalies §1.

Method

DSR has three stages §3.7:

1. Quantized encoder and general decoder (trained on ImageNet, then frozen) §3.7. The encoder maps the image to features at two resolutions. Each feature vector is replaced by its nearest codebook entry §3.1 eq. 1. The codebooks are K_{hi} and K_{lo}. The hi-resolution features are 4× smaller than the input, and the lo-resolution features are 8× smaller §3.1. The general decoder upsamples Q_{lo}, concatenates it with Q_{hi}, applies residual blocks and transposed convolutions, and outputs I_{gen} §3.2.
2. Object-specific decoder and anomaly detector (trained per object class). Two subspace restriction modules, which are encoder-decoder networks, map each quantized level to continuous features \tilde F. A vector-quantization projection with the same codebooks then gives \tilde Q. A convolutional decoder turns \tilde Q into the restored image I_{spc} §3.3 fig. 4. The two reconstructions are concatenated and passed to a U-Net that outputs the anomaly map M §3.4. The image-level score is the global maximum of the map after a 21 \times 21 average filter §3.4.
3. Upsampling module (trained last). A U-Net-like network refines the map to full resolution §3.5. The paper trains it on images with copy-pasted smudges, using a focal loss against the known pixel mask §3.7.

The anomaly generator works on the quantized features of an anomaly-free image §3.6. A Perlin noise mask M_{gt} marks where to create an anomaly §3.6. Inside the mask, each feature is replaced by a codebook entry chosen from its near neighbours. The bound \lambda_s excludes the 5% most similar codebook entries, to avoid turning normal features into false alarms §3.6.

Training loss in stage 1 (ImageNet) eq. 2:

\mathcal{L}_{st1} = L_2(I, I_{gen}) + L_2(sg[F], Q) + \lambda_1 L_2(F, sg[Q]) eq. 2

Training loss in stage 2 (detector and object-specific decoder, on generated anomalies) eq. 3:

\mathcal{L}_{st2} = \mathcal{L}_{foc}(M_{gt}, M) + \lambda_2 L_2(\tilde F, Q) + \lambda_3 L_2(I, I_{spc}) eq. 3

The quantization that produces Q from F is eq. 1:

Q_{i,j} = \arg\min_{e_l \in K} \lVert F_{i,j} - e_l \rVert eq. 1

The anomaly generator, as the paper describes it §3.6:

# Feature-space anomaly generation  [§3.6]
# Input:  features (B, h, w, C) = Q; codebook (N, C) = K; mask (B, h, w)
# Output: anomalous features (B, h, w, C) = Q_a
# lambda_s is 0.05 in all experiments  [§3.6]

def generate_anomalies(features, codebook, mask, lambda_s):
    # N: number of codebook vectors  [§3.6]
    num_codes = codebook.shape[0]
    anomalous = features.copy()
    for image in range(batch_size):
        # One similarity bound per image: n ~ U[lambda_s N, N]  [§3.6]
        bound = uniform(lambda_s * num_codes, num_codes)
        for position in where(mask[image] == 1):
            # Rank codebook vectors by distance to this feature  [§3.6]
            ranked = argsort(distance(features[image, position], codebook))
            # Draw rank k ~ U[lambda_s N, n] and copy that code  [§3.6]
            rank = uniform_int(lambda_s * num_codes, bound)
            anomalous[image, position] = codebook[ranked[rank]]
    return anomalous

Inference, following the released test code, which is what the reported numbers come from test_dsr.py:

# Inference as the released code does it  [§3.1]-[§3.5]
# Input:  image (B, 3, 224, 224) after resize to 256 and centre crop
# Output: image score (B,), pixel map (B, 224, 224)

def detect(image):
    # Quantize the image at two resolutions  [§3.1] [eq. 1]
    quant_hi, quant_lo = quantized_encoder(image)
    # Restore normal features, then rebuild the object  [§3.3]
    restored_hi, restored_lo = subspace_restriction(quant_hi, quant_lo)
    rebuilt = object_decoder(restored_hi, restored_lo)
    # Rebuild with the general decoder  [§3.2]
    general = general_decoder(quant_hi, quant_lo)
    # Anomaly map from both rebuilds  [§3.4]
    anomaly_map = detection_unet(rebuilt, general)
    # Refine to pixel resolution  [§3.5]
    pixel_map = upsampling_unet(rebuilt, general, anomaly_map)
    # Image score: 21x21 mean filter, then global max  [§3.4]
    smoothed = average_filter(anomaly_map, size=21)
    return max(smoothed), pixel_map

Hyperparameters

Results

Author-reported. Setup: KSDD2 with 2085 anomaly-free and 246 anomalous training images. The test set has 894 anomaly-free and 110 anomalous images §4.2. MVTec AD has 15 classes; image-level AUROC and pixel-level AP are reported §4.3. Hardware is not stated. The reported speed is 58 FPS §5.

KSDD2, unsupervised tab. 1: AP for detection (AP_{det}) and localization (AP_{loc}), in %.

KSDD2, with N anomalous training images that have masks tab. 2. The authors list AP_{det} and AP_{loc} in %. The row labelled KSDD2 baseline is the supervised method of the dataset paper, bovzivc2021mixed.

MVTec AD, image-level AUROC (%) tab. 3:

MVTec AD, pixel-level AP (%) tab. 4:

Similarity bound on MVTec, image-level AUROC (%) tab. 5: \lambda_s = 0.01: 97.5; 0.02: 98.0; 0.05: 98.2; 0.1: 97.7; 0.2: 97.1; 0.5: 95.5.

Ablation on MVTec tab. 6:

Qualitative results: the KSDD2 examples fig. 6, the KSDD2 predictions fig. 7, the MVTec outputs fig. 8, and a comparison of reconstructions against DRAEM fig. 9. The paper describes these in words only.

Paper vs code

Code: official repository at commit 672bfb81434fd2a6c5ef00db858cef8834c54f28 (Apache-2.0) LICENSE. It has the training and test scripts and the MVTec classes. It has no KSDD2 code, no ImageNet stage 1 training code, and no timing code. The pretrained stage 1 weights come from a Google Drive link in the README README.md.

• Similarity bound. Paper: n \sim U[\lambda_s N, N] with \lambda_s = 0.05 §3.6. Code: the sampled strength is uniform on [0.2, 1], so the bound starts at 0.2 N train_dsr.py train_dsr.py. The 5% exclusion, and the uniform draw of rank k in [0.05N, n), do match the paper train_dsr.py.
• Anomaly levels. Paper: features are replaced at both levels §3.6. Code: in 50% of samples both levels are replaced, in 25% only the hi level, and in 25% only the lo level train_dsr.py.
• Loss norm. Paper: L_2 is the Euclidean distance §3.7. Code: mean squared error for the feature and image terms train_dsr.py. The feature term is summed over both levels, each with weight 1 train_dsr.py.
• VQ term. The subspace restriction modules return a VQ commitment loss, but the stage 2 total does not include it train_dsr.py train_dsr.py. The paper's eq. 3 has no such term, so this is consistent with the paper.
• Map resolution. Paper: the anomaly map is at feature resolution §3 §3.5. Code: the anomaly U-Net takes two RGB images (6 channels) and has three pooling steps followed by three upsampling steps, so its map is at the full input resolution dsr_model.py dsr_model.py. The stage 2 ground truth is max-pooled to the feature grid and then nearest-upsampled back to full size, so it is blocky train_dsr.py.
• Upsampling input. Paper: the input is the image I concatenated with the bilinearly upsampled mask §3.5. Code: the input is the two reconstructions and the full-resolution map. No bilinear step is applied train_dsr.py dsr_model.py.
• Subspace restriction depth. Paper: three downsampling and three upsampling blocks §3.3. Code: two of each dsr_model.py dsr_model.py. The image reconstruction network does match the paper's two downsampling and four transposed convolutions dsr_model.py.
• Stage 3 iterations. Paper: 20,000 iterations §4.1. Code: epochs//2 epochs of 1000 iterations, which is 50,000 with --epochs 100 train_dsr.py data_loader.py.
• Stage 2 iterations. Consistent. 100 epochs of 1000 iterations is 100,000, with the decay at epoch 80 train_dsr.py train_dsr.py.
• Stage 3 smudges. Paper: "copy-pasted smudges" §3.7. Code: solid colour patches under a Perlin mask, blended with the image at a random strength up to 0.8, with half the samples left anomaly-free data_loader.py. The augmenter list (gamma, hue, solarize, equalize, affine, and others) is defined but never used data_loader.py.
• Checkpoint names. Training saves the upsampling module to dsr_<lr>_<epochs>_bs<bs>_anom_par<p>_ + _upsample.pckl with no object name, so each class overwrites the previous one train_dsr.py train_dsr.py. Testing loads <run>_upsample_module_<obj>.pckl test_dsr.py. These names do not match. The released checkpoint archive is not in the repository, so the name it uses could not be checked.
• Silent load. The image reconstruction network is loaded with strict=False, so missing or mismatched keys are ignored test_dsr.py.
• Evaluation resolution. Test images are resized to 256 and centre-cropped to 224, and the pixel metrics use the 224 crop test_dsr.py test_dsr.py test_dsr.py. The paper does not state the input size.
• Stage 1 and speed. No stage 1 training code. The 58 FPS figure has no timing code and no hardware in the paper §5.

Checked and consistent: the codebook size 4096 and dimension 128, the \lambda_2 = 1 and \lambda_3 = 10 weights, the 5% exclusion in the sampler, and the stage 2 schedule (100,000 iterations, decay at 80,000).

Pitfalls

• Abstract percentages are relative. The abstract says "10% AP" and "35% AP". From Table 1, DSR's AP_{det} is 87.2 against 79.3 for MAD, which is +7.9 points, about +10% relative. Its AP_{loc} is 61.4 against 45.3 for PaDim, which is +16.1 points, about +35% relative. These are our calculations from Table 1 tab. 1; the paper's own text gives 7.9 points §4.2.1.
• Upsampling checkpoint. The naming mismatch above means the released test script may not run from a freshly trained model. Rename or save per class before evaluating.
• Unstated details. The paper never gives the input size, the optimizer, the Perlin threshold or scales, the focal-loss parameters, or the weight initialisation. The code's values are in the hyperparameter table and should be treated as the working configuration, not as paper-stated values.
• Image score borders. The 21×21 average uses zero padding of 10 pixels, so averages near the image border are diluted by zeros test_dsr.py. The effect on border defects was not measured.
• Sampling bound is sensitive. The paper reports robustness to \lambda_s across 0.01 to 0.2, with a drop at 0.5 tab. 5. The code uses 0.2 as the lower bound of the sampled strength, so the effective candidate pool is smaller and differs from the paper's stated range.
• Anomaly-free samples. About half of the stage 2 samples have an all-zero Perlin mask, so they contain no generated anomaly data_loader.py. The paper does not mention this, and it lowers the number of anomaly examples the detector sees.
• Frozen codebook. The VQ codebook is updated by EMA only in training mode. The model is put in eval mode before stage 2 and stage 3, so the codebook stays fixed train_dsr.py discrete_model.py. The straight-through estimator passes gradients through the quantizer discrete_model.py.
• One model per class. Stage 2 and 3 run per object class. The test runs a separate model per class, and the paper's comparisons use one model per class. There is no multi-class model in the repository.
• Reproduction limits. KSDD2 numbers cannot be reproduced from the repository, since no KSDD2 code is present. The MVTec numbers need the pretrained release, which is hosted on Google Drive and was not downloaded or checked.

Lineage

• Builds on: VQ-VAE (Neural Discrete Representation Learning, NeurIPS) for the quantized encoder and the codebook training §3.1 §3.7. VQ-VAE-2 (Generating Diverse High-Fidelity Images with VQ-VAE-2, NeurIPS) for two-level quantization §3.1. Perlin noise (An image synthesizer, SIGGRAPH 1985) for the anomaly masks §3.6. None of these has a card on the wiki.
• Compared against: DRAEM §4.2.1 §4.3.1, which also synthesises anomalies but in image space; MAD, PaDim, US (uninformed students), CutPaste, RIAD §4.2.1 §4.3.1; and the supervised KSDD2 method (bovzivc2021mixed) §4.2.2.
• Superseded by: not known from wiki content.