AnomalyDINO

repository·main·Indexed 42 days ago

https://github.com/dammsi/anomalydino

A framework for boosting patch-based few-shot and batched zero-shot anomaly detection using DINOv2 features. It supports standard datasets like MVTec-AD and VisA, providing tools for feature extraction via DINOv2 and ViT backbones, background masking, and similarity search using FAISS or scikit-learn.

Tokens
6.5K
Snippets
18
Records
24
Agent score
54%

What's inside AnomalyDINO

  1. Perform Few-shot Anomaly Detection

    main

    Use the run_anomalydino.py script to evaluate few-shot anomaly detection. You can specify the number of shots, repetitions (seeds), and preprocessing method.

    Preprocessing Options:

    • agnostic: (Default) Applies masking where PCA-based masking works well and augments reference samples by rotations.
    • informed: Use this for faster inspection.

    Preprocessing logic can be customized in src/utils.py within the get_dataset_info function.

    Key CLI Arguments:

    • --dataset: The dataset name (e.g., MVTec or VisA).
    • --shots: A list of shot counts (e.g., 1 2 4 8 16).
    • --num_seeds: Number of repetitions.
    • --preprocess: Preprocessing strategy (agnostic or informed).
    • --data_root: Path to the dataset root.
    • --eval_segm: Activate this flag to enable segmentation metric evaluation.

    Results are saved in the respective results directory as metrics_seed={seed}.json for each seed. The script automatically generates example plots and anomaly maps.

    # Reproduce default settings for MVTec
    python run_anomalydino.py --dataset MVTec --shots 1 2 4 8 16 --num_seeds 3 --preprocess agnostic --data_root data/mvtec_anomaly_detection
    
    # Faster inspection for VisA
    python run_anomalydino.py --dataset VisA --shots 1 --num_seeds 1 --preprocess informed --data_root data/VisA_pytorch/1cls/
  2. Prepare datasets for AnomalyDINO

    main

    AnomalyDINO supports MVTec-AD and VisA datasets. You must download them from their official sources and organize them according to the following structure.

    Default Data Roots:

    • MVTec-AD: data/mvtec_anomaly_detection
    • VisA: data/VisA_pytorch/1cls/

    Required Directory Structure:

    your_data_root
    ├── object1
    │   ├── ground_truth        # anomaly annotations per anomaly type (optional)
    │   │   ├── anomaly_type1
    │   │   └── ...
    │   ├── test                # test images per anomaly type & 'good' (if applicable)
    │   │   ├── anomaly_type1    
    │   │   └── good
    │   └── train               # train/reference images (without anomalies)
    │       └── good
    ├── object2
    │   └── ...

    Note: If no 'good' test set is available, only inference is performed and segmentation/detection metrics cannot be evaluated.

  3. Perform Batched-Zero-Shot Anomaly Detection

    main

    To reproduce results for the batched zero-shot scenario, use the run_anomalydino_batched.py script.

    # Batched zero-shot for MVTec
    python run_anomalydino_batched.py --dataset MVTec --data_root data/mvtec_anomaly_detection
    
    # Batched zero-shot for VisA
    python run_anomalydino_batched.py --dataset VisA --data_root data/VisA_pytorch/1cls/
  4. Install AnomalyDINO and its dependencies

    main

    To set up AnomalyDINO, create a virtual environment and install the required dependencies using pip.

    Note on FAISS GPU acceleration: If you require GPU-accelerated similarity search via faiss, it is recommended to use a Conda environment instead of pip. If you prefer to perform similarity search on the CPU, use the --faiss_on_cpu flag during execution.

    python -m venv .venvAnomalyDINO
    source .venvAnomalyDINO/bin/activate
    pip install -r requirements.txt
  5. How few-shot and full-shot settings work

    main

    The run_anomaly_detection function handles different evaluation modes via the n_ref_samples parameter:

    • Full-shot Setting: Set n_ref_samples = -1. The function will use all images found in {data_root}/{object_name}/train/good/ to build the memory bank.
    • Few-shot Setting: Set n_ref_samples to a positive integer (e.g., 5). The function selects a deterministic subset of images from the training folder based on the seed parameter. The selection logic is: sorted(os.listdir(img_ref_folder))[seed*n_ref_samples:(seed + 1)*n_ref_samples].
  6. Initialize AnomalyDINO model

    main

    To use AnomalyDINO, you must first load a DINOv2 backbone using get_model. You need to specify the model name (e.g., 'dinov2_vits14'), the device (e.g., 'cuda'), and optionally the smaller_edge_size to control the input resolution.

    from src.backbones import get_model
    
    # Load DINOv2 model
    model = get_model('dinov2_vits14', 'cuda', smaller_edge_size=672)
  7. Perform Few-shot Anomaly Detection

    main

    Few-shot anomaly detection with AnomalyDINO follows these steps:

    1. Extract nominal features: Load reference (good) samples. You can optionally augment them (e.g., using augment_image for rotation). Note that the project recommends not applying masking to reference features to reduce false positives.
    2. Extract test features: Load the test image and extract its patch features.
    3. Masking (Optional): Use model.compute_background_mask_from_image to filter out irrelevant patches (background) from the test image based on a threshold and masking_type.
    4. Compute distances: Calculate the distance between test patches and their closest nominal patch using a 1-Nearest Neighbour approach.
    5. Score: The anomaly score can be derived from the patch distances (e.g., the mean of the top 1% highest distances).
  8. Generate batch sample plots with create_sample_plots

    main

    Use create_sample_plots to automatically generate a large overview image containing grids of anomaly maps. For each object, it creates a plot showing:

    1. The first five samples of every detected anomaly type.
    2. The first five 'good' samples for baseline comparison.

    Parameters:

    • experiment_path (str): The directory where the resulting .png plot will be saved.
    • anomaly_maps_dir (str): The directory containing the .npy distance files.
    • seed (int/str): A seed value used in the output filename.
    • dataset (str): The dataset name (used to fetch object/anomaly info via get_dataset_info).
    • data_root (str): The root directory of the dataset images.
    create_sample_plots(
        experiment_path="./plots",
        anomaly_maps_dir="./exp_results",
        seed=42,
        dataset="MVTec",
        data_root="./data/MVTec"
    )
  9. Perform classification evaluation with `eval_classification`

    main

    Use eval_classification to calculate image-level metrics (AUROC, AP, and F1) by aggregating pixel-wise anomaly scores into single image scores.

    Arguments:

    • gt_filenames: List of paths to ground truth mask files.
    • prediction_filenames: List of paths to predicted anomaly maps (expected as .npy files).
    • aggregation_statistics: Method to reduce anomaly maps to image scores. Supported values:
      • "meantop1p": Mean of the top 1% of distances.
      • "max_anomaly_map": Maximum value in the reconstructed anomaly map.
      • "max_patch_distance": Maximum distance in the flattened array.
    from src.post_eval import eval_classification
    
    # gt_filenames and prediction_filenames should be lists of strings
    auroc_clf, ap_clf, f1_clf = eval_classification(
        gt_filenames=my_gt_list,
        prediction_filenames=my_pred_list,
        aggregation_statistics="meantop1p"
    )
  10. Parse MVTec AD dataset files with `parse_dataset_files`

    main

    This utility parses the directory structure of the MVTec AD dataset to align test images with their corresponding ground truth masks and predicted anomaly maps.

    Arguments:

    • object_name: The specific object category (e.g., "bottle", "cable").
    • dataset_base_dir: The root directory containing the MVTec dataset.
    • anomaly_maps_dir: The directory where your generated anomaly maps are stored.
    • dataset: (Default "MVTec") The dataset type. If set to "MVTec", it expects .png extensions and specific mask naming conventions (_mask.png).
    from src.post_eval import parse_dataset_files
    
    gt_files, pred_files = parse_dataset_files(
        object_name="bottle",
        dataset_base_dir="/data/MVTec",
        anomaly_maps_dir="/data/predictions"
    )
  11. Extract features and compute background masks

    main

    The model provides low-level methods for feature extraction and mask generation:

    • model.prepare_image(image): Prepares an image for the model, returning the image_tensor and the grid_size.
    • model.extract_features(image_tensor): Extracts patch-level features from the prepared tensor.
    • model.compute_background_mask(features, grid_size, threshold, masking_type): Computes a mask to filter out background patches. threshold controls sensitivity, and masking_type determines the masking strategy.
    • model.get_embedding_visualization(features, grid_size): Generates a PCA-based visualization of the feature embeddings.
  12. Run anomaly detection with run_anomaly_detection()

    main

    The run_anomaly_detection function is the primary entrypoint for evaluating the anomaly detection performance of a specific object or product. It supports both full-shot (using all available training data) and few-shot (using a subset of training data) settings.

    Key Workflow:

    1. Memory Bank Construction: It extracts features from reference (training) images to build a searchable index using FAISS.
    2. Inference: It processes test images, extracts patch-level features, and computes distances to the nearest neighbors in the memory bank.
    3. Masking: If masking is enabled, it uses the model to estimate a foreground mask to discard background patches, focusing detection on the object.
    4. Output: It returns anomaly scores (based on the mean of the top 1% of patch distances), memory bank setup time, and per-image inference times.

    Data Structure Requirements:

    • data_root must follow a specific directory structure:
      • {data_root}/{object_name}/train/good/ (Reference images)
      • {data_root}/{object_name}/test/{anomaly_type}/ (Test images)
    • object_anomalies should be a dictionary mapping object_name to a list of anomaly type strings.
    # Example usage sketch
    anomaly_scores, setup_time, inf_times = run_anomaly_detection(
        model=my_dino_model,
        object_name='product_a',
        data_root='/path/to/dataset',
        n_ref_samples=5,  # Few-shot: use 5 samples
        object_anomalies={'product_a': ['scratch', 'dent']},
        plots_dir='./results',
        masking=True,
        knn_metric='L2_normalized'
    )