File size: 3,096 Bytes
1a54fc5 7c5d998 1a54fc5 7c5d998 1a54fc5 7c5d998 1a54fc5 7c5d998 1a54fc5 7c5d998 1a54fc5 7c5d998 1a54fc5 7c5d998 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | 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,
}
|