Upload full code for Space

README.md

@@ -1,14 +1 @@
----
-title: AnomalyDetection
-emoji: 📈
-colorFrom: green
-colorTo: yellow
-sdk: gradio
-sdk_version: 6.1.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: Anomaly detection using Dino v3
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## To run and reproduce the PatchKNN model, use interactive demo notebook

backbones/__init__.py

@@ -0,0 +1,24 @@
+import timm
+
+from backbones.dino_v2 import build_dinov2_small, build_dinov2_base, build_dinov2_large
+from backbones.dino_v3 import build_dinov3_small, build_dinov3_base, build_dinov3_large
+
+"""
+Model registry for backbones
+"""
+
+
+_BACKBONES = {
+    "dinov2_small": build_dinov2_small,
+    "dinov2_base": build_dinov2_base,
+    "dinov2_large": build_dinov2_large,
+    "dinov3_small": build_dinov3_small,
+    "dinov3_base": build_dinov3_base,
+    "dinov3_large": build_dinov3_large,
+
+}
+
+def get_backbone(name: str, **kwargs):
+    if name not in _BACKBONES:
+        raise ValueError(f"Unknown backbone '{name}'. Available: {list(_BACKBONES)}")
+    return _BACKBONES[name](**kwargs)

backbones/dino_v2.py

@@ -0,0 +1,22 @@
+
+import timm
+"""
+DinoV2 backbones
+"""
+def build_dinov2_small(**kwargs):
+
+    model = timm.create_model('vit_small_patch14_reg4_dinov2.lvd142m',
+                              pretrained=True, num_classes=0)
+    return model
+
+def build_dinov2_base(**kwargs):
+
+    model = timm.create_model('vit_base_patch14_reg4_dinov2.lvd142m',
+                              pretrained=True, num_classes=0)
+    return model
+
+def build_dinov2_large(**kwargs):
+
+    model = timm.create_model('vit_large_patch14_dinov2.lvd142m',
+                              pretrained=True, num_classes=0)
+    return model

backbones/dino_v3.py

@@ -0,0 +1,33 @@
+import timm
+"""
+DinoV3 backbones
+"""
+
+def build_dinov3_small(**kwargs):
+    model = timm.create_model(
+        "vit_small_plus_patch16_dinov3.lvd1689m",
+        pretrained=True,
+        num_classes=0,
+        **kwargs,
+    )
+    return model
+
+
+def build_dinov3_base(**kwargs):
+    model = timm.create_model(
+        "vit_base_patch16_dinov3.lvd1689m",
+        pretrained=True,
+        num_classes=0,
+        **kwargs,
+    )
+    return model
+
+
+def build_dinov3_large(**kwargs):
+    model = timm.create_model(
+        "vit_large_patch16_dinov3.lvd1689m",
+        pretrained=True,
+        num_classes=0,
+        **kwargs,
+    )
+    return model

dataset/dataloader.py

@@ -0,0 +1,28 @@
+import os
+from typing import List, Tuple
+
+from anomalib.data import MVTec, MVTecAD
+from anomalib.data.datasets.image import MVTecDataset
+"""Loading logic :
+   Download MVTech AD dataset from the website and place it in data folder (Anomalib installation doesnt work, sorry you have to donwload manually)
+
+"""
+
+
+def load_mvtec(category: str, root: str = "./datasets/MVTec",) :
+    train_ds = MVTecDataset(
+        root=root,
+        category=category,
+        split="train"
+    )
+
+    test_ds = MVTecDataset(
+        root=root,
+        category=category,
+        split="test"
+    )
+    
+    train_paths = train_ds.samples["image_path"].tolist()
+    test_paths = test_ds.samples["image_path"].tolist()
+
+    return train_paths, test_paths

evaluation/anomaly_evaluator.py

@@ -0,0 +1,111 @@
+import numpy as np
+import cv2
+from sklearn.metrics import roc_auc_score, precision_recall_curve, average_precision_score
+from skimage.measure import label, regionprops
+
+class AnomalyEvaluator:
+    def __init__(self, pixel_subsample_rate=0.01, compute_pro=False):
+        self.subsample_rate = pixel_subsample_rate
+        self.compute_pro = compute_pro
+        self.reset()
+
+    def reset(self):
+        self.img_preds = []
+        self.img_labels = []
+        
+        self.pix_preds = [] 
+        self.pix_labels = [] 
+        
+        self.full_amaps = []
+        self.full_masks = []
+
+    def update(self, image_score, gt_label, anomaly_map=None, gt_mask=None):
+
+
+        self.img_preds.append(image_score)
+        self.img_labels.append(gt_label)
+
+        if anomaly_map is not None and gt_mask is not None:
+            self._update_pixel_metrics(anomaly_map, gt_mask)
+
+    def _update_pixel_metrics(self, amap, mask):
+        if mask.shape != amap.shape:
+            mask = cv2.resize(mask, (amap.shape[1], amap.shape[0]), interpolation=cv2.INTER_NEAREST)
+        
+        mask = (mask > 0).astype(int)
+
+        if self.compute_pro:
+            self.full_amaps.append(amap)
+            self.full_masks.append(mask)
+
+        flat_amap = amap.flatten()
+        flat_mask = mask.flatten()
+
+        if self.compute_pro or self.subsample_rate >= 1.0:
+            self.pix_preds.extend(flat_amap)
+            self.pix_labels.extend(flat_mask)
+        else:
+            # Random Subsampling to save memory
+            num_pixels = len(flat_mask)
+            sample_size = int(num_pixels * self.subsample_rate)
+            indices = np.random.choice(num_pixels, sample_size, replace=False)
+            self.pix_preds.extend(flat_amap[indices])
+            self.pix_labels.extend(flat_mask[indices])
+
+    def compute(self):
+        results = {}
+        y_true = np.array(self.img_labels)
+        y_score = np.array(self.img_preds)
+        
+        results['image_auroc'] = roc_auc_score(y_true, y_score)
+        
+        results['image_ap'] = average_precision_score(y_true, y_score)
+
+        prec, rec, _ = precision_recall_curve(y_true, y_score)
+        f1_scores = 2 * (prec * rec) / (prec + rec + 1e-8)
+        results['image_f1_max'] = np.max(f1_scores)
+
+        if len(self.pix_labels) > 0:
+            pix_true = np.array(self.pix_labels)
+            pix_score = np.array(self.pix_preds)
+            
+            results['pixel_auroc'] = roc_auc_score(pix_true, pix_score)
+            
+            prec_p, rec_p, thresholds_p = precision_recall_curve(pix_true, pix_score)
+            f1_p = 2 * (prec_p * rec_p) / (prec_p + rec_p + 1e-8)
+            best_idx = np.argmax(f1_p)
+            best_threshold = thresholds_p[best_idx] if best_idx < len(thresholds_p) else 0.5
+            
+            results['pixel_f1_max'] = np.max(f1_p)
+
+            if self.compute_pro:
+                results['pixel_pro'] = self._compute_pro(best_threshold)
+
+        return results
+
+    def _compute_pro(self, threshold):
+
+        total_pro = 0
+        n_defects = 0
+
+        for i in range(len(self.full_amaps)):
+            gt = self.full_masks[i]
+            # Skip normal images
+            if np.sum(gt) == 0:
+                continue
+
+            pred_mask = (self.full_amaps[i] >= threshold).astype(int)
+
+            # Label connected components in Ground Truth
+            labeled_gt = label(gt)
+            regions = regionprops(labeled_gt)
+
+            for region in regions:
+                n_defects += 1
+                blob_mask = (labeled_gt == region.label)
+                overlap_pixels = np.sum(pred_mask & blob_mask)
+                blob_area = region.area
+
+                total_pro += (overlap_pixels / blob_area)
+
+        return total_pro / n_defects if n_defects > 0 else 0.0

main.py

@@ -0,0 +1,135 @@
+
+from __future__ import annotations
+import os
+import cv2
+import numpy as np
+import torch
+
+from segmenters import PCASegmenter
+from segmenters.sam3 import SAM3Segmenter
+from backbones import get_backbone
+from dataset.dataloader import load_mvtec
+from models.model_bank_knn import PatchKNNDetector
+from evaluation.anomaly_evaluator import AnomalyEvaluator 
+
+def main(
+    category: str = "bottle",
+    root: str | None = None,
+    backbone_name: str = "dinov3_small",
+    use_sam3: bool = True,
+    use_pca: bool = False,
+    pca_backbone_name: str | None = None,
+    return_results: bool = False,
+    backbone_model=None,
+    segmenter_obj=None,
+    n_ref: int = 1,
+):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # 1. Setup
+    root = root or os.path.join("dataset", "mvtec_anomaly_detection")
+    train_paths, test_paths = load_mvtec(category=category, root=root)
+    test_paths = test_paths[::5]
+    train_paths = train_paths[::5]
+    print(f"{category}: {len(train_paths)} train, {len(test_paths)} test images")
+
+    # 2. Initialize Evaluator with PRO support
+    evaluator = AnomalyEvaluator(pixel_subsample_rate=0.01, compute_pro=False)
+
+    # 3. Model Init
+    segmenter = segmenter_obj
+    if segmenter is None:
+        if use_pca:
+            pca_backbone = pca_backbone_name or backbone_name
+            segmenter = PCASegmenter(backbone_name=pca_backbone, device=device)
+            print(f"Using PCA segmenter (backbone={pca_backbone})")
+        elif use_sam3:
+            segmenter = SAM3Segmenter(text_prompt=category, device=device)
+            print("Using SAM3 segmenter")
+        else:
+            print("No segmenter selected; using full-image foreground.")
+        
+
+    backbone = backbone_model or get_backbone(backbone_name)
+
+    model = PatchKNNDetector(
+        backbone=backbone,
+        segmenter=segmenter,
+        device=device,
+        k_neighbors=1,
+    )
+
+    print(f"Fitting model... (backbone={backbone_name}, sam3={use_sam3})")
+    model.fit(train_paths, n_ref=n_ref)
+
+    # 4. Evaluation Loop
+    print(f"Starting evaluation on {len(test_paths)} images...")
+
+    for i, path in enumerate(test_paths):
+        # Predict
+        image, amap, score = model.predict(path)
+
+        # Ground Truth Logic
+        is_anomaly = 0 if "good" in path else 1
+        
+        if is_anomaly == 0:
+            gt_mask = np.zeros_like(amap)
+        else:
+            mask_path = path.replace("test", "ground_truth").replace(".png", "_mask.png")
+            if os.path.exists(mask_path):
+                gt_mask = cv2.imread(mask_path, 0)
+                if gt_mask.shape != amap.shape:
+                    gt_mask = cv2.resize(gt_mask, (amap.shape[1], amap.shape[0]), interpolation=cv2.INTER_NEAREST)
+                gt_mask = (gt_mask > 0).astype(int)
+            else:
+                gt_mask = np.zeros_like(amap)
+
+        # Update Evaluator
+        evaluator.update(image_score=score, gt_label=is_anomaly, anomaly_map=amap, gt_mask=gt_mask)
+
+        if i % 20 == 0:
+            print(f"Processed {i}/{len(test_paths)}...")
+
+    # 5. Compute & Print Results
+    results = evaluator.compute()
+    
+    print("\n" + "="*40)
+    print(f"FINAL RESULTS: {category}")
+    print("-" * 40)
+    
+    # Image Level
+    print(f"Image AUROC:   {results['image_auroc']:.4f}")
+    print(f"Image F1-Max:  {results['image_f1_max']:.4f}")
+    print(f"Image AP:      {results['image_ap']:.4f}")
+    print("-" * 40)
+    
+    # Pixel Level
+    if 'pixel_auroc' in results:
+        print(f"Pixel AUROC:   {results['pixel_auroc']:.4f}")
+        print(f"Pixel F1-Max:  {results['pixel_f1_max']:.4f}")
+        #print(f"PRO Score:     {results['pixel_pro']:.4f}") 
+    print("="*40)
+
+    return results if return_results else None
+
+if __name__ == "__main__":
+
+
+    category = os.environ.get("MVTec_CATEGORY", "bottle")
+    root = os.environ.get("MVTec_ROOT", None)
+    backbone_name = os.environ.get("BACKBONE_NAME", "dinov3_small")
+    use_sam3_env = os.environ.get("USE_SAM3", "0").lower()
+    use_sam3 = use_sam3_env not in {"0", "false", "no"}
+    use_pca_env = os.environ.get("USE_PCA", "1").lower()
+    use_pca = use_pca_env in {"1", "true", "yes"}
+    pca_backbone_name = os.environ.get("PCA_BACKBONE", None)
+
+    main(
+        category=category,
+        root=root,
+        backbone_name=backbone_name,
+        use_sam3=use_sam3,
+        use_pca=use_pca,
+        pca_backbone_name=pca_backbone_name,
+        n_ref=int(os.environ.get("N_REF", 1)),
+    )

models/model_bank_knn.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+from typing import Iterable, Tuple, Optional
+
+import cv2
+import numpy as np
+from PIL import Image
+
+import torch
+from timm.data import resolve_data_config
+import torchvision.transforms as T
+from torchvision.transforms import InterpolationMode
+
+from segmenters import BaseSegmenter
+from utils.visualize import visualize_segmentation
+
+
+class PatchKNNDetector:
+
+    def __init__(self, backbone, segmenter = None, device = "cuda", k_neighbors = 1):
+
+
+        self.device = device
+        self.backbone = backbone.to(device)
+        #Switch backbone to inference mode 
+        self.backbone.eval()
+
+        self.segmenter = segmenter
+        self.k_neighbors = k_neighbors
+
+        # Prepare resize/normalize augmentations shared by DINO and SAM
+        data_cfg = resolve_data_config({}, model=self.backbone)
+        _, self.img_size, _ = data_cfg["input_size"]
+        interp = data_cfg.get("interpolation", "bicubic")
+
+        self.transform = T.Compose(
+            [
+                T.Resize(self.img_size, interpolation=getattr(InterpolationMode, interp.upper(), InterpolationMode.BICUBIC)),
+                T.ToTensor(),
+                T.Normalize(mean=data_cfg.get("mean", (0.485, 0.456, 0.406)),
+                            std=data_cfg.get("std", (0.229, 0.224, 0.225))),
+            ]
+        )
+
+        self.num_register_tokens = getattr(self.backbone, "num_register_tokens", 0)
+
+        # Memory bank of foreground patch embeddings
+        self.memory_bank = None
+        self.patch_grid_size = None
+
+
+    def fit(self, train_image_paths, n_ref = 1):
+        """Populate memory bank with references """
+
+        selected_paths = list(train_image_paths)[:n_ref]
+        all_patches = []
+
+        for path in selected_paths:
+
+            #Extracting features
+            image = self._load_image(path)
+            patch_feats, grid_size = self._extract_patch_features(image)
+
+            # Applying foreground mask
+            patch_mask = self._compute_patch_mask(image, grid_size)
+            patch_feats_fg = patch_feats[patch_mask]
+
+            all_patches.append(patch_feats_fg)
+            self.patch_grid_size = grid_size
+
+        self.memory_bank = np.concatenate(all_patches, axis=0)
+
+    def predict(self, image_path) :
+        """
+        Run anomaly detection inference
+        """
+        image = self._load_image(image_path)
+        patch_feats, grid_size = self._extract_patch_features(image)
+
+        patch_mask = self._compute_patch_mask(image, grid_size)
+
+        # Compute distances only on foreground patches
+        scores_fg = self._knn_distances(patch_feats[patch_mask])
+
+        # Put scores back into full patch grid
+        scores_all = np.zeros(patch_feats.shape[0], dtype=np.float32)
+        scores_all[patch_mask] = scores_fg
+        patch_map = scores_all.reshape(grid_size)
+
+        # Upsample to full image ( just for visualization)
+        h, w = image.shape[:2]
+        anomaly_map = cv2.resize(
+            patch_map,
+            (w, h),
+            interpolation=cv2.INTER_CUBIC,
+        ).astype(np.float32)
+
+        # Using image-level score - mean of top 1% patch scores
+        image_score = self._mean_top_percent(scores_fg, top_percent=1.0)
+
+        return image, anomaly_map, image_score
+
+
+    def _load_image(self, path):
+        bgr = cv2.imread(path, cv2.IMREAD_COLOR)
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        return rgb
+
+    @staticmethod
+    def _l2_normalize(feats: np.ndarray, eps: float = 1e-12) -> np.ndarray:
+        norm = np.linalg.norm(feats, axis=1, keepdims=True)
+        return feats / np.maximum(norm, eps)
+
+    def _extract_patch_features(self, image: np.ndarray) :
+        """
+        Run backbone on a single image and return patch features
+        """
+        pil_resized, _ = self._resize_for_model(image)
+        x = self.transform(pil_resized).unsqueeze(0).to(self.device)
+
+        with torch.inference_mode():
+            out = self.backbone.forward_features(x)
+
+        tokens = out.get("x_norm_patchtokens") if isinstance(out, dict) else out
+        
+        if tokens is None and isinstance(out, dict):
+            tokens = out.get("x")
+        if tokens is not None and tokens.ndim == 4:
+            B, C, Hf, Wf = tokens.shape
+            tokens = tokens.permute(0, 2, 3, 1).reshape(B, Hf * Wf, C)
+
+        B, N, C = tokens.shape
+
+        if hasattr(self.backbone, "patch_embed") and hasattr(self.backbone.patch_embed, "grid_size"):
+            gh, gw = self.backbone.patch_embed.grid_size
+        else:
+            gh = int(np.sqrt(N))
+            gw = max(1, N // max(1, gh))
+
+        n_patches = gh * gw
+        patch_tokens = tokens[:, -n_patches:, :]
+
+        # Flatten and normalize
+        feats = (
+            patch_tokens.reshape(B * n_patches, C)
+            .detach()
+            .cpu()
+            .numpy()
+            .astype("float32")
+        )
+        feats = self._l2_normalize(feats)
+
+        grid_size = (gh, gw)
+        return feats, grid_size
+
+    def _compute_patch_mask(self,image,grid_size) :
+        """
+        Convert a pixel-level mask to patch-level mask.
+        """
+        h_p, w_p = grid_size
+        n_patches = h_p * w_p
+
+        if self.segmenter is None:
+            return np.ones(n_patches, dtype=bool)
+
+        # Resize image same way as in the backbone before sending to SAM
+        pil_resized, np_resized = self._resize_for_model(image)
+        full_mask = self.segmenter.get_object_mask(np_resized) 
+
+        # Optionally visualize in resized space
+        visualize_segmentation(
+            np_resized,
+            full_mask,
+            grid_size=None,
+            title=f"Segmentation debug (resized {self.img_size})",
+        )
+
+        full_mask_uint8 = (full_mask.astype(np.uint8) * 255).astype(np.float32)
+
+        # Downsample to patch grid with area interpolation for coverage
+        mask_small = cv2.resize(
+            full_mask_uint8,
+            (w_p, h_p),
+            interpolation=cv2.INTER_AREA,
+        ) / 255.0
+
+        patch_mask = (mask_small >= 0.5).reshape(-1)
+
+        # Fallback if mask collapses
+        fg_ratio = patch_mask.mean()
+        if fg_ratio < 0.01 or fg_ratio > 0.99:
+            patch_mask = np.ones(n_patches, dtype=bool)
+
+        return patch_mask
+
+    def _resize_for_model(self, image):
+        pil = Image.fromarray(image)
+        pil_resized = pil.resize((self.img_size, self.img_size), Image.BICUBIC)
+        np_resized = np.array(pil_resized)
+        return pil_resized, np_resized
+
+    def _knn_distances(self, feats: np.ndarray) -> np.ndarray:
+        """
+        Compute distance of each query feature to its nearest neighbors in the memory bank.
+
+        """
+        if self.memory_bank is None:
+            raise RuntimeError("Memory bank is empty.")
+
+        
+        a = feats 
+        b = self.memory_bank 
+
+        # vectorize version of L2 distances 
+        a2 = np.sum(a**2, axis=1, keepdims=True)        
+        b2 = np.sum(b**2, axis=1, keepdims=True).T    
+        ab = a @ b.T                                   
+        # Clip to avoid negative values
+        d2 = np.clip(a2 + b2 - 2.0 * ab, a_min=0.0, a_max=None)
+        d = np.sqrt(d2) 
+
+        # kNN: take mean of k smallest distances per patch
+        k = min(self.k_neighbors, d.shape[1])
+        if k == 1:
+            min_d = d.min(axis=1)
+        else:
+            # partial sort for efficiency
+            part = np.partition(d, kth=k - 1, axis=1)[:, :k]
+            min_d = part.mean(axis=1)
+
+        return min_d.astype(np.float32)
+
+    @staticmethod
+    def _mean_top_percent(values: np.ndarray, top_percent: float = 1.0) -> float:
+        """Mean of top p% values used as image level anomaly score."""
+        if values.size == 0:
+            return 0.0
+        k = max(1, int(round(values.size * (top_percent / 100.0))))
+        part = np.partition(values, -k)[-k:]
+        return float(part.mean())

segmenters/__init__.py

@@ -0,0 +1,21 @@
+from segmenters.base_segmenter import BaseSegmenter
+
+
+from segmenters.sam3 import SAM3Segmenter
+
+from segmenters.pca_segmenter import PCASegmenter
+
+"""
+Model registry for segmenters
+"""
+_SEGMENTERS = {}
+if SAM3Segmenter is not None:
+    _SEGMENTERS["sam3"] = SAM3Segmenter
+if PCASegmenter is not None:
+    _SEGMENTERS["pca"] = PCASegmenter
+
+
+def get_segmenter(name: str, **kwargs):
+    if name not in _SEGMENTERS:
+        raise ValueError(f"Unknown segmenter '{name}'. Available: {list(_SEGMENTERS)}")
+    return _SEGMENTERS[name](**kwargs)

segmenters/base_segmenter.py

@@ -0,0 +1,14 @@
+
+import numpy as np
+
+class BaseSegmenter:
+    """Base class for segmentation models """
+    def get_object_mask(self, image: np.ndarray) -> np.ndarray:
+        """
+        Args:
+            image
+        Returns:
+            bool mask of shape, where True = foreground object.
+        """
+        raise NotImplementedError
+

segmenters/pca_segmenter.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+import numpy as np
+import torch
+import cv2
+from PIL import Image
+import torchvision.transforms as T
+from timm.data import resolve_data_config
+
+from backbones import get_backbone
+from segmenters import BaseSegmenter
+
+
+class PCASegmenter(BaseSegmenter):
+    def __init__(
+        self,
+        backbone_name: str = "dinov3_base",
+        device: str | None = None,
+        threshold: float = 2.5,
+        kernel_size: int = 5,
+        border: float = 0.2,
+    ):
+        super().__init__()
+        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        self.model = get_backbone(backbone_name).to(self.device)
+        self.model.eval()
+
+
+        cfg = resolve_data_config({}, model=self.model)
+        _, img_size, _ = cfg["input_size"]
+        arch = getattr(getattr(self.model, "pretrained_cfg", {}), "get", lambda k, d=None: {})(  # type: ignore[arg-type]
+            "architecture", ""
+        )
+        if isinstance(arch, str) and "dinov3" in arch:
+            img_size = max(img_size, 512)
+
+        self.img_size = img_size
+        interp = cfg.get("interpolation", "bicubic")
+        self.transform = T.Compose(
+            [
+                T.Resize((self.img_size, self.img_size), interpolation=getattr(T.InterpolationMode, interp.upper(), T.InterpolationMode.BICUBIC)),
+                T.ToTensor(),
+                T.Normalize(mean=cfg.get("mean", (0.485, 0.456, 0.406)), std=cfg.get("std", (0.229, 0.224, 0.225))),
+            ]
+        )
+        self.threshold = threshold
+        self.border = border
+        self.kernel = np.ones((kernel_size, kernel_size), dtype=np.uint8)
+
+    def get_object_mask(self, image: np.ndarray) -> np.ndarray:
+        h0, w0 = image.shape[:2]
+        pil = Image.fromarray(image.astype(np.uint8))
+        x = self.transform(pil).unsqueeze(0).to(self.device)
+
+        with torch.inference_mode():
+            out = self.model.forward_features(x)
+
+        tokens = out.get("x_norm_patchtokens") if isinstance(out, dict) else out
+
+        if tokens is None and isinstance(out, dict):
+            tokens = out.get("x")
+        if tokens is not None and tokens.ndim == 4:
+            B, C, Hf, Wf = tokens.shape
+            tokens = tokens.permute(0, 2, 3, 1).reshape(B, Hf * Wf, C)
+
+        gh_dyn = int(np.sqrt(tokens.shape[1]))
+        gw_dyn = max(1, tokens.shape[1] // max(1, gh_dyn))
+        gh, gw = gh_dyn, gw_dyn
+
+        if hasattr(self.model, "patch_embed") and hasattr(self.model.patch_embed, "grid_size"):
+            gh0, gw0 = self.model.patch_embed.grid_size
+            if gh0 * gw0 == tokens.shape[1]:
+                gh, gw = gh0, gw0
+        n_patches = gh * gw
+        tokens = tokens[:, -n_patches:, :]
+
+        feats = tokens.squeeze(0).detach().cpu().numpy().astype(np.float32)
+        feats -= feats.mean(0, keepdims=True)
+        u, s, vh = np.linalg.svd(feats, full_matrices=False)
+        pc1 = vh[0]
+        scores = feats @ pc1
+        mask = scores > self.threshold
+        m_grid = mask.reshape(gh, gw)
+        bh = int(gh * self.border)
+        bw = int(gw * self.border)
+        inner = m_grid[bh : gh - bh, bw : gw - bw]
+        if inner.size > 0 and inner.mean() <= 0.35:
+            mask = scores < -self.threshold
+            m_grid = mask.reshape(gh, gw)
+        mask = m_grid.astype(np.uint8)
+        mask = cv2.dilate(mask, self.kernel, iterations=1)
+        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, self.kernel)
+        mask = cv2.resize(mask, (w0, h0), interpolation=cv2.INTER_NEAREST)
+        return mask.astype(bool)

segmenters/sam2.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+import numpy as np
+import torch
+from PIL import Image
+from transformers import Sam2Processor, Sam2Model
+
+from segmenters import BaseSegmenter
+
+
+class SAM2Segmenter(BaseSegmenter):
+    """
+    SAM2 wrapper.
+
+    - Uses Sam2Model (e.g. `facebook/sam2.1-hiera-large`).
+    - Segments (approximately) all objects in the image by prompting
+      with a full-image bounding box and returns a single boolean mask
+      given by the union of all predicted masks.
+    """
+
+    def __init__(
+        self,
+        text_prompt: str | None = None,
+        model_name: str = "facebook/sam2.1-hiera-large",
+        device: str = "cuda",
+        mask_threshold: float = 0.5,
+    ) -> None:
+        """
+        Args:
+            text_prompt: kept for compatibility with SAM3Segmenter, but unused.
+            model_name: HF repo id for the SAM2 model, e.g. "facebook/sam2.1-hiera-large".
+            device: "cuda" or "cpu".
+            mask_threshold: pixel threshold for masks (after SAM2 post-processing).
+        """
+        super().__init__()
+
+        if torch.cuda.is_available() and device.startswith("cuda"):
+            self.device = torch.device(device)
+        else:
+            self.device = torch.device("cpu")
+
+        # Load SAM2 model + processor
+        self.model = Sam2Model.from_pretrained(model_name).to(self.device)
+        self.model.eval()
+        self.processor = Sam2Processor.from_pretrained(model_name)
+
+    def get_object_mask(self, image: np.ndarray) -> np.ndarray:
+        """
+        Run SAM2 and return a single foreground mask.
+
+        - Convert image to PIL.
+        - Use a single bounding box covering the whole image as prompt.
+        - Run SAM2, post-process masks to image resolution.
+        - Threshold and union all masks into one boolean (H, W) array.
+
+        """
+        # Ensure PIL image
+        if isinstance(image, np.ndarray):
+            pil_image = Image.fromarray(image.astype(np.uint8)).convert("RGB")
+        else:
+            pil_image = image
+
+        W, H = pil_image.size  # PIL: (W, H)
+
+        # Full image bounding box: [x_min, y_min, x_max, y_max]
+        input_boxes = [[[0, 0, W, H]]]
+
+        # Build inputs for SAM2
+        inputs = self.processor(
+            images=pil_image,
+            input_boxes=input_boxes,
+            return_tensors="pt",
+        ).to(self.device)
+
+        with torch.no_grad():
+            # multimask_output=False → one mask per box
+            outputs = self.model(**inputs, multimask_output=False)
+
+        # Post-process masks to original resolution
+        masks = self.processor.post_process_masks(
+            outputs.pred_masks.cpu(),  # (B, num_masks, H', W')
+            inputs["original_sizes"],
+        )[0]
+
+        # Shapes can be:
+        # - (num_masks, H, W)
+        # - or (1, num_masks, H, W) depending on version
+        if masks.ndim == 4:
+            # (B, num_masks, H, W) -> (num_masks, H, W) for B=1
+            masks = masks[0]
+
+        if masks.ndim == 2:
+            # Single mask: (H, W)
+            full_mask = (masks > self.mask_threshold).numpy().astype(bool)
+            return full_mask
+
+        if masks.ndim != 3:
+            # Failsafe: if something weird happens, keep everything
+            return np.ones((H, W), dtype=bool)
+
+        # masks: (num_masks, H, W)
+        masks_bin = masks > self.mask_threshold
+        combined = masks_bin.any(dim=0)  # (H, W)
+        full_mask = combined.numpy().astype(bool)
+
+        return full_mask

segmenters/sam3.py

@@ -0,0 +1,97 @@
+
+from __future__ import annotations
+import numpy as np
+import torch
+from PIL import Image
+from transformers import Sam3Processor, Sam3Model
+
+from segmenters import BaseSegmenter
+
+
+class SAM3Segmenter(BaseSegmenter):
+    """
+    SAM3 wrapper using a text prompt of object type
+    """
+
+    def __init__(
+        self,
+        text_prompt: str,
+        model_name: str = "facebook/sam3",
+        device: str = "cuda",
+        score_threshold: float = 0.5,
+        mask_threshold: float = 0.5 ):
+        """
+        Args:
+            text_prompt: stuff we want to segment.
+            model_name: HF repo id for the SAM3 model.
+            device: "cuda" or "cpu".
+            score_threshold: min detection score to keep an instance.
+            mask_threshold: pixel threshold for masks.
+        """
+        super().__init__()
+
+        if torch.cuda.is_available() and device.startswith("cuda"):
+            self.device = torch.device(device)
+        else:
+            self.device = torch.device("cpu")
+
+        # preprocess text prompt so metal_nut is processed as metal nut
+        self.text_prompt = text_prompt.replace("_", " ")
+        self.score_threshold = score_threshold
+        self.mask_threshold = mask_threshold
+
+        # Loading model model + defining processor
+        self.model = Sam3Model.from_pretrained(model_name).to(self.device)
+        self.model.eval()
+        self.processor = Sam3Processor.from_pretrained(model_name)
+
+    def get_object_mask(self, image: np.ndarray) -> np.ndarray:
+        """
+        Running SAM3 and returning a single foreground mask.
+        """
+        # Pill image stuff - probably there is less idiotic way, but it is wat ChatGPT suggested
+        if isinstance(image, np.ndarray):
+            pil_image = Image.fromarray(image.astype(np.uint8)).convert("RGB")
+        else:
+            pil_image = image
+
+        # defining preprocessor with text prompt
+        inputs = self.processor(
+            images=pil_image,
+            text=self.text_prompt,
+            return_tensors="pt",
+        ).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+
+        # Post-process instance segmentation 
+        target_sizes = inputs.get("original_sizes").tolist()
+        results = self.processor.post_process_instance_segmentation(
+            outputs,
+            threshold=self.score_threshold,
+            mask_threshold=self.mask_threshold,
+            target_sizes=target_sizes,
+        )[0]
+
+        masks = results.get("masks", None)
+        scores = results.get("scores", None)
+
+        # If SAM completely fails we keep everything
+        if masks is None or masks.numel() == 0:
+            H, W = pil_image.size[1], pil_image.size[0] 
+            return np.ones((H, W), dtype=bool)
+
+        if scores is not None:
+            keep = scores >= self.score_threshold
+            if keep.sum() == 0:
+                H, W = pil_image.size[1], pil_image.size[0]
+                return np.ones((H, W), dtype=bool)
+            masks = masks[keep]
+
+        # check if mask passes mask treshold
+        masks_bin = (masks > self.mask_threshold)
+        combined = masks_bin.any(dim=0)
+        full_mask = combined.cpu().numpy().astype(bool)
+
+        return full_mask

utils/visualize.py

@@ -0,0 +1,114 @@
+
+import numpy as np
+import matplotlib.pyplot as plt
+import cv2
+
+#Completeöy vibecoded
+def visualize_prediction(
+    image: np.ndarray,
+    anomaly_map: np.ndarray,
+    image_score: float,
+    threshold_percentile: float = 95.0,
+    title: str | None = None,
+) -> None:
+    """
+    Show:
+        - original image
+        - heatmap overlay
+        - binary mask overlay (thresholded)
+    """
+    # Normalize anomaly map to [0, 1] for visualization
+    amap = anomaly_map.astype(np.float32)
+    amap -= amap.min()
+    if amap.max() > 0:
+        amap /= amap.max()
+
+    thresh = np.percentile(amap, threshold_percentile)
+    binary = amap >= thresh
+
+    fig, axes = plt.subplots(1, 3, figsize=(14, 4))
+
+    axes[0].imshow(image)
+    axes[0].set_title("Input image")
+    axes[0].axis("off")
+
+    axes[1].imshow(image)
+    im = axes[1].imshow(amap, cmap="jet", alpha=0.5)
+    axes[1].set_title("Anomaly heatmap")
+    axes[1].axis("off")
+    fig.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04)
+
+    axes[2].imshow(image)
+    axes[2].imshow(binary, cmap="gray", alpha=0.5)
+    axes[2].set_title(f"Thresholded (>{threshold_percentile:.0f}%)")
+    axes[2].axis("off")
+
+    if title is None:
+        title = f"Image anomaly score: {image_score:.3f}"
+
+    fig.suptitle(title)
+    plt.tight_layout()
+    plt.show()
+
+def visualize_segmentation(
+    image: np.ndarray,
+    full_mask: np.ndarray,
+    grid_size: tuple[int, int] | None = None,
+    title: str | None = None,
+) -> None:
+    """
+    Visualize SAM2 segmentation.
+
+    Args:
+        image: (H, W, 3) RGB uint8
+        full_mask: (H, W) bool or 0/1 array from SAM2
+        grid_size: optional (H_patches, W_patches) to also show patch-level mask
+        title: optional title string
+    """
+    img = image
+    mask = full_mask.astype(bool)
+    H, W = mask.shape
+
+    # Prepare figure layout
+    n_cols = 3 if grid_size is None else 4
+    fig, axes = plt.subplots(1, n_cols, figsize=(4 * n_cols, 4))
+
+    # 1) input image
+    axes[0].imshow(img)
+    axes[0].set_title("Input image")
+    axes[0].axis("off")
+
+    # 2) raw binary mask
+    axes[1].imshow(mask, cmap="gray")
+    axes[1].set_title("SAM2 mask (full-res)")
+    axes[1].axis("off")
+
+    # 3) overlay mask on image
+    axes[2].imshow(img)
+    axes[2].imshow(mask, cmap="Reds", alpha=0.4)
+    axes[2].set_title("Mask overlay")
+    axes[2].axis("off")
+
+    # 4) optional patch-level mask (after downsampling)
+    if grid_size is not None:
+        gh, gw = grid_size
+        # downsample full mask to patch grid and back up to image size
+        patch_mask_small = cv2.resize(
+            mask.astype(np.uint8), (gw, gh), interpolation=cv2.INTER_NEAREST
+        ).astype(bool)
+        patch_mask_full = cv2.resize(
+            patch_mask_small.astype(np.uint8),
+            (W, H),
+            interpolation=cv2.INTER_NEAREST,
+        ).astype(bool)
+
+        axes[3].imshow(img)
+        axes[3].imshow(patch_mask_full, cmap="Blues", alpha=0.4)
+        axes[3].set_title("Patch-level mask (after downsample)")
+        axes[3].axis("off")
+
+    if title is not None:
+        fig.suptitle(title)
+
+    plt.tight_layout()
+    plt.show()