[Deep Learning] 이미지 세그멘테이션, U-Net 정리

---

1. U-Net이란?

• 배경: 2015년 Olaf Ronneberger 연구팀이 제안.
• 특징: 적은 양의 데이터로도 매우 정교한 세그멘테이션이 가능함.
• 용도: 의료 영상(세포/장기 검출), 자율 주행, 위성 이미지 분석, 심박 그래프 추출 등.

---

2. U-Net의 구조: 왜 'U'인가?

① 인코더 (Contracting Path - 수축 경로)

• 이미지의 전체적인 맥락(Context)을 파악하는 단계입니다.
• 일반적인 CNN처럼 컨볼루션과 풀링을 반복하며 이미지 크기는 줄이고 채널(특징)은 늘립니다.

② 디코더 (Expanding Path - 확장 경로)

• 줄어든 이미지를 다시 원래 크기로 복원하며 정확한 위치(Localization) 정보를 찾아냅니다.
• Up-convolution을 통해 해상도를 높입니다.

③ 스킵 커넥션 (Skip Connection)

• 인코더의 고해상도 특징 맵을 디코더의 상응하는 층에 직접 이어붙이는(Concatenate) 방식입니다.
• 효과: 인코더가 압축 과정에서 잃어버린 미세한 위치 정보를 디코더가 다시 참고하게 하여, 얇은 선이나 복잡한 경계도 정교하게 복원합니다.

---

3. PyTorch 핵심 코드 구현

1) 모델 정의 (UNet.py)

import torch
import torch.nn as nn

class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()

        def conv_block(in_c, out_c):
            return nn.Sequential(
                nn.Conv2d(in_c, out_c, kernel_size=3, padding=1),
                nn.BatchNorm2d(out_c),
                nn.ReLU(inplace=True),
                nn.Conv2d(out_c, out_c, kernel_size=3, padding=1),
                nn.BatchNorm2d(out_c),
                nn.ReLU(inplace=True)
            )

        # Encoder (Downsampling)
        self.enc1 = conv_block(in_channels, 64)
        self.pool1 = nn.MaxPool2d(2)
        self.enc2 = conv_block(64, 128)
        self.pool2 = nn.MaxPool2d(2)
        self.enc3 = conv_block(128, 256)
        self.pool3 = nn.MaxPool2d(2)
        self.enc4 = conv_block(256, 512)
        self.pool4 = nn.MaxPool2d(2)

        # Bottleneck
        self.bottleneck = conv_block(512, 1024)

        # Decoder (Upsampling)
        self.up4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.dec4 = conv_block(1024, 512) # Skip Connection 반영
        self.up3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.dec3 = conv_block(512, 256)
        self.up2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.dec2 = conv_block(256, 128)
        self.up1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.dec1 = conv_block(128, 64)

        self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool1(e1))
        e3 = self.enc3(self.pool2(e2))
        e4 = self.enc4(self.pool3(e3))

        b = self.bottleneck(self.pool4(e4))

        d4 = self.up4(b)
        d4 = torch.cat((d4, e4), dim=1) # 스킵 커넥션 결합
        d4 = self.dec4(d4)

        d3 = self.up3(d4)
        d3 = torch.cat((d3, e3), dim=1)
        d3 = self.dec3(d3)

        d2 = self.up2(d3)
        d2 = torch.cat((d2, e2), dim=1)
        d2 = self.dec2(d2)

        d1 = self.up1(d2)
        d1 = torch.cat((d1, e1), dim=1)
        d1 = self.dec1(d1)

        return self.final_conv(d1)

2) 학습 루프 (Training)

# 모델 및 하이퍼파라미터 설정
model = UNet(in_channels=3, out_channels=1).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.BCEWithLogitsLoss() # 이진 분류를 위한 손실 함수

for epoch in range(num_epochs):
    model.train()
    for images, masks in train_loader:
        images, masks = images.to(device), masks.to(device)

        # 순전파 (Forward)
        outputs = model(images)
        loss = criterion(outputs, masks)

        # 역전파 및 최적화 (Backward)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

---

4. 실전 활용: 심박 그래프 추출

U-Net을 활용하면 복잡한 배경(격자, 노이즈)이 있는 심박지 이미지에서 그래프 선만 깨끗하게 추출할 수 있습니다.

• 핵심 원리:
  • 인코더가 격자와 그래프의 형태학적 특징을 구분하는 법을 학습합니다.
  • 스킵 커넥션이 1픽셀 수준의 아주 얇은 그래프 선 정보를 디코더에 전달하여 뭉개짐 없이 복원합니다.
• Output 처리: 모델이 내뱉은 확률값을 임계값(Threshold, 0.5)으로 끊어주면 검은색 바탕에 흰색 줄만 남은 결과물을 얻게 됩니다.

---

5. 마치며

U-Net은 단순히 채널을 줄이고 늘리는 구조를 넘어, 스킵 커넥션이라는 혁신적인 아이디어로 세그멘테이션의 정밀도를 한 차원 끌어올린 모델입니다.