pip install segmentation-models-pytorch albumentations opencv-python-headless

import segmentation_models_pytorch as smp

# UNet + ResNet34 백본 (ImageNet 사전학습)
model = smp.Unet(
    encoder_name='resnet34',
    encoder_weights='imagenet',
    in_channels=3,
    classes=3,  # 클래스 수 (배경 포함)
)

print(f"파라미터 수: {sum(p.numel() for p in model.parameters()):,}")

# DeepLabV3+ (ASPP + Decoder)
model = smp.DeepLabV3Plus(
    encoder_name='resnet50',
    encoder_weights='imagenet',
    in_channels=3,
    classes=3,
)

import torch
from torch.utils.data import Dataset, DataLoader
import cv2
import numpy as np
import albumentations as A

class SegmentationDataset(Dataset):
    """세그멘테이션 데이터셋입니다."""

    def __init__(self, image_paths, mask_paths, transform=None):
        self.image_paths = image_paths
        self.mask_paths = mask_paths
        self.transform = transform

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

    def __getitem__(self, idx):
        image = cv2.imread(self.image_paths[idx])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.mask_paths[idx], cv2.IMREAD_GRAYSCALE)

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

        # 정규화 및 텐서 변환
        image = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        mask = torch.from_numpy(mask).long()

        return image, mask

# 증강 파이프라인
train_transform = A.Compose([
    A.Resize(256, 256),
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.2),
    A.RandomBrightnessContrast(p=0.3),
])

import torch.nn as nn
import torch.optim as optim

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Dice Loss + CrossEntropy 조합
criterion = smp.losses.DiceLoss(mode='multiclass')
ce_loss = nn.CrossEntropyLoss()

optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50)

num_epochs = 50
for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0.0

    for images, masks in train_loader:
        images = images.to(device)
        masks = masks.to(device)

        optimizer.zero_grad()
        outputs = model(images)  # [B, C, H, W]

        loss = criterion(outputs, masks) + ce_loss(outputs, masks)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

    scheduler.step()

    # 검증
    model.eval()
    val_iou = 0.0
    with torch.no_grad():
        for images, masks in val_loader:
            images = images.to(device)
            masks = masks.to(device)
            outputs = model(images)
            preds = outputs.argmax(dim=1)
            val_iou += compute_iou(preds, masks)

    print(f"Epoch {epoch+1}/{num_epochs} | Loss: {epoch_loss/len(train_loader):.4f} | Val mIoU: {val_iou/len(val_loader):.4f}")

import matplotlib.pyplot as plt

def predict_and_visualize(model, image_path, device):
    """이미지에 대해 세그멘테이션을 수행하고 결과를 시각화합니다."""
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    resized = cv2.resize(image, (256, 256))

    input_tensor = torch.from_numpy(resized).permute(2, 0, 1).float() / 255.0
    input_tensor = input_tensor.unsqueeze(0).to(device)

    model.eval()
    with torch.no_grad():
        output = model(input_tensor)
        pred_mask = output.argmax(dim=1).squeeze().cpu().numpy()

    # 시각화
    fig, axes = plt.subplots(1, 2, figsize=(10, 5))
    axes[0].imshow(resized)
    axes[0].set_title('원본')
    axes[1].imshow(pred_mask, cmap='jet')
    axes[1].set_title('예측 마스크')
    plt.tight_layout()
    plt.show()