CVPR 2026 SubspaceAD：单样本 97.1% AUROC，无需训练即达少样本异常检测新 SOTA

# 1. Create environment
conda create -n subspacead python=3.10
conda activate subspacead

# 2. Install dependencies and the package
pip install -r requirements.txt
pip install -e .

datasets/mvtec-ad/

python main.py \
    --dataset_name mvtec_ad \
    --dataset_path datasets/mvtec-ad \
    --categories bottle \
    --model_ckpt facebook/dinov2-with-registers-giant \
    --image_res 672 \
    --k_shot 5 \
    --aug_count 30 \
    --pca_ev 0.99 \
    --outdir results/debug_run

python predict_all_image.py  --normal_dir ./datasets/mvtec_ad/bottle/train/good   --query_image ./datasets/mvtec_ad/bottle/test/broken_large/   --model_ckpt facebook/dinov2-with-registers-giant  --image_res 320   --patch_size 224    --patch_overlap 0.2      --pca_ev 0.99     --outdir ./prediction_output

#!/usr/bin/env python3
"""
predict_single_image.py

Standalone single-image anomaly detection script for SubspaceAD.

Given:
  - a folder (or list) of normal reference images, and
  - one query image path,
fit a PCA subspace on the normal DINOv2 patch features and predict:
  - an image-level anomaly score
  - a pixel-level anomaly heatmap

Usage example:
    python predict_single_image.py \
        --normal_dir ./datasets/mvtec-ad/bottle/train/good \
        --query_image ./datasets/mvtec-ad/bottle/test/broken_large/000.png \
        --model_ckpt facebook/dinov2-with-registers-giant \
        --image_res 672 \
        --pca_ev 0.99 \
        --outdir ./prediction_output
"""

import argparse
import logging
import math
import os
import sys
from pathlib import Path
from types import SimpleNamespace

import cv2
import numpy as np
import torch
from PIL import Image

# ---------------------------------------------------------------------------
# Imports: try installed package first, otherwise use local src/ folder
# ---------------------------------------------------------------------------
try:
    from subspacead.config import parse_layer_indices, parse_grouped_layers
    from subspacead.core.extractor import FeatureExtractor
    from subspacead.core.pca import PCAModel
    from subspacead.core.patching import process_image_patched, get_patch_coords
    from subspacead.post_process.scoring import (
        aggregate_image_score,
        calculate_anomaly_scores,
        post_process_map,
    )
except ImportError:
    repo_root = Path(__file__).resolve().parent
    src_path = str(repo_root / "src")
    if src_path not in sys.path:
        sys.path.insert(0, src_path)
    from subspacead.config import parse_layer_indices, parse_grouped_layers
    from subspacead.core.extractor import FeatureExtractor
    from subspacead.core.pca import PCAModel
    from subspacead.core.patching import process_image_patched, get_patch_coords
    from subspacead.post_process.scoring import (
        aggregate_image_score,
        calculate_anomaly_scores,
        post_process_map,
    )


DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)


def _collect_image_paths(path_source, extensions=(".png", ".jpg", ".jpeg", ".bmp", ".tiff")):
    """Return a list of image paths from a directory or a single file path."""
    path_source = Path(path_source)
    if path_source.is_file():
        return [str(path_source)]
    if not path_source.is_dir():
        raise ValueError(f"Path must be an existing directory or file: {path_source}")

    paths = sorted(
        p for p in path_source.iterdir()
        if p.is_file() and p.suffix.lower() in extensions
    )
    if not paths:
        raise ValueError(f"No images found in {path_source} with extensions {extensions}")
    return [str(p) for p in paths]


def _build_config(args):
    """Convert argparse Namespace to the attribute-style config used by core code."""
    return SimpleNamespace(
        model_ckpt=args.model_ckpt,
        image_res=args.image_res,
        layers=args.layers,
        agg_method=args.agg_method,
        grouped_layers=args.grouped_layers,
        docrop=args.docrop,
        use_clahe=args.use_clahe,
        score_method=args.score_method,
        drop_k=args.drop_k,
        img_score_agg=args.img_score_agg,
        bg_mask_method=args.bg_mask_method,
        mask_threshold_method=args.mask_threshold_method,
        percentile_threshold=args.percentile_threshold,
        dino_saliency_layer=args.dino_saliency_layer,
        patch_size=args.patch_size,
        patch_overlap=args.patch_overlap,
        batch_size=args.batch_size,
    )


def _get_mask_background(saliency_map, cfg):
    """Return a boolean background mask from a DINO saliency map."""
    background_mask = np.zeros_like(saliency_map, dtype=bool)
    try:
        if cfg.mask_threshold_method == "percentile":
            threshold = np.percentile(saliency_map, cfg.percentile_threshold * 100)
            background_mask = saliency_map < threshold
        else:  # otsu
            norm_mask = cv2.normalize(
                saliency_map, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U
            )
            _, binary_mask = cv2.threshold(
                norm_mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
            )
            background_mask = binary_mask == 0
    except Exception as e:
        logging.warning(f"Background masking failed: {e}. Returning empty mask.")
    return background_mask


def fit_pca_model(normal_paths, extractor, cfg, pca_ev=0.99, pca_dim=None, whiten=False):
    """
    Fit a PCAModel on the normal reference images.

    Returns:
        pca_params (dict): parameters required by calculate_anomaly_scores
        feature_dim (int): feature dimensionality
        h_p (int), w_p (int): token grid size
    """
    logging.info(f"Fitting PCA on {len(normal_paths)} normal reference images...")

    # Determine feature dimension and token grid from one reference image
    temp_img = Image.open(normal_paths[0]).convert("RGB")
    layers = parse_layer_indices(cfg.layers)
    grouped_layers = parse_grouped_layers(cfg.grouped_layers) if cfg.agg_method == "group" else []

    temp_tokens, (h_p, w_p), _ = extractor.extract_tokens(
        [temp_img],
        cfg.image_res,
        layers,
        cfg.agg_method,
        grouped_layers,
        cfg.docrop,
        use_clahe=cfg.use_clahe,
        dino_saliency_layer=cfg.dino_saliency_layer,
    )
    feature_dim = temp_tokens.shape[-1]

    if cfg.patch_size:
        # ---------------- Patch-based fitting ----------------
        temp_patch = temp_img.crop((0, 0, cfg.patch_size, cfg.patch_size))
        temp_tokens, (h_p, w_p), _ = extractor.extract_tokens(
            [temp_patch],
            cfg.image_res,
            layers,
            cfg.agg_method,
            grouped_layers,
            cfg.docrop,
            use_clahe=cfg.use_clahe,
            dino_saliency_layer=cfg.dino_saliency_layer,
        )
        feature_dim = temp_tokens.shape[-1]
        tokens_per_patch = h_p * w_p

        total_patches = 0
        num_batches = 0
        for path in normal_paths:
            img = Image.open(path).convert("RGB")
            patch_coords = get_patch_coords(img.height, img.width, cfg.patch_size, cfg.patch_overlap)
            total_patches += len(patch_coords)
            num_batches += math.ceil(len(patch_coords) / cfg.batch_size)
        total_tokens = total_patches * tokens_per_patch

        def feature_generator():
            for path in normal_paths:
                pil_img = Image.open(path).convert("RGB")
                patch_coords = get_patch_coords(
                    pil_img.height, pil_img.width, cfg.patch_size, cfg.patch_overlap
                )
                for i in range(0, len(patch_coords), cfg.batch_size):
                    coord_batch = patch_coords[i : i + cfg.batch_size]
                    patch_batch = [pil_img.crop(c) for c in coord_batch]
                    tokens_batch, _, saliency_masks_batch = extractor.extract_tokens(
                        patch_batch,
                        cfg.image_res,
                        layers,
                        cfg.agg_method,
                        grouped_layers,
                        cfg.docrop,
                        use_clahe=cfg.use_clahe,
                        dino_saliency_layer=cfg.dino_saliency_layer,
                    )
                    tokens_flat = tokens_batch.reshape(-1, feature_dim)

                    if cfg.bg_mask_method == "dino_saliency":
                        masks_flat = saliency_masks_batch.reshape(-1)
                        foreground_tokens = tokens_flat[_get_mask_background(masks_flat, cfg) == False]
                        yield foreground_tokens if foreground_tokens.shape[0] > 0 else tokens_flat
                    else:
                        yield tokens_flat
    else:
        # ---------------- Full-image fitting ----------------
        total_train_images = len(normal_paths)
        total_tokens = total_train_images * h_p * w_p
        num_batches = math.ceil(total_train_images / cfg.batch_size)

        def feature_generator():
            for i in range(0, len(normal_paths), cfg.batch_size):
                path_batch = normal_paths[i : i + cfg.batch_size]
                img_batch = [Image.open(p).convert("RGB") for p in path_batch]
                tokens_batch, _, saliency_masks_batch = extractor.extract_tokens(
                    img_batch,
                    cfg.image_res,
                    layers,
                    cfg.agg_method,
                    grouped_layers,
                    cfg.docrop,
                    use_clahe=cfg.use_clahe,
                    dino_saliency_layer=cfg.dino_saliency_layer,
                )
                tokens_flat = tokens_batch.reshape(-1, feature_dim)

                if cfg.bg_mask_method == "dino_saliency":
                    masks_flat = saliency_masks_batch.reshape(-1)
                    foreground_tokens = tokens_flat[_get_mask_background(masks_flat, cfg) == False]
                    yield foreground_tokens if foreground_tokens.shape[0] > 0 else tokens_flat
                else:
                    yield tokens_flat

    pca_model = PCAModel(k=pca_dim, ev=pca_ev, whiten=whiten)
    pca_params = pca_model.fit(feature_generator, feature_dim, total_tokens, num_batches)
    logging.info(f"PCA fitted: kept k={pca_params['k']} components.")
    return pca_params, feature_dim, h_p, w_p



def save_visualization(pil_img, anomaly_map, out_path, colormap=cv2.COLORMAP_JET):
    """Save original image, anomaly heatmap, and an overlay."""
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    img_np = np.array(pil_img)
    anomaly_norm = (anomaly_map - anomaly_map.min()) / (anomaly_map.max() - anomaly_map.min() + 1e-8)
    anomaly_u8 = (anomaly_norm * 255).astype(np.uint8)

    heatmap = cv2.applyColorMap(anomaly_u8, colormap)
    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)

    # Resize heatmap to match image if needed
    if img_np.shape[:2] != heatmap.shape[:2]:
        heatmap = cv2.resize(heatmap, (img_np.shape[1], img_np.shape[0]), interpolation=cv2.INTER_LINEAR)

    alpha = 0.5
    overlay = (img_np.astype(np.float32) * (1 - alpha) + heatmap.astype(np.float32) * alpha).astype(np.uint8)

    # Stack horizontally
    vis = np.concatenate([img_np, heatmap, overlay], axis=1)
    Image.fromarray(vis).save(out_path)
    logging.info(f"Visualization saved to {out_path}")


def main():
    parser = argparse.ArgumentParser(description="SubspaceAD single-image prediction")

    # Data
    parser.add_argument("--normal_dir", type=str, required=True,
                        help="Directory (or single image path) containing normal reference images.")
    parser.add_argument("--query_image", type=str, required=True,
                        help="Path to a query image or a directory containing query images.")
    parser.add_argument("--outdir", type=str, default="./prediction_output",
                        help="Directory where the output heatmap will be saved.")
    parser.add_argument("--save_heatmap", action="store_true", default=True,
                        help="Save a visualization of the anomaly heatmap.")

    # Model
    parser.add_argument("--model_ckpt", type=str,
                        default="facebook/dinov2-with-registers-large",
                        help="HuggingFace DINOv2 checkpoint.")
    parser.add_argument("--image_res", type=int, default=256,
                        help="Input resolution for the model.")
    parser.add_argument("--batch_size", type=int, default=4,
                        help="Batch size for feature extraction.")

    # Feature aggregation
    parser.add_argument("--layers", type=str, default="-12,-13,-14,-15,-16,-17,-18",
                        help="Comma-separated layer indices to aggregate.")
    parser.add_argument("--agg_method", type=str, default="mean",
                        choices=["concat", "mean", "group"],
                        help="How to aggregate selected layers.")
    parser.add_argument("--grouped_layers", type=str, default=None,
                        help="Layer groups for 'group' agg. Format: '-1,-2:-3,-4'.")
    parser.add_argument("--docrop", action="store_true",
                        help="Apply center crop during preprocessing.")
    parser.add_argument("--use_clahe", action="store_true",
                        help="Apply CLAHE contrast enhancement.")

    # PCA
    parser.add_argument("--pca_ev", type=float, default=0.99,
                        help="Explained-variance ratio for PCA. Ignored if --pca_dim is set.")
    parser.add_argument("--pca_dim", type=int, default=None,
                        help="Fixed number of PCA components.")
    parser.add_argument("--whiten", action="store_true",
                        help="Whitening in PCA.")

    # Scoring
    parser.add_argument("--score_method", type=str, default="reconstruction",
                        choices=["reconstruction", "mahalanobis", "cosine", "euclidean"],
                        help="Anomaly scoring method.")
    parser.add_argument("--drop_k", type=int, default=0,
                        help="Drop first k principal components during scoring.")
    parser.add_argument("--img_score_agg", type=str, default="mtop1p",
                        choices=["max", "mean", "p99", "mtop5", "mtop1p"],
                        help="Aggregation method for the image-level score.")

    # Background masking
    parser.add_argument("--bg_mask_method", type=str, default=None,
                        choices=[None, "dino_saliency"],
                        help="Optional background removal. 'pca_normality' is omitted for simplicity.")
    parser.add_argument("--mask_threshold_method", type=str, default="percentile",
                        choices=["percentile", "otsu"],
                        help="Binarization method for the saliency mask.")
    parser.add_argument("--percentile_threshold", type=float, default=0.15,
                        help="Percentile threshold when using percentile masking.")
    parser.add_argument("--dino_saliency_layer", type=int, default=6,
                        help="Transformer layer index for DINO saliency mask.")

    # Patching (optional)
    parser.add_argument("--patch_size", type=int, default=None,
                        help="If set, process images in overlapping patches of this size.")
    parser.add_argument("--patch_overlap", type=float, default=0.0,
                        help="Overlap ratio between patches (0.0-1.0).")

    args = parser.parse_args()
    cfg = _build_config(args)

    logging.info(f"Using device: {DEVICE}")

    # Load data
    normal_paths = _collect_image_paths(args.normal_dir)
    logging.info(f"Found {len(normal_paths)} normal reference images.")

    query_paths = _collect_image_paths(args.query_image)
    logging.info(f"Found {len(query_paths)} query image(s).")

    # Build feature extractor
    extractor = FeatureExtractor(args.model_ckpt)

    # Fit PCA model
    pca_params, feature_dim, h_p, w_p = fit_pca_model(
        normal_paths,
        extractor,
        cfg,
        pca_ev=args.pca_ev,
        pca_dim=args.pca_dim,
        whiten=args.whiten,
    )

    # Predict all query images
    results = []
    for query_path in query_paths:
        logging.info(f"Predicting anomaly for {query_path}...")
        anomaly_map, img_score, pil_img = predict_single_image(
            query_path, extractor, pca_params, cfg, feature_dim, h_p, w_p
        )

        print(f"\n[{Path(query_path).name}] Image-level anomaly score ({args.img_score_agg}): {img_score:.6f}")
        results.append((Path(query_path).name, img_score))

        query_stem = Path(query_path).stem

        # Save heatmap visualization
        if args.save_heatmap:
            out_path = Path(args.outdir) / f"{query_stem}_anomaly.png"
            save_visualization(pil_img, anomaly_map, out_path)

        # Save raw anomaly map as .npy
        raw_out_path = Path(args.outdir) / f"{query_stem}_anomaly_map.npy"
        np.save(raw_out_path, anomaly_map)
        logging.info(f"Raw anomaly map saved to {raw_out_path}")

    if len(results) > 1:
        print("\n" + "=" * 60)
        print("Summary of image-level anomaly scores:")
        for name, score in results:
            print(f"  {name:<40s} {score:.6f}")
        print("=" * 60 + "\n")


if __name__ == "__main__":
    main()