| from pathlib import Path |
| import cv2 |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torchvision.transforms as T |
| from PIL import Image |
| from torchvision import models |
|
|
|
|
| class ParcelGuardPredictor: |
| def __init__(self, backbone_path, threshold=0.70, device=None): |
| self.device = torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu')) |
| self.threshold = float(threshold) |
|
|
| backbone = models.resnet18(weights=None) |
| self.extractor = nn.Sequential( |
| backbone.conv1, backbone.bn1, backbone.relu, backbone.maxpool, |
| backbone.layer1, backbone.layer2 |
| ) |
| state = torch.load(backbone_path, map_location='cpu', weights_only=True) |
| self.extractor.load_state_dict(state, strict=True) |
| self.extractor.to(self.device).eval() |
|
|
| self.transform = T.Compose([ |
| T.Resize((224, 224)), |
| T.ToTensor(), |
| T.Normalize(mean=[0.485, 0.456, 0.406], |
| std=[0.229, 0.224, 0.225]), |
| ]) |
|
|
| @torch.inference_mode() |
| def _features(self, image_path): |
| path = Path(image_path) |
| if not path.is_file(): |
| raise FileNotFoundError(f'Image not found: {path}') |
| image = Image.open(path).convert('RGB') |
| tensor = self.transform(image).unsqueeze(0).to(self.device) |
| return self.extractor(tensor)[0].cpu().numpy() |
|
|
| @staticmethod |
| def _distance(query_features, reference_features): |
| q = query_features.transpose(1, 2, 0) |
| r = reference_features.transpose(1, 2, 0) |
| q = q / (np.linalg.norm(q, axis=-1, keepdims=True) + 1e-8) |
| r = r / (np.linalg.norm(r, axis=-1, keepdims=True) + 1e-8) |
| return 1.0 - np.sum(q * r, axis=-1) |
|
|
| def predict(self, query_image, reference_images): |
| if not reference_images: |
| raise ValueError('Provide at least one reference image.') |
| query_features = self._features(query_image) |
| maps = [] |
| for ref_path in reference_images: |
| ref_features = self._features(ref_path) |
| maps.append(self._distance(query_features, ref_features)) |
| anomaly_map = np.min(np.stack(maps), axis=0) |
| score = float(np.percentile(anomaly_map, 95)) |
| prediction = int(score >= self.threshold) |
|
|
| image = cv2.imread(str(query_image)) |
| if image is None: |
| raise ValueError('OpenCV could not read the query image.') |
| heatmap = cv2.resize(anomaly_map.astype(np.float32), |
| (image.shape[1], image.shape[0])) |
| heatmap = cv2.normalize(heatmap, None, 0, 255, |
| cv2.NORM_MINMAX).astype(np.uint8) |
| color = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) |
| overlay = cv2.addWeighted(image, 0.60, color, 0.40, 0) |
|
|
| return { |
| 'anomaly_score_p95': score, |
| 'predicted_defect': prediction, |
| 'label': 'DEFECT' if prediction else 'NORMAL', |
| 'threshold': self.threshold, |
| 'anomaly_map': anomaly_map, |
| 'overlay_bgr': overlay, |
| } |
|
|
