딥러닝 기반 불량 검출

import torch import torch.nn as nn from torchvision import models, transforms class DefectClassifier(nn.Module): """제조 결함 분류기 (Transfer Learning)""" def __init__(self, num_classes: int, freeze_backbone: bool = True): super().__init__() # EfficientNet-B0 백본 (ImageNet 사전학습) self.backbone = models.efficientnet_b0(pretrained=True) # 백본 동결 (Feature Extraction 모드) if freeze_backbone: for param in self.backbone.parameters(): param.requires_grad = False # 분류기 교체 (제조 결함용) in_features = self.backbone.classifier[1].in_features self.backbone.classifier = nn.Sequential( nn.Dropout(p=0.3), nn.Linear(in_features, 256), nn.ReLU(), nn.Dropout(p=0.2), nn.Linear(256, num_classes) ) def forward(self, x): return self.backbone(x) # 데이터 전처리 파이프라인 train_transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 사용 예시 model = DefectClassifier(num_classes=5) # OK, Scratch, Dent, Stain, Crack device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device)

import torch.nn as nn # 가중 CrossEntropy def get_weighted_loss(class_counts): """클래스별 빈도 역수로 가중치 계산""" total = sum(class_counts) weights = [total / c for c in class_counts] weights = torch.FloatTensor(weights) weights = weights / weights.sum() * len(weights) return nn.CrossEntropyLoss(weight=weights) # 예: OK 9000장, Scratch 500장, Dent 300장 class_counts = [9000, 500, 300, 150, 50] criterion = get_weighted_loss(class_counts)

from torch.utils.data import DataLoader from tqdm import tqdm def train_epoch(model, dataloader, criterion, optimizer, device): """1 에포크 학습""" model.train() running_loss = 0.0 correct = 0 total = 0 for images, labels in tqdm(dataloader, desc="Training"): images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() return running_loss / len(dataloader), correct / total def validate(model, dataloader, criterion, device): """검증""" model.eval() running_loss = 0.0 all_preds, all_labels = [], [] with torch.no_grad(): for images, labels in dataloader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) running_loss += loss.item() _, predicted = outputs.max(1) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) accuracy = np.mean(np.array(all_preds) == np.array(all_labels)) return running_loss / len(dataloader), accuracy

class DefectInferenceEngine: """결함 분류 추론 엔진""" def __init__(self, model_path: str, class_names: list, device='cuda'): self.device = torch.device(device if torch.cuda.is_available() else 'cpu') self.class_names = class_names # 모델 로드 self.model = DefectClassifier(num_classes=len(class_names)) checkpoint = torch.load(model_path, map_location=self.device) self.model.load_state_dict(checkpoint['model_state_dict']) self.model.to(self.device) self.model.eval() def predict(self, image_path: str): """단일 이미지 추론""" image = Image.open(image_path).convert('RGB') input_tensor = self.transform(image).unsqueeze(0).to(self.device) with torch.no_grad(): outputs = self.model(input_tensor) probabilities = torch.softmax(outputs, dim=1)[0] predicted_idx = probabilities.argmax().item() return { 'predicted_class': self.class_names[predicted_idx], 'confidence': probabilities[predicted_idx].item(), 'is_defect': predicted_idx != 0 # OK가 0번 클래스 } # 사용 예시 engine = DefectInferenceEngine( model_path='best_model.pth', class_names=['OK', 'Scratch', 'Dent', 'Stain', 'Crack'] ) result = engine.predict('test_image.jpg') print(f"결과: {result['predicted_class']} ({result['confidence']*100:.1f}%)")