PSAD: Few Shot Part Segmentation Reveals Compositional Logic for Industrial Anomaly Detection

This repository includes a pytorch implementation of the paper "Few Shot Part Segmentation Reveals Compositional Logic for Industrial Anomaly Detection" accepted in AAAI2024.

Abstract

Logical anomalies (LA) refer to data violating underlying logical constraints e.g., the quantity, arrangement, or composition of components within an image. Detecting accurately such anomalies requires models to reason about various component types through segmentation. However, curation of pixel-level annotations for semantic segmentation is both time-consuming and expensive. Although there are some prior few-shot or unsupervised co-part segmentation algorithms, they often fail on images with industrial object. These images have components with similar textures and shapes, and a precise differentiation proves challenging. In this study, we introduce a novel component segmentation model for LA detection that leverages a few labeled samples and unlabeled images sharing logical constraints. To ensure consistent segmentation across unlabeled images, we employ a histogram matching loss in conjunction with an entropy loss. As segmentation predictions play a crucial role, we propose to enhance both local and global sample validity detection by capturing key aspects from visual semantics via three memory banks: class histograms, component composition embeddings and patch-level representations. For effective LA detection, we propose an adaptive scaling strategy to standardize anomaly scores from different memory banks in inference. Extensive experiments on the public benchmark MVTec LOCO AD reveal our method achieves 98.1% AUROC in LA detection vs. 89.6% from competing methods.

Few Shot Part Segmentation

PSAD

Dataset

• We use MVTec LOCO AD Dataset for evaluation. For details of this dataset, please see here.

Training

To implement our proposed Part Segmentation-based Anomaly Detection (PSAD), follow procedure shown as below. If you want to skip the preprocessing and segmentation process 1~4, download preprocessed images, segmentation maps and use them as inputs.

1. Annotate few labeled images per category.

We used an annotation tool to make annotations for few images. In most cases, we used 5 labeled images. But when there are multiples product types (e.g., juicie bottle and splicing connectors), we used 1 labeled images per type, i.e., total 3 images.

2. Preprocess images.

Convert label format, crop and resize the image. Please check the code in preprocess directory.

3. Train a few shot part segmentation model using few labeld and numerous unlabeled images. Using different levels of features of pretrained encoder can lead to different segmentation results.

   CUDA_VISIBLE_DEVICES=[GPU_ID] python finetune_cnn_coord.py 
   --n_shot 5 
   --num_epochs 100 
   --obj_name [OBJ_NAME] 
   --snapshot_dir [DIR]

4. Train a segmentation model again using predicted pseudo label.

   run_unet.sh [GPU_ID] [DIR]

5. Save predictions of PatchCore. Please check patchcore directory. Or you can use the scores already obtained using PatchCore. It includes anomaly scores of all train and test data predicted using PatchCore without the coreset sampling, and the scores of train data by treating each as a test data for 'adaptive scaling'. Each .pth file (Ex. dict = train/good/000.pt) has a dictionary with 5 keys ['distance', 'anomaly_map_interpolate', 'anomaly_maps', 'anomaly_scores', 'label', 'name']. 'anomaly_scores' is the score for the adaptive scaling.

6. Train and test PSAD based on segmentation results. You can choose memory type like "hcp". h,c, and p denote histogram, composition, and patch memory banks.

   ./run_ad.sh [GPU_ID] [DIR] "hcp" "max"

Testing

Our proposed PSAD and other comparison methods are evaluated using image AUROC scores. As there is no ground truth for segmentation task, we indirectly evaluated the segmentation performance using anomaly detection performance. In general, accurate segmentation correlates with anomaly detection performance.

Results

• Comparison with state-of-the-art methods

• PSAD performance depending on segmentation model

• Qualitative evaluation on FSS models

• Ablation study on multiple memory banks and adaptive scaling (AS)

• PSAD performance using less normal images (same segmentation model is used.)

• Visualizing hitrograms from different memory banks

Acknowledgments

Our work was inspired by many previous works related to industrial anomaly detection and few shot segmentation including PatchCore, RePRI. Thanks to their inspiring works.