product-defect-classification/
├── data/
│   ├── train/
│   │   ├── good/          # 양품 이미지
│   │   └── defect/        # 불량 이미지
│   └── val/
│       ├── good/
│       └── defect/
├── outputs/
│   ├── checkpoints/       # 모델 체크포인트
│   └── results/           # 평가 결과
├── train.py
├── evaluate.py
└── config.py

# config.py — 프로젝트 설정

from dataclasses import dataclass

@dataclass
class Config:
    """프로젝트 설정을 정의합니다."""
    # 데이터
    data_dir: str = 'data'
    image_size: int = 224
    num_classes: int = 2
    class_names: list = None

    # 모델
    model_name: str = 'efficientnet_b0'
    pretrained: bool = True

    # 학습
    batch_size: int = 32
    epochs: int = 30
    learning_rate: float = 1e-4
    weight_decay: float = 1e-4

    # 장치
    device: str = 'cuda'

    def __post_init__(self):
        self.class_names = self.class_names or ['good', 'defect']

cfg = Config()

import os
import torch
from torch.utils.data import DataLoader
from torchvision import datasets
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np
from PIL import Image

# 증강 파이프라인 정의
def get_transforms(phase='train'):
    """학습/검증용 증강 파이프라인을 반환합니다."""
    if phase == 'train':
        return A.Compose([
            A.Resize(256, 256),
            A.RandomCrop(224, 224),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(p=0.5),
            A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1, p=0.5),
            A.GaussNoise(var_limit=(10, 50), p=0.3),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2(),
        ])
    else:
        return A.Compose([
            A.Resize(224, 224),
            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ToTensorV2(),
        ])

# Albumentations을 적용하는 커스텀 Dataset
class ProductDataset(torch.utils.data.Dataset):
    """제품 이미지 데이터셋입니다."""

    def __init__(self, root_dir, transform=None):
        self.dataset = datasets.ImageFolder(root_dir)
        self.transform = transform

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        image, label = self.dataset[idx]
        image = np.array(image)

        if self.transform:
            augmented = self.transform(image=image)
            image = augmented['image']

        return image, label

# DataLoader 생성
train_dataset = ProductDataset('data/train', transform=get_transforms('train'))
val_dataset = ProductDataset('data/val', transform=get_transforms('val'))

train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=cfg.batch_size, shuffle=False, num_workers=4)

# 데이터 분포 확인
print(f"학습 데이터: {len(train_dataset)}장")
print(f"검증 데이터: {len(val_dataset)}장")
for cls_name, cls_idx in train_dataset.dataset.class_to_idx.items():
    count = sum(1 for _, l in train_dataset.dataset.samples if l == cls_idx)
    print(f"  {cls_name}: {count}장")

from torch.utils.data import WeightedRandomSampler

# 클래스별 가중치 계산
def get_class_weights(dataset):
    """클래스 불균형에 대한 가중치를 계산합니다."""
    targets = [label for _, label in dataset.dataset.samples]
    class_counts = np.bincount(targets)
    total = len(targets)
    weights = total / (len(class_counts) * class_counts)
    return torch.FloatTensor(weights)

class_weights = get_class_weights(train_dataset)
print(f"클래스 가중치: {class_weights}")
# 예: good=0.55, defect=4.5 (불량이 적으므로 높은 가중치)

# 방법 1: 가중 손실 함수
criterion = torch.nn.CrossEntropyLoss(weight=class_weights.to(cfg.device))

# 방법 2: WeightedRandomSampler (오버샘플링 효과)
sample_weights = [class_weights[label] for _, label in train_dataset.dataset.samples]
sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights))

train_loader = DataLoader(
    train_dataset,
    batch_size=cfg.batch_size,
    sampler=sampler,       # shuffle 대신 sampler 사용
    num_workers=4,
)

import timm
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

# 모델 생성
model = timm.create_model(cfg.model_name, pretrained=cfg.pretrained, num_classes=cfg.num_classes)
model = model.to(cfg.device)

# 옵티마이저와 스케줄러
optimizer = AdamW(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs, eta_min=1e-6)

# 학습 함수
def train_one_epoch(model, loader, criterion, optimizer, device):
    """한 에포크를 학습합니다."""
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += labels.size(0)

    return total_loss / total, correct / total

# 검증 함수
@torch.no_grad()
def evaluate(model, loader, criterion, device):
    """모델을 검증합니다."""
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    all_preds = []
    all_labels = []

    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)

        total_loss += loss.item() * images.size(0)
        _, predicted = outputs.max(1)
        correct += predicted.eq(labels).sum().item()
        total += labels.size(0)

        all_preds.extend(predicted.cpu().numpy())
        all_labels.extend(labels.cpu().numpy())

    return total_loss / total, correct / total, all_preds, all_labels

# 학습 루프
best_acc = 0
for epoch in range(cfg.epochs):
    train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, cfg.device)
    val_loss, val_acc, _, _ = evaluate(model, val_loader, criterion, cfg.device)
    scheduler.step()

    print(f"Epoch [{epoch+1}/{cfg.epochs}] "
          f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
          f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")

    # 최고 성능 모델 저장
    if val_acc > best_acc:
        best_acc = val_acc
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'val_acc': val_acc,
        }, 'outputs/checkpoints/best_model.pth')
        print(f"  ✓ 최고 성능 모델 저장 (Acc: {val_acc:.4f})")

import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix

# 최고 모델 로드
checkpoint = torch.load('outputs/checkpoints/best_model.pth')
model.load_state_dict(checkpoint['model_state_dict'])

# 상세 평가
_, _, preds, labels = evaluate(model, val_loader, criterion, cfg.device)

# 분류 리포트
print(classification_report(labels, preds, target_names=cfg.class_names))

# 혼동 행렬 시각화
cm = confusion_matrix(labels, preds)
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
            xticklabels=cfg.class_names, yticklabels=cfg.class_names)
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.title('Confusion Matrix')
plt.tight_layout()
plt.savefig('outputs/results/confusion_matrix.png', dpi=150)

from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
import cv2

# 마지막 Conv 레이어 확인
# EfficientNet의 경우 features의 마지막 블록
target_layers = [model.conv_head]  # timm EfficientNet

cam = GradCAM(model=model, target_layers=target_layers)

# 불량 샘플에 대한 Grad-CAM
def visualize_gradcam(image_path, model, cam, class_idx):
    """Grad-CAM 히트맵을 생성합니다."""
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img_resized = cv2.resize(img, (224, 224))
    rgb_img = img_resized / 255.0

    # 텐서 변환
    input_tensor = torch.from_numpy(img_resized).permute(2, 0, 1).float() / 255.0
    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
    input_tensor = ((input_tensor - mean) / std).unsqueeze(0)

    targets = [ClassifierOutputTarget(class_idx)]
    grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0, :]
    visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)

    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    axes[0].imshow(img_resized)
    axes[0].set_title('원본')
    axes[1].imshow(grayscale_cam, cmap='jet')
    axes[1].set_title('Grad-CAM')
    axes[2].imshow(visualization)
    axes[2].set_title('오버레이')
    for ax in axes:
        ax.axis('off')
    plt.tight_layout()
    plt.savefig('outputs/results/gradcam_defect.png', dpi=150)

# 불량 이미지에 대해 시각화
visualize_gradcam('data/val/defect/sample.jpg', model, cam, class_idx=1)

# ONNX 변환
model.eval()
dummy_input = torch.randn(1, 3, 224, 224).to(cfg.device)

torch.onnx.export(
    model, dummy_input, 'outputs/model.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
    opset_version=17,
)
print("ONNX 모델 내보내기 완료")

# ONNX Runtime 검증
import onnxruntime as ort
import numpy as np

session = ort.InferenceSession('outputs/model.onnx')
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {'input': test_input})
print(f"ONNX 출력 형태: {outputs[0].shape}")