Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,15 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/bottle_defect.png filter=lfs diff=lfs merge=lfs -text
+examples/bottle_good.png filter=lfs diff=lfs merge=lfs -text
+examples/carpet_defect.png filter=lfs diff=lfs merge=lfs -text
+examples/carpet_good.png filter=lfs diff=lfs merge=lfs -text
+examples/hazelnut_defect.png filter=lfs diff=lfs merge=lfs -text
+examples/hazelnut_good.png filter=lfs diff=lfs merge=lfs -text
+examples/leather_defect.png filter=lfs diff=lfs merge=lfs -text
+examples/leather_good.png filter=lfs diff=lfs merge=lfs -text
+examples/tile_defect.png filter=lfs diff=lfs merge=lfs -text
+examples/tile_good.png filter=lfs diff=lfs merge=lfs -text
+examples/wood_defect.png filter=lfs diff=lfs merge=lfs -text
+examples/wood_good.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,28 @@
 ---
-title: Patchcore Anomaly Detection
-emoji: 📊
-colorFrom: green
-colorTo: purple
+title: Industrial Anomaly Detection — PatchCore Demo
+emoji: 🔍
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
-sdk_version: 6.17.3
-python_version: '3.13'
+sdk_version: "6.17.3"
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Industrial Anomaly Detection with PatchCore
+
+Detects surface defects in industrial images using **PatchCore** (Roth et al., CVPR 2022) —
+a training-free method that builds a memory bank of normal patch features and flags anomalies
+as regions far from that bank.
+
+**Supported categories (15):** bottle, cable, capsule, carpet, grid, hazelnut, leather,
+metal_nut, pill, screw, tile, toothbrush, transistor, wood, zipper
+
+**How it works:** Upload any product image → nearest-neighbour distance from patch features
+to the memory bank → pixel-level anomaly heatmap + image-level score and pass/fail verdict.
+
+**Backbone:** WideResNet-101-2 (ImageNet pretrained, layers 2 + 3 concatenated → 1536-dim descriptors)
+
+**Note:** Thresholds are not stored in `metrics.json` (only AUROC/PRO metrics are saved).
+The verdict currently uses a default threshold of 0.5 — calibrate per category from the
+score distributions in `results/{category}/score_distribution.png`.

app.py

@@ -0,0 +1,186 @@
+"""
+PatchCore Anomaly Detection — Gradio demo.
+
+Loads per-category models from results/{category}/patchcore_model.pt and runs
+CPU inference. Thresholds are not stored in metrics.json (only AUROC/PRO metrics
+are saved), so the verdict threshold defaults to 0.5 — adjust FALLBACK_THRESHOLD
+below after inspecting per-category score distributions.
+"""
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import matplotlib.cm as cm
+import torch
+import torchvision.transforms as T
+from PIL import Image
+
+
+THRESHOLDS_FILE = Path(__file__).parent / "thresholds.json"
+_thresholds = json.loads(THRESHOLDS_FILE.read_text()) if THRESHOLDS_FILE.exists() else {}
+
+def load_threshold(category: str) -> float:
+    if category in _thresholds:
+        return _thresholds[category]["threshold"]   # calibrated from train/good/
+    return 0.5   # fallback — run calibrate_thresholds.py to fix this
+
+REPO_ROOT   = Path(__file__).parent
+RESULTS_DIR = REPO_ROOT / "results"
+DATA_DIR    = REPO_ROOT / "anomaly_ds"
+
+sys.path.insert(0, str(REPO_ROOT / "src"))
+from patchcore import PatchCore
+
+
+# ---------------------------------------------------------------------------
+# Category discovery
+# ---------------------------------------------------------------------------
+
+def discover_categories() -> list[str]:
+    return sorted(
+        d.name for d in RESULTS_DIR.iterdir()
+        if d.is_dir() and (d / "patchcore_model.pt").exists()
+    )
+
+
+AVAILABLE_CATEGORIES = discover_categories()
+
+
+# ---------------------------------------------------------------------------
+# Model cache (loaded on first use, stays in RAM)
+# ---------------------------------------------------------------------------
+
+_model_cache: dict = {}
+
+
+def load_model(category: str) -> PatchCore:
+    if category in _model_cache:
+        return _model_cache[category]
+
+    model_path = str(RESULTS_DIR / category / "patchcore_model.pt")
+    model = PatchCore(device="cpu", faiss_gpu=False)
+    model.load(model_path)
+    _model_cache[category] = model
+    return model
+
+
+# ---------------------------------------------------------------------------
+# Preprocessing — must match IMAGE_TRANSFORM in src/dataset.py exactly:
+#   Resize(256) → CenterCrop(224) → ToTensor → Normalize(ImageNet)
+# ---------------------------------------------------------------------------
+
+_transform = T.Compose([
+    T.Resize(256),
+    T.CenterCrop(224),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]),
+])
+
+
+# ---------------------------------------------------------------------------
+# Heatmap overlay (visualize.py only exposes a private numpy helper;
+# this PIL version is self-contained and avoids the Agg backend conflict)
+# ---------------------------------------------------------------------------
+
+def overlay_heatmap(image: Image.Image, anomaly_map: np.ndarray, alpha: float = 0.45) -> Image.Image:
+    image_rgb = image.resize((224, 224)).convert("RGB")
+    norm = (anomaly_map - anomaly_map.min()) / (anomaly_map.max() - anomaly_map.min() + 1e-8)
+    heatmap_rgba = cm.jet(norm)
+    heatmap_rgb = (heatmap_rgba[:, :, :3] * 255).astype(np.uint8)
+    heatmap_pil = Image.fromarray(heatmap_rgb).resize(image_rgb.size)
+    return Image.blend(image_rgb, heatmap_pil, alpha)
+
+
+# ---------------------------------------------------------------------------
+# Main predict function
+# ---------------------------------------------------------------------------
+
+def predict(category: str, image: Image.Image):
+    if image is None:
+        return None, "No image provided.", None
+
+    model = load_model(category)
+    threshold = load_threshold(category)
+
+    tensor = _transform(image.convert("RGB")).unsqueeze(0)  # [1, 3, 224, 224]
+
+    with torch.no_grad():
+        # predict() returns (image_score: float, anomaly_map: np.ndarray [224,224])
+        image_score, anomaly_map = model.predict(tensor)
+
+    verdict = "✅  NORMAL" if image_score < threshold else "❌  ANOMALY"
+    heatmap = overlay_heatmap(image, anomaly_map)
+
+    return float(image_score), verdict, heatmap
+
+
+# ---------------------------------------------------------------------------
+# Example images (first 4 categories, one normal + one defective each)
+# ---------------------------------------------------------------------------
+
+def build_examples() -> list:
+    examples_dir = REPO_ROOT / "examples"
+    if not examples_dir.exists():
+        return []
+    examples = []
+    for img_path in sorted(examples_dir.glob("*.png")):
+        cat = img_path.stem.rsplit("_", 1)[0]   # "bottle_good" → "bottle"
+        if cat in AVAILABLE_CATEGORIES:
+            examples.append([cat, str(img_path)])
+    return examples
+
+
+EXAMPLES = build_examples()
+
+
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
+
+import gradio as gr
+
+with gr.Blocks(title="PatchCore — Industrial Anomaly Detection") as demo:
+    gr.Markdown("## PatchCore — Industrial Anomaly Detection")
+    gr.Markdown(
+        "Select a product category, upload an image (or use your camera), "
+        "and see the anomaly score and a pixel-level heatmap.\n\n"
+        f"**Available categories ({len(AVAILABLE_CATEGORIES)}):** "
+        + ", ".join(AVAILABLE_CATEGORIES)
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            category_dd = gr.Dropdown(
+                choices=AVAILABLE_CATEGORIES,
+                value=AVAILABLE_CATEGORIES[0],
+                label="Product category",
+            )
+            image_in = gr.Image(
+                type="pil",
+                label="Input image",
+                sources=["upload", "webcam"],
+            )
+            run_btn = gr.Button("Detect Anomalies", variant="primary")
+
+        with gr.Column(scale=1):
+            score_out   = gr.Number(label="Anomaly score")
+            verdict_out = gr.Textbox(label="Verdict")
+            heatmap_out = gr.Image(type="pil", label="Anomaly heatmap overlay")
+
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[category_dd, image_in],
+        label="Try an example",
+    )
+
+    run_btn.click(
+        fn=predict,
+        inputs=[category_dd, image_in],
+        outputs=[score_out, verdict_out, heatmap_out],
+    )
+
+if __name__ == "__main__":
+    demo.launch(theme=gr.themes.Soft())

requirements.txt

@@ -0,0 +1,10 @@
+gradio>=4.0
+torch==2.1.0+cpu
+torchvision==0.16.0+cpu
+--extra-index-url https://download.pytorch.org/whl/cpu
+faiss-cpu>=1.7.4
+timm>=0.9.0
+Pillow>=9.0
+numpy>=1.24
+scipy>=1.10
+matplotlib>=3.7

results/bottle/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7a9e581289b5d32a4530cf9952c220c07551503bfbb81711685ae6a11f63cc3
+size 20138721

results/cable/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1bdd948541d11e4d7ab20522053630e8c7ef7fa1459e0289c044a701830f49a
+size 21582561

results/capsule/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcd0c1ca16f83054e63553647235524a44baccfe49b25d54ab0b5e1a038f28a8
+size 21100257

results/carpet/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:191caefcd4a2b469ca4e7002af5d941a459f134dde065c1f04e298fcf6833e02
+size 26976993

results/grid/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d5c934d56e7a9902e4058b88601f210b26d7a1a48439611dc6144f9d3598bd9
+size 25437921

results/hazelnut/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74aed0cb018220a4a512452a2abe2ac1faf49f2009087df47bdf64e5d42c3ca6
+size 37670625

results/leather/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3861c6717c9c3581cf5730215446d6e5aead5ca4b4366dfc4ee3d93c1eb15040
+size 23607009

results/metal_nut/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35194312f53d8c4b1f94e35932bf1be18c5c8dab00df6d63bbae9b64d59585e9
+size 21198561

results/pill/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5f31b19dcd8768ac66566d17f52abaff4c0f344c17f908513e305f1678eccba
+size 25726689

results/screw/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:340dd308c267495f1ebf41f8ac5ddda2b1e54564dd4846d52527b7c0b7631103
+size 30832353

results/tile/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a49a4172b2fbf0bbdc4f0251042fc4ae9d01975fe715e0d8bacee389a79fab57
+size 22160097

results/toothbrush/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0046ec9960f4b08714fbd2ca2fe12d66aa691b146559d7dbffbbb70d48228fcb
+size 5783265

results/transistor/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:235c110540b698cf7cfdbbbbd2b3e4bd558091cc408621ad84dd38cb50c6cf7f
+size 20522721

results/wood/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28f3d64f6ddf1acd45ca177fc8b79ca6d77044dbee08088f019ee078b22745b7
+size 23797473

results/zipper/patchcore_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b8d90f9958306b868985bcabc2a02079875713a4921bc43d7dbb6dd6695f7cd
+size 23124705

src/coreset.py

@@ -0,0 +1,90 @@
+"""
+Greedy coreset subsampling via Johnson-Lindenstrauss projection + furthest-point selection.
+Pure PyTorch — no sklearn dependency at subsampling time.
+
+Algorithm:
+  1. Project features: F_proj = F @ R,  R ~ N(0, 1/sqrt(128)),  shape [N, 128]
+  2. Pick a random starting point, seed the coreset C
+  3. Maintain min-distance-to-coreset vector `min_dists` for all remaining points
+  4. Greedily add the point with maximum min-distance
+  5. Repeat until |C| = ceil(ratio * N)
+
+"""
+
+import math
+import torch
+
+
+PROJ_DIM   = 128
+CHUNK_SIZE = 1000   # rows of the full set to process per distance batch
+
+
+def _pairwise_l2sq_chunked(
+    query: torch.Tensor,   # [M, D]
+    bank:  torch.Tensor,   # [K, D]
+) -> torch.Tensor:
+    """
+    Returns squared L2 distance matrix [M, K] without materializing a huge intermediate.
+    Uses the (a-b)^2 = a^2 - 2ab + b^2 expansion.
+    """
+    # ||a||^2  [M, 1]
+    a2 = (query ** 2).sum(dim=1, keepdim=True)
+    # ||b||^2  [1, K]
+    b2 = (bank ** 2).sum(dim=1, keepdim=True).t()
+    # -2 <a, b>  [M, K]
+    ab = torch.mm(query, bank.t())
+    dist2 = a2 + b2 - 2 * ab
+    dist2.clamp_(min=0.0)   # numerical safety
+    return dist2
+
+
+def subsample_coreset(features: torch.Tensor, ratio: float) -> torch.Tensor:
+    """
+    Greedy furthest-point coreset subsampling.
+
+    Args:
+        features : [N, D] float32 tensor on GPU (or CPU)
+        ratio    : fraction of points to keep (e.g. 0.01 = 1%)
+
+    Returns:
+        [M, D] coreset tensor on the same device as `features`,
+        where M = ceil(ratio * N).
+    """
+    N, D = features.shape
+    M = max(1, math.ceil(ratio * N))
+    device = features.device
+
+    # --- Johnson-Lindenstrauss projection for distance computation only ---
+    torch.manual_seed(42)
+    R = torch.randn(D, PROJ_DIM, device=device) / math.sqrt(PROJ_DIM)
+    F_proj = features @ R          # [N, PROJ_DIM]  (used only for distances)
+
+    # --- Initialise with a random seed point ---
+    start_idx = torch.randint(0, N, (1,)).item()
+    coreset_indices = [start_idx]
+
+    # min distance from each point to the current coreset (in projected space)
+    # Initialise with distance to the first coreset point
+    first = F_proj[start_idx].unsqueeze(0)             # [1, PROJ_DIM]
+    min_dists = _pairwise_l2sq_chunked(F_proj, first).squeeze(1)  # [N]
+
+    # --- Greedy furthest-point loop ---
+    for _ in range(M - 1):
+        new_idx = int(min_dists.argmax().item())
+        coreset_indices.append(new_idx)
+
+        # Update min_dists: distance from every point to the newly added coreset member
+        new_pt = F_proj[new_idx].unsqueeze(0)  # [1, PROJ_DIM]
+
+        # Chunked to avoid OOM
+        for start in range(0, N, CHUNK_SIZE):
+            end  = min(start + CHUNK_SIZE, N)
+            chunk_dists = _pairwise_l2sq_chunked(
+                F_proj[start:end], new_pt
+            ).squeeze(1)                           # [chunk]
+            min_dists[start:end] = torch.minimum(
+                min_dists[start:end], chunk_dists
+            )
+
+    idx_tensor = torch.tensor(coreset_indices, dtype=torch.long, device=device)
+    return features[idx_tensor]                    # [M, D] full-dim vectors

src/feature_extractor.py

@@ -0,0 +1,79 @@
+"""
+Feature extractor using WideResNet-101-2 backbone (timm).
+Extracts multi-scale patch descriptors from layer2 (stride 8) and layer3 (stride 16).
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+
+
+class PatchFeatureExtractor(nn.Module):
+    """
+    Extracts L2-normalized patch descriptors from WideResNet-101-2.
+
+    For a 224×224 input:
+      - layer2 output: [B, 512,  28, 28]
+      - layer3 output: [B, 1024, 14, 14]
+
+    Both are spatially aligned to 28×28 via adaptive avg pooling on layer3,
+    then concatenated → [B, 1536, 28, 28].
+    Each of the 28×28 = 784 spatial locations yields a 1536-dim descriptor.
+
+    No training — runs in eval mode with torch.no_grad().
+    """
+
+    def __init__(self, backbone: str = "wide_resnet101_2", pretrained: bool = True):
+        super().__init__()
+        # features_only=True exposes intermediate feature maps
+        self.backbone = timm.create_model(
+            backbone,
+            pretrained=pretrained,
+            features_only=True,
+            out_indices=(1, 2),   # layer2=index1 (512ch), layer3=index2 (1024ch)
+        )
+        self.backbone.eval()
+        for param in self.backbone.parameters():
+            param.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, 3, 224, 224] ImageNet-normalised images.
+        Returns:
+            patch_features: [B, H*W, 1536] L2-normalised patch descriptors
+                            where H=W=28 for 224px input.
+        """
+        features = self.backbone(x)   # list: [layer2, layer3]
+        f2 = features[0]              # [B, 512,  28, 28]
+        f3 = features[1]              # [B, 1024, 14, 14]
+
+        # Upsample layer3 to match layer2 spatial size
+        f3_up = F.adaptive_avg_pool2d(f3, output_size=f2.shape[-2:])  # [B, 1024, 28, 28]
+
+        # Concatenate along channel dim
+        combined = torch.cat([f2, f3_up], dim=1)  # [B, 1536, 28, 28]
+
+        B, C, H, W = combined.shape
+        # Reshape to [B, H*W, C]
+        patch_features = combined.permute(0, 2, 3, 1).reshape(B, H * W, C)
+
+        # L2 normalize each descriptor
+        patch_features = F.normalize(patch_features, p=2, dim=-1)
+
+        return patch_features
+
+    @torch.no_grad()
+    def extract_patch_features(self, images: torch.Tensor) -> torch.Tensor:
+        """
+        Convenience method: flatten all patches across batch.
+        Args:
+            images: [B, 3, 224, 224]
+        Returns:
+            [B * H * W, 1536]  — ready for coreset subsampling or memory bank
+        """
+        patch_features = self.forward(images)  # [B, 784, 1536]
+        B, N, C = patch_features.shape
+        return patch_features.reshape(B * N, C)

src/patchcore.py

@@ -0,0 +1,232 @@
+"""
+PatchCore model: fit (feature extraction + coreset indexing) and predict (KNN scoring).
+
+No training loop — PatchCore is purely:
+  fit  : one forward pass over normal images → coreset memory bank
+  predict: KNN distance query for each test patch → anomaly map + image score
+"""
+
+import warnings
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy.ndimage import gaussian_filter
+from tqdm import tqdm
+
+from feature_extractor import PatchFeatureExtractor
+from coreset import subsample_coreset
+
+
+class PatchCore:
+    """
+    PatchCore anomaly detector.
+
+    Args:
+        backbone      : timm model name (default: 'wide_resnet101_2')
+        coreset_ratio : fraction of patches to keep in memory bank (default: 0.01)
+        device        : torch device string (default: 'cuda')
+        faiss_gpu     : use faiss GPU index (default: True, fallback to CPU)
+        gaussian_sigma: sigma for anomaly map smoothing (default: 4)
+    """
+
+    PATCH_GRID = 28    # spatial grid size for 224px input (both axes)
+    FEAT_DIM   = 1536  # layer2 (512) + layer3 (1024)
+
+    def __init__(
+        self,
+        backbone: str = "wide_resnet101_2",
+        coreset_ratio: float = 0.01,
+        device: str = "cuda",
+        faiss_gpu: bool = True,
+        gaussian_sigma: float = 4.0,
+    ):
+        self.backbone      = backbone
+        self.coreset_ratio = coreset_ratio
+        self.device        = torch.device(device if torch.cuda.is_available() else "cpu")
+        self.faiss_gpu     = faiss_gpu
+        self.gaussian_sigma = gaussian_sigma
+
+        self.extractor = PatchFeatureExtractor(backbone=backbone, pretrained=True)
+        self.extractor = self.extractor.to(self.device)
+
+        self.memory_bank: torch.Tensor = None  # [M, 1536]
+        self._faiss_index = None
+        self._index_backend = None
+
+    # ------------------------------------------------------------------
+    # Fit
+    # ------------------------------------------------------------------
+
+    def fit(self, train_loader) -> None:
+        """
+        Extract patch features from all normal training images,
+        run coreset subsampling, build faiss index.
+        """
+        print("[PatchCore] Extracting features from training images …")
+        all_features = []
+
+        self.extractor.eval()
+        with torch.no_grad():
+            for batch in tqdm(train_loader, desc="  Feature extraction", leave=False):
+                if isinstance(batch, (list, tuple)):
+                    imgs = batch[0]
+                else:
+                    imgs = batch
+                imgs = imgs.to(self.device)
+                feats = self.extractor.extract_patch_features(imgs)  # [B*784, 1536]
+                all_features.append(feats.cpu())
+
+        all_features = torch.cat(all_features, dim=0)  # [N_total, 1536]
+        print(f"[PatchCore] Total patches before coreset: {len(all_features):,}")
+
+        # Move to GPU for fast coreset computation
+        all_features = all_features.to(self.device)
+
+        print(f"[PatchCore] Running coreset subsampling (ratio={self.coreset_ratio}) …")
+        self.memory_bank = subsample_coreset(all_features, self.coreset_ratio)
+        print(f"[PatchCore] Memory bank size after coreset: {len(self.memory_bank):,}  "
+              f"({self.memory_bank.element_size() * self.memory_bank.numel() / 1e6:.1f} MB)")
+
+        self._build_faiss_index()
+
+    def _build_faiss_index(self) -> None:
+        """Build KNN backend from the memory bank (faiss preferred, torch fallback)."""
+        bank_np = self.memory_bank.cpu().numpy().astype(np.float32)
+
+        try:
+            import faiss
+
+            d = bank_np.shape[1]
+            index_flat = faiss.IndexFlatL2(d)
+
+            if self.faiss_gpu and torch.cuda.is_available():
+                try:
+                    res = faiss.StandardGpuResources()
+                    self._faiss_index = faiss.index_cpu_to_gpu(res, 0, index_flat)
+                    self._index_backend = "faiss-gpu"
+                    print("[PatchCore] Using faiss GPU index.")
+                except Exception as e:
+                    warnings.warn(f"[PatchCore] faiss GPU index failed ({e}); falling back to faiss CPU.")
+                    self._faiss_index = index_flat
+                    self._index_backend = "faiss-cpu"
+            else:
+                self._faiss_index = index_flat
+                self._index_backend = "faiss-cpu"
+
+            self._faiss_index.add(bank_np)
+
+        except Exception as e:
+            warnings.warn(
+                "[PatchCore] faiss is unavailable/incompatible; "
+                f"falling back to torch KNN search. Details: {e}"
+            )
+            # Keep a contiguous float bank on the configured device for fallback KNN.
+            self.memory_bank = self.memory_bank.float().contiguous().to(self.device)
+            self._faiss_index = None
+            self._index_backend = "torch"
+
+    def _search_knn(self, feats_np: np.ndarray, k: int = 1) -> Tuple[np.ndarray, np.ndarray]:
+        """Return (squared_l2_distances, indices) for nearest neighbours."""
+        if self._index_backend in ("faiss-gpu", "faiss-cpu"):
+            return self._faiss_index.search(feats_np, k=k)
+
+        # Torch fallback for environments where faiss import fails (e.g., NumPy ABI mismatch).
+        query = torch.from_numpy(feats_np).to(self.device, dtype=torch.float32)
+        bank = self.memory_bank
+
+        # Chunking keeps peak memory usage bounded for larger memory banks.
+        chunk_size = 1024
+        dist_chunks = []
+        idx_chunks = []
+
+        with torch.no_grad():
+            for start in range(0, query.shape[0], chunk_size):
+                q = query[start:start + chunk_size]  # [Q, D]
+                d2 = torch.sum((q[:, None, :] - bank[None, :, :]) ** 2, dim=-1)  # [Q, M]
+                vals, idxs = torch.topk(d2, k=k, dim=1, largest=False)
+                dist_chunks.append(vals)
+                idx_chunks.append(idxs)
+
+        distances = torch.cat(dist_chunks, dim=0).detach().cpu().numpy().astype(np.float32)
+        indices = torch.cat(idx_chunks, dim=0).detach().cpu().numpy().astype(np.int64)
+        return distances, indices
+
+    # ------------------------------------------------------------------
+    # Predict
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    def predict(self, image_tensor: torch.Tensor) -> Tuple[float, np.ndarray]:
+        """
+        Compute anomaly score and pixel-level anomaly map for a single image.
+
+        Args:
+            image_tensor: [1, 3, 224, 224] normalised image tensor
+
+        Returns:
+            image_score : float — max patch distance (image-level anomaly score)
+            anomaly_map : np.ndarray [224, 224] — smoothed, upsampled patch distance map
+        """
+        image_tensor = image_tensor.to(self.device)
+
+        # Extract patch features: [1*784, 1536]
+        feats = self.extractor.extract_patch_features(image_tensor)
+        feats_np = feats.cpu().numpy().astype(np.float32)
+
+        # KNN distance query (k=1 nearest neighbour)
+        distances, _ = self._search_knn(feats_np, k=1)  # [N_patches, 1]
+        patch_scores = distances[:, 0]  # [N_patches]  squared L2 distances
+
+        # Compute actual grid size from number of patches (robust to different feature extractors)
+        num_patches = len(patch_scores)
+        patch_grid = int(np.sqrt(num_patches))
+        assert patch_grid * patch_grid == num_patches, \
+            f"Non-square patch grid: {num_patches} patches (expected {patch_grid}²)"
+
+        # Reshape to spatial grid
+        score_map = patch_scores.reshape(patch_grid, patch_grid)
+
+        # Upsample to 224×224 via bilinear interpolation
+        score_tensor = torch.from_numpy(score_map).unsqueeze(0).unsqueeze(0)  # [1,1,G,G]
+        score_upsampled = F.interpolate(
+            score_tensor, size=(224, 224), mode="bilinear", align_corners=False
+        ).squeeze().numpy()  # [224, 224]
+
+        # Gaussian smoothing
+        anomaly_map = gaussian_filter(score_upsampled, sigma=self.gaussian_sigma)
+
+        # Image-level score: max patch distance
+        image_score = float(patch_scores.max())
+
+        return image_score, anomaly_map
+
+    # ------------------------------------------------------------------
+    # Save / Load
+    # ------------------------------------------------------------------
+
+    def save(self, path: str) -> None:
+        """Serialize memory bank and config to a .pt file."""
+        torch.save(
+            {
+                "memory_bank": self.memory_bank.cpu(),
+                "backbone": self.backbone,
+                "coreset_ratio": self.coreset_ratio,
+                "gaussian_sigma": self.gaussian_sigma,
+            },
+            path,
+        )
+        print(f"[PatchCore] Model saved to {path}")
+
+    def load(self, path: str) -> None:
+        """Load memory bank and config from a .pt file, rebuild faiss index."""
+        ckpt = torch.load(path, map_location="cpu")
+        self.memory_bank    = ckpt["memory_bank"].to(self.device)
+        self.backbone       = ckpt["backbone"]
+        self.coreset_ratio  = ckpt["coreset_ratio"]
+        self.gaussian_sigma = ckpt["gaussian_sigma"]
+        self._build_faiss_index()
+        print(f"[PatchCore] Model loaded from {path}  "
+              f"(memory bank: {len(self.memory_bank):,} vectors)")

src/visualize.py

@@ -0,0 +1,191 @@
+"""
+Visualization utilities for PatchCore anomaly detection results.
+
+Produces:
+  - anomaly_grid.png    : 8 anomalous + 8 normal samples with heatmap overlays
+  - roc_curves.png      : Image AUROC | Pixel AUROC | PRO curve (3-subplot figure)
+  - score_distribution.png : KDE of anomaly scores, normal vs anomalous
+"""
+
+import os
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.cm as cm
+import seaborn as sns
+import cv2
+
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406])
+IMAGENET_STD  = np.array([0.229, 0.224, 0.225])
+
+
+def _denormalize(tensor_img) -> np.ndarray:
+    """Convert [3, H, W] normalised tensor to [H, W, 3] uint8 numpy image."""
+    img = tensor_img.cpu().numpy().transpose(1, 2, 0)
+    img = img * IMAGENET_STD + IMAGENET_MEAN
+    img = np.clip(img, 0, 1)
+    return (img * 255).astype(np.uint8)
+
+
+def _heatmap_overlay(img_np: np.ndarray, anomaly_map: np.ndarray, alpha: float = 0.4) -> np.ndarray:
+    """Overlay jet colormap anomaly map onto RGB image."""
+    score_norm = (anomaly_map - anomaly_map.min()) / (anomaly_map.max() - anomaly_map.min() + 1e-8)
+    heatmap = (cm.jet(score_norm)[:, :, :3] * 255).astype(np.uint8)
+    overlay = (alpha * heatmap + (1 - alpha) * img_np).astype(np.uint8)
+    return overlay
+
+
+def save_anomaly_grid(
+    test_dataset,
+    anomaly_maps: List[np.ndarray],
+    scores: List[float],
+    labels: List[int],
+    save_path: str,
+    n: int = 16,
+) -> None:
+    """
+    Select 8 anomalous + 8 normal samples.
+    For each: original image | GT mask | anomaly heatmap overlay.
+    """
+    os.makedirs(Path(save_path).parent, exist_ok=True)
+
+    anomaly_idx = [i for i, l in enumerate(labels) if l == 1]
+    normal_idx  = [i for i, l in enumerate(labels) if l == 0]
+
+    n_each = n // 2
+    # Sample uniformly if more samples than needed
+    if len(anomaly_idx) > n_each:
+        step = len(anomaly_idx) // n_each
+        anomaly_idx = anomaly_idx[::step][:n_each]
+    if len(normal_idx) > n_each:
+        step = len(normal_idx) // n_each
+        normal_idx = normal_idx[::step][:n_each]
+
+    selected = anomaly_idx[:n_each] + normal_idx[:n_each]
+    actual_n = len(selected)
+
+    fig, axes = plt.subplots(actual_n, 3, figsize=(9, actual_n * 3))
+    if actual_n == 1:
+        axes = axes[np.newaxis, :]
+
+    for row, idx in enumerate(selected):
+        img_tensor, label, gt_mask_tensor, defect_type = test_dataset[idx]
+
+        img_np   = _denormalize(img_tensor)
+        gt_np    = gt_mask_tensor.squeeze().cpu().numpy()
+        amap     = anomaly_maps[idx]
+        score    = scores[idx]
+        tag      = "ANOMALY" if label == 1 else "NORMAL"
+        color    = "red"    if label == 1 else "green"
+
+        overlay  = _heatmap_overlay(img_np, amap)
+
+        axes[row, 0].imshow(img_np)
+        axes[row, 0].set_title(f"[{tag}] score={score:.3f}", color=color, fontsize=8)
+        axes[row, 0].axis("off")
+
+        axes[row, 1].imshow(gt_np, cmap="gray", vmin=0, vmax=1)
+        axes[row, 1].set_title("GT Mask", fontsize=8)
+        axes[row, 1].axis("off")
+
+        axes[row, 2].imshow(overlay)
+        axes[row, 2].set_title("Anomaly Heatmap", fontsize=8)
+        axes[row, 2].axis("off")
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  Saved anomaly grid → {save_path}")
+
+
+def save_roc_curves(
+    image_roc: dict,
+    pixel_roc_auroc: float,
+    pro_curve: dict,
+    save_path: str,
+) -> None:
+    """
+    3-subplot figure: Image AUROC | Pixel AUROC | PRO curve.
+    """
+    os.makedirs(Path(save_path).parent, exist_ok=True)
+
+    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+
+    # --- Image AUROC ---
+    ax = axes[0]
+    ax.plot(image_roc["fpr"], image_roc["tpr"], lw=2,
+            label=f'AUROC={image_roc["auroc"]:.3f}')
+    ax.plot([0, 1], [0, 1], "k--", lw=1)
+    # Mark operating point (Youden's J)
+    thr = image_roc["threshold"]
+    fpr_arr, tpr_arr = image_roc["fpr"], image_roc["tpr"]
+    j = tpr_arr - fpr_arr
+    best = int(np.argmax(j))
+    ax.scatter(fpr_arr[best], tpr_arr[best], marker="o", color="red", zorder=5,
+               label=f"Threshold={thr:.3f}")
+    ax.set_xlabel("FPR")
+    ax.set_ylabel("TPR")
+    ax.set_title("Image-level ROC")
+    ax.legend(fontsize=8)
+    ax.grid(True, alpha=0.3)
+
+    # --- Pixel AUROC (text only — full pixel ROC arrays not stored to save memory) ---
+    ax = axes[1]
+    ax.text(0.5, 0.5, f"Pixel AUROC\n{pixel_roc_auroc:.4f}",
+            ha="center", va="center", fontsize=20, transform=ax.transAxes,
+            bbox=dict(boxstyle="round", facecolor="lightblue", alpha=0.5))
+    ax.set_title("Pixel-level AUROC")
+    ax.axis("off")
+
+    # --- PRO curve ---
+    ax = axes[2]
+    fpr_pro = pro_curve["fpr_array"]
+    pro_pro = pro_curve["pro_array"]
+    ax.plot(fpr_pro, pro_pro, lw=2, label=f'PRO-AUC={pro_curve["pro_auc"]:.3f}')
+    ax.axvline(x=0.3, color="red", linestyle="--", lw=1, label="FPR=0.3 cutoff")
+    ax.set_xlabel("Mean FPR")
+    ax.set_ylabel("Mean Per-Region Overlap")
+    ax.set_title("PRO Curve")
+    ax.legend(fontsize=8)
+    ax.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  Saved ROC curves    → {save_path}")
+
+
+def save_score_distribution(
+    scores: List[float],
+    labels: List[int],
+    save_path: str,
+) -> None:
+    """KDE plot of anomaly score distributions: normal vs anomalous with overlap shading."""
+    os.makedirs(Path(save_path).parent, exist_ok=True)
+
+    scores = np.array(scores)
+    labels = np.array(labels)
+    normal_scores   = scores[labels == 0]
+    anomaly_scores  = scores[labels == 1]
+
+    fig, ax = plt.subplots(figsize=(8, 5))
+
+    if len(normal_scores) > 1:
+        sns.kdeplot(normal_scores, ax=ax, label="Normal", color="green", fill=True, alpha=0.4)
+    if len(anomaly_scores) > 1:
+        sns.kdeplot(anomaly_scores, ax=ax, label="Anomalous", color="red", fill=True, alpha=0.4)
+
+    ax.set_xlabel("Anomaly Score (max patch distance)")
+    ax.set_ylabel("Density")
+    ax.set_title("Score Distribution: Normal vs Anomalous")
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=300, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  Saved score dist    → {save_path}")

thresholds.json

@@ -0,0 +1,122 @@
+{
+  "bottle": {
+    "threshold": 0.2890068769454956,
+    "percentile": 95,
+    "n_images": 209,
+    "score_min": 0.24581491947174072,
+    "score_mean": 0.2680253663131495,
+    "score_max": 0.3055186867713928
+  },
+  "cable": {
+    "threshold": 0.5094758927822113,
+    "percentile": 95,
+    "n_images": 224,
+    "score_min": 0.4513741135597229,
+    "score_mean": 0.4833795888615506,
+    "score_max": 0.536372184753418
+  },
+  "capsule": {
+    "threshold": 0.1864305555820465,
+    "percentile": 95,
+    "n_images": 219,
+    "score_min": 0.1464100480079651,
+    "score_mean": 0.16722383836632995,
+    "score_max": 0.20267021656036377
+  },
+  "carpet": {
+    "threshold": 0.2884562313556671,
+    "percentile": 95,
+    "n_images": 280,
+    "score_min": 0.25853419303894043,
+    "score_mean": 0.27355455990348543,
+    "score_max": 0.3000943660736084
+  },
+  "grid": {
+    "threshold": 0.2998784899711609,
+    "percentile": 95,
+    "n_images": 264,
+    "score_min": 0.25811439752578735,
+    "score_mean": 0.28114455396478827,
+    "score_max": 0.3175504207611084
+  },
+  "hazelnut": {
+    "threshold": 0.38583904504776,
+    "percentile": 95,
+    "n_images": 391,
+    "score_min": 0.32867467403411865,
+    "score_mean": 0.36304525600369936,
+    "score_max": 0.4044198989868164
+  },
+  "leather": {
+    "threshold": 0.2573680400848389,
+    "percentile": 95,
+    "n_images": 245,
+    "score_min": 0.22839152812957764,
+    "score_mean": 0.244711997071091,
+    "score_max": 0.27069514989852905
+  },
+  "metal_nut": {
+    "threshold": 0.4237459897994995,
+    "percentile": 95,
+    "n_images": 220,
+    "score_min": 0.3752959370613098,
+    "score_mean": 0.40138462700627064,
+    "score_max": 0.4571906328201294
+  },
+  "pill": {
+    "threshold": 0.3175112664699554,
+    "percentile": 95,
+    "n_images": 267,
+    "score_min": 0.2706294059753418,
+    "score_mean": 0.295599180214414,
+    "score_max": 0.3364473581314087
+  },
+  "screw": {
+    "threshold": 0.3218137860298157,
+    "percentile": 95,
+    "n_images": 320,
+    "score_min": 0.275600790977478,
+    "score_mean": 0.30216500014066694,
+    "score_max": 0.35416650772094727
+  },
+  "tile": {
+    "threshold": 0.4006050288677215,
+    "percentile": 95,
+    "n_images": 230,
+    "score_min": 0.35888969898223877,
+    "score_mean": 0.3827742672484854,
+    "score_max": 0.4205925464630127
+  },
+  "toothbrush": {
+    "threshold": 0.4440391033887863,
+    "percentile": 95,
+    "n_images": 60,
+    "score_min": 0.38699811697006226,
+    "score_mean": 0.4185704320669174,
+    "score_max": 0.4575169086456299
+  },
+  "transistor": {
+    "threshold": 0.4263874530792236,
+    "percentile": 95,
+    "n_images": 213,
+    "score_min": 0.3608095049858093,
+    "score_mean": 0.4002147546396569,
+    "score_max": 0.44323837757110596
+  },
+  "wood": {
+    "threshold": 0.3080090880393982,
+    "percentile": 95,
+    "n_images": 247,
+    "score_min": 0.2672466039657593,
+    "score_mean": 0.2905003445351172,
+    "score_max": 0.3450666666030884
+  },
+  "zipper": {
+    "threshold": 0.18321026563644408,
+    "percentile": 95,
+    "n_images": 240,
+    "score_min": 0.15239810943603516,
+    "score_mean": 0.16628144855300586,
+    "score_max": 0.19208133220672607
+  }
+}