Spaces:

morinousagi
/

pytorch-vit-defect-detect

Sleeping

pytorch-vit-defect-detect / src /train_eval.py

Upload 4 files

6742b9c verified 6 months ago

2.44 kB

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

	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