hostabee/marfik-controle-qualite

1"""2Détection d'anomalies visuelles — PatchCore simplifié (Roth et al., 2022).3 4Principe : on n'a besoin QUE d'images de pièces conformes. Un réseau de vision5pré-entraîné (Wide-ResNet-50, torchvision, licence BSD) extrait des descripteurs6locaux ("patches") ; on les stocke dans une banque mémoire. Une pièce à7contrôler est comparée patch par patch : plus la distance au plus proche voisin8est grande, plus la zone est anormale. On obtient un score global + une carte9de chaleur localisant le défaut. Fonctionne sur CPU.10"""11from __future__ import annotations12 13from pathlib import Path14 15import numpy as np16import torch17import torch.nn.functional as F18from PIL import Image19from torchvision import transforms20from torchvision.models import Wide_ResNet50_2_Weights, wide_resnet50_221 22IMG = 22423TF = transforms.Compose([24    transforms.Resize((IMG, IMG)),25    transforms.ToTensor(),26    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),27])28 29 30class PatchCore:31    def __init__(self, device: str = "cpu"):32        self.device = device33        net = wide_resnet50_2(weights=Wide_ResNet50_2_Weights.IMAGENET1K_V1).eval().to(device)34        self.stem = torch.nn.Sequential(net.conv1, net.bn1, net.relu, net.maxpool, net.layer1)35        self.l2, self.l3 = net.layer2, net.layer336        self.bank: torch.Tensor | None = None37        self.threshold: float | None = None38        self.calib_scores: list[float] = []39 40    @torch.no_grad()41    def _features(self, x: torch.Tensor) -> torch.Tensor:42        f2 = self.l2(self.stem(x))            # (B, 512, 28, 28)43        f3 = self.l3(f2)                      # (B, 1024, 14, 14)44        f3 = F.interpolate(f3, size=f2.shape[-2:], mode="bilinear", align_corners=False)45        f = torch.cat([f2, f3], 1)46        f = F.avg_pool2d(f, 3, 1, 1)          # agrégation locale47        return f                              # (B, 1536, 28, 28)48 49    def _load(self, path: Path) -> torch.Tensor:50        return TF(Image.open(path).convert("RGB")).unsqueeze(0).to(self.device)51 52    def fit(self, good_images: list[Path], sampling: float = 0.1) -> None:53        feats = []54        for p in good_images:55            f = self._features(self._load(p))56            feats.append(f.permute(0, 2, 3, 1).reshape(-1, f.shape[1]))57        bank = torch.cat(feats)58        # sous-échantillonnage aléatoire (le coreset greedy est trop lent sur CPU)59        n = max(2000, int(len(bank) * sampling))60        idx = torch.randperm(len(bank))[:n]61        self.bank = bank[idx]62        # seuil : sur les images conformes, score max + marge63        self.calib_scores = [self.score(p)[0] for p in good_images[:20]]64        mu, sd = float(np.mean(self.calib_scores)), float(np.std(self.calib_scores))65        self.threshold = mu + 3 * sd66 67    @torch.no_grad()68    def score(self, path: Path) -> tuple[float, np.ndarray]:69        f = self._features(self._load(path))70        h, w = f.shape[-2:]71        patches = f.permute(0, 2, 3, 1).reshape(-1, f.shape[1])72        d = torch.cdist(patches, self.bank).min(dim=1).values      # plus proche voisin73        amap = d.reshape(h, w)74        amap = F.interpolate(amap[None, None], size=(IMG, IMG), mode="bilinear", align_corners=False)[0, 0]75        amap = torch.from_numpy(_blur(amap.numpy()))76        return float(amap.max()), amap.numpy()77 78 79def _blur(a: np.ndarray, k: int = 9) -> np.ndarray:80    from scipy.ndimage import gaussian_filter81    return gaussian_filter(a, sigma=k / 3)82