File size: 2,435 Bytes
4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 6742b9c 4c775f9 | 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 | import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F # for Peak Score logic
from tqdm import tqdm
import numpy as np
# Added 'device' argument to match your pipeline.py call
def train_model(model, train_loader, epochs=10, device=None):
# Use passed device, or fall back to auto-detection
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
optimizer = optim.AdamW(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()
for epoch in range(epochs):
model.train()
total_loss = 0
for images, _ in tqdm(train_loader, desc=f"Epoch {epoch+1}"):
images = images.to(device)
optimizer.zero_grad()
reconstructed = model(images)
loss = criterion(reconstructed, images)
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Epoch {epoch+1} complete. Avg Loss: {total_loss/len(train_loader):.6f}")
# Added 'device' argument to match your pipeline.py call
def evaluate_anomaly(model, test_loader, device=None):
if device is None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
errors = []
labels = []
print("Computing anomaly scores (Peak Score Method)...")
with torch.no_grad():
for images, label in tqdm(test_loader):
images = images.to(device)
reconstructed = model(images)
# 1. Calculate Squared Error map [Batch, C, H, W]
diff = (images - reconstructed)**2
# 2. Convert to Grayscale (Mean across channels) [Batch, 1, H, W]
diff_map = torch.mean(diff, dim=1, keepdim=True)
# 3. Peak Score Logic: Apply a 15x15 blur (Average Pooling)
# Aggregates the 116p defect signal and reduces single-pixel noise.
smoothed_diff = F.avg_pool2d(diff_map, kernel_size=15, stride=1, padding=7)
# 4. Get the Max value for each image in the batch (Hotspot detection)
batch_error, _ = torch.max(smoothed_diff.view(images.size(0), -1), dim=1)
errors.extend(batch_error.cpu().numpy())
labels.extend(label.numpy())
return errors, labels |