|
|
""" |
|
|
Deep SVDD Anomaly Detection Model |
|
|
Trained on CIFAR-10, CIFAR-100, and STL-10 |
|
|
""" |
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
from torchvision import transforms |
|
|
from PIL import Image |
|
|
import pickle |
|
|
import json |
|
|
from pathlib import Path |
|
|
|
|
|
|
|
|
class ResidualBlock(nn.Module): |
|
|
def __init__(self, in_ch: int, out_ch: int, stride: int = 1): |
|
|
super().__init__() |
|
|
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False) |
|
|
self.bn1 = nn.BatchNorm2d(out_ch) |
|
|
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, stride=1, padding=1, bias=False) |
|
|
self.bn2 = nn.BatchNorm2d(out_ch) |
|
|
|
|
|
self.shortcut = nn.Sequential() |
|
|
if stride != 1 or in_ch != out_ch: |
|
|
self.shortcut = nn.Sequential( |
|
|
nn.Conv2d(in_ch, out_ch, 1, stride=stride, bias=False), |
|
|
nn.BatchNorm2d(out_ch) |
|
|
) |
|
|
|
|
|
def forward(self, x): |
|
|
out = F.relu(self.bn1(self.conv1(x))) |
|
|
out = self.bn2(self.conv2(out)) |
|
|
out += self.shortcut(x) |
|
|
return F.relu(out) |
|
|
|
|
|
|
|
|
class DeepSVDDEncoder(nn.Module): |
|
|
def __init__(self, latent_dim: int = 512): |
|
|
super().__init__() |
|
|
self.conv1 = nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False) |
|
|
self.bn1 = nn.BatchNorm2d(64) |
|
|
self.layer1 = self._make_layer(64, 128, stride=2) |
|
|
self.layer2 = self._make_layer(128, 256, stride=2) |
|
|
self.layer3 = self._make_layer(256, 512, stride=2) |
|
|
self.layer4 = self._make_layer(512, 512, stride=2) |
|
|
self.fc = nn.Linear(512 * 4 * 4, latent_dim, bias=False) |
|
|
|
|
|
def _make_layer(self, in_ch: int, out_ch: int, stride: int = 1): |
|
|
return nn.Sequential( |
|
|
ResidualBlock(in_ch, out_ch, stride), |
|
|
ResidualBlock(out_ch, out_ch, 1) |
|
|
) |
|
|
|
|
|
def forward(self, x): |
|
|
x = F.relu(self.bn1(self.conv1(x))) |
|
|
x = self.layer1(x) |
|
|
x = self.layer2(x) |
|
|
x = self.layer3(x) |
|
|
x = self.layer4(x) |
|
|
x = x.view(x.size(0), -1) |
|
|
return self.fc(x) |
|
|
|
|
|
|
|
|
class DeepSVDDAnomalyDetector: |
|
|
""" |
|
|
Deep SVDD Anomaly Detection Model |
|
|
|
|
|
Usage: |
|
|
from model import DeepSVDDAnomalyDetector |
|
|
|
|
|
detector = DeepSVDDAnomalyDetector.from_pretrained('.') |
|
|
score, is_anomaly = detector.predict('image.jpg') |
|
|
""" |
|
|
|
|
|
def __init__(self, model_path, thresholds_path, config_path, device='cuda'): |
|
|
self.device = torch.device(device if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
|
|
|
with open(config_path, 'r') as f: |
|
|
self.config = json.load(f) |
|
|
|
|
|
|
|
|
checkpoint = torch.load(model_path, map_location=self.device) |
|
|
self.latent_dim = checkpoint['latent_dim'] |
|
|
self.center = checkpoint['center'].to(self.device) |
|
|
self.radius = checkpoint['radius'].item() |
|
|
|
|
|
self.encoder = DeepSVDDEncoder(self.latent_dim).to(self.device) |
|
|
self.encoder.load_state_dict(checkpoint['encoder_state_dict']) |
|
|
self.encoder.eval() |
|
|
|
|
|
|
|
|
with open(thresholds_path, 'rb') as f: |
|
|
thresholds = pickle.load(f) |
|
|
|
|
|
self.threshold_95 = thresholds['95th_percentile'] |
|
|
self.threshold_99 = thresholds['99th_percentile'] |
|
|
self.threshold_optimal = thresholds['optimal_f1'] |
|
|
self.threshold = self.threshold_optimal |
|
|
|
|
|
|
|
|
self.transform = transforms.Compose([ |
|
|
transforms.Resize((128, 128)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
|
|
]) |
|
|
|
|
|
@classmethod |
|
|
def from_pretrained(cls, model_path='.', device='cuda'): |
|
|
"""Load pretrained model from directory or HuggingFace Hub""" |
|
|
model_path = Path(model_path) |
|
|
return cls( |
|
|
model_path=model_path / 'deepsvdd_model.pth', |
|
|
thresholds_path=model_path / 'thresholds.pkl', |
|
|
config_path=model_path / 'config.json', |
|
|
device=device |
|
|
) |
|
|
|
|
|
def set_threshold(self, threshold_type='optimal'): |
|
|
"""Set threshold: 'optimal', '95th', or '99th'""" |
|
|
if threshold_type == 'optimal': |
|
|
self.threshold = self.threshold_optimal |
|
|
elif threshold_type == '95th': |
|
|
self.threshold = self.threshold_95 |
|
|
elif threshold_type == '99th': |
|
|
self.threshold = self.threshold_99 |
|
|
|
|
|
@torch.no_grad() |
|
|
def predict(self, image_path): |
|
|
"""Predict if image is anomaly""" |
|
|
if isinstance(image_path, (str, Path)): |
|
|
image = Image.open(image_path).convert('RGB') |
|
|
else: |
|
|
image = image_path |
|
|
|
|
|
image_tensor = self.transform(image).unsqueeze(0).to(self.device) |
|
|
embeddings = self.encoder(image_tensor) |
|
|
score = torch.sum((embeddings - self.center) ** 2, dim=1).item() |
|
|
is_anomaly = score > self.threshold |
|
|
|
|
|
return score, is_anomaly |
|
|
|
