U-Net 의 이해

U-Net은 이미지 분할(Image Segmentation) 문제를 해결하기 위해 개발된 합성곱 신경망(CNN) 기반의 모델이다. 이 모델은 2015년에 Olaf Ronneberger와 그의 동료들에 의해 의료 영상 분석을 위해 제안되었으며, 주로 생물학적 이미지 분할에 사용되었다. U-Net은 그 단순하지만 강력한 구조로 인해 다양한 컴퓨터 비전 문제에 널리 사용되고 있다.

U-Net: Convolutional Networks for Biomedical Image Segmentation

 

U-Net의 등장 배경

U-Net은 주로 의료 영상 분석에서 사용되는 이미지 분할 문제를 해결하기 위해 개발되었다. 전통적인 이미지 분할 방법들은 복잡한 전처리나 특징 추출 과정을 거쳐야 했지만, U-Net은 이러한 과정 없이 이미지에서 중요한 정보를 효율적으로 학습할 수 있도록 설계되었다. 특히 의료 영상에서 병변이나 세포와 같은 작은 구조들을 정확하게 분할하는 것이 중요한데, U-Net은 이를 효과적으로 수행할 수 있다.

 

---

U-Net의 아키텍처

U-Net의 아키텍처는 U자 형태를 띠고 있어 그 이름이 붙여졌다. 이 아키텍처는 크게 인코더(Encoder)와 디코더(Decoder)로 나눌 수 있다. 인코더는 이미지를 점점 축소하면서 중요한 특징을 추출하고, 디코더는 이를 다시 원래의 해상도로 복원하여 분할된 이미지를 생성한다.

 

 

 

1. 인코더(Contracting Path)

인코더는 전형적인 CNN 구조를 사용하여 이미지의 특징을 추출한다. 각 단계에서 컨볼루션 층을 사용하여 이미지의 공간 정보를 축소하고, 최대 풀링(Max Pooling) 연산을 통해 해상도를 줄이면서 특징의 수를 증가시킨다. 이를 통해 이미지의 중요한 특징을 효과적으로 압축할 수 있다.

 

 

위의 부분만 코드를 작성을 하면 아래와 같다.

import torch
import torch.nn as nn
import torch.nn.functional as F

class InitialEncoder(nn.Module):
    def __init__(self, in_channels):
        super(InitialEncoder, self).__init__()
        # 첫 번째 인코딩 단계: 입력 이미지를 64개의 채널로 변환
        self.encoder1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3),  # 572x572 -> 570x570
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3),           # 570x570 -> 568x568
            nn.ReLU(inplace=True)
        )
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)  # 568x568 -> 284x284
    
    def forward(self, x):
        # 첫 번째 인코딩 블록
        e1 = self.encoder1(x)
        p1 = self.pool1(e1)
        return e1, p1

# 모델 테스트
x = torch.randn(1, 1, 572, 572)  # 입력 이미지 (1, 채널, 높이, 너비)
encoder = InitialEncoder(in_channels=1)
e1, p1 = encoder(x)

print("Encoder1 Output Shape:", e1.shape)  # (1, 64, 568, 568)
print("Pooled Output Shape:", p1.shape)    # (1, 64, 284, 284)

 

컨볼루션 레이어를 여러 번 거친 후에 이미지 크기를 줄이는 것은 중요한 특징을 더 잘 학습하기 위해서 의도된 설계이다.

구체적으로, 각 컨볼루션 레이어는 입력 이미지로부터 다양한 크기와 복잡도를 가진 특징을 추출해낸다. 여러 번의 3x3 컨볼루션을 거치는 동안 네트워크는 더욱 풍부한 로컬 정보를 학습하게 된다. 특히 첫 번째 컨볼루션 레이어는 저수준 특징(예: 경계, 질감)을 추출하고, 그 후의 레이어들은 점점 더 복잡하고 추상적인 특징을 학습하게 된다.

이 과정을 거친 후에 Max Pooling을 통해 공간 해상도를 줄임으로써, 네트워크는 더 큰 Receptive Field을 가지게 되고 이미지의 전반적인 구조나 맥락을 파악할 수 있게 된다. 이를 통해 모델은 작은 세부 사항뿐만 아니라 더 넓은 맥락에서 이미지의 중요한 특징을 학습하게 된다.

 

이미지 크기가 32가 될 때까지 인코더의 각 단계를 계속 수행하며, 인코더는 작은 세부 사항에서 더 추상적인 표현까지 다양한 정보를 학습하게 된다. 이후에는 디코더를 통해 원래 크기로 복원하면서 분할 결과를 얻을 수 있게 된다.

 

---

 

2. 디코더(Expansive Path)

디코더는 인코더에서 추출한 특징 맵을 사용하여 이미지를 다시 원래의 해상도로 복원하는 역할을 한다. 인코더에서는 이미지의 공간 정보를 점점 줄여가며 중요한 특징을 추출하게 되지만, 이 과정에서 해상도가 줄어들면서 공간적인 세부 정보가 손실된다. 디코더는 이러한 압축된 특징을 다시 복원하여 픽셀 단위로 이미지 분할을 수행할 수 있도록 돕는다.
즉, 디코더는 인코더에서 얻은 고수준의 추상적 특징을 이용해 이미지의 원래 구조를 점진적으로 복원하며, 분할된 출력 지도를 생성한다. 이를 통해 우리는 인풋 이미지에서 객체의 위치와 경계를 픽셀 단위로 식별할 수 있는 분할 결과를 얻을 수 있다.

 

"분할된 출력 지도"라는 표현은 이미지의 각 픽셀이 특정 클래스에 속하는지를 나타내는 출력 이미지를 의미한다.
예를 들어, 의료 영상에서 특정 병변을 분할한다고 할 때, U-Net의 출력은 입력 이미지와 같은 크기의 이미지를 생성하지만, 각 픽셀은 병변 영역에 속하는지 여부를 나타내는 값을 가지게 된다.

이렇게 생성된 출력 지도는 입력 이미지의 각 픽셀을 클래스별로 할당한 결과로, 예를 들어 0은 배경, 1은 병변을 나타낼 수 있다. 이 과정은 이미지 내에서 우리가 관심 있는 객체들을 정확히 분리해내는 데 사용된다. 이 결과는 이후에 의료 진단, 객체 탐지, 또는 다른 컴퓨터 비전 문제에서 활용될 수 있다.

 

디코더에서 사용하는 업샘플링(Upsampling)**과 스킵 연결(Skip Connection)은 복원을 더 정밀하게 수행하는 데 중요한 역할을 한다. 업샘플링을 통해 이미지의 해상도를 늘리고, 스킵 연결을 통해 인코더에서 얻은 고해상도 정보를 디코더에 전달함으로써 더 나은 복원 성능을 제공한다.

 

 

위의 이미지에서 인코더에서 디코더로 넘어간 부분을 코드로 보자.

import torch
import torch.nn as nn
import torch.nn.functional as F

class BottleneckDecoder(nn.Module):
    def __init__(self):
        super(BottleneckDecoder, self).__init__()
        
        # 중심부 정의
        self.bottleneck = nn.Sequential(
            nn.Conv2d(512, 1024, kernel_size=3, padding=1),  # 28x28 -> 28x28
            nn.ReLU(inplace=True),
            nn.Conv2d(1024, 1024, kernel_size=3, padding=1), # 28x28 -> 28x28
            nn.ReLU(inplace=True)
        )
        
        # 디코더 업샘플링 및 특징 복원
        self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)  # 28x28 -> 56x56
        self.decoder4 = nn.Sequential(
            nn.Conv2d(1024, 512, kernel_size=3, padding=1),  # 56x56 -> 56x56
            nn.ReLU(inplace=True),
            nn.Conv2d(512, 512, kernel_size=3, padding=1),   # 56x56 -> 56x56
            nn.ReLU(inplace=True)
        )
    
    def forward(self, x, skip):
        # 중심부 연산
        b = self.bottleneck(x)
        
        # 디코더의 첫 번째 업샘플링 및 특징 결합
        d4 = self.upconv4(b)  # 업샘플링
        d4 = torch.cat((d4, skip), dim=1)  # 스킵 연결
        d4 = self.decoder4(d4)  # 디코더 블록
        return d4

# 모델 테스트
x = torch.randn(1, 512, 28, 28)  # 인코더에서 내려온 마지막 특징 맵
skip = torch.randn(1, 512, 56, 56)  # 인코더에서 스킵된 특징 맵

bottleneck_decoder = BottleneckDecoder()
output = bottleneck_decoder(x, skip)

print("Bottleneck and Decoder Output Shape:", output.shape)  # 예상 출력: (1, 512, 56, 56)

 

 

중심부 (Bottleneck)

• 중심부는 512 채널을 1024 채널로 확장하여 더 깊은 특징을 학습한다. 이 단계는 인코더의 마지막 출력을 사용해 더 풍부한 고차원 표현을 학습하는 역할을 한다.

 

업샘플링 및 디코더 초기 단계

• 업샘플링(ConvTranspose2d)을 통해 특징 맵의 해상도를 두 배로 확장하여, 크기를 28x28에서 56x56으로 복원한다
• 스킵 연결을 통해 인코더에서 추출한 고해상도 정보를 결합하여 더 세밀한 정보를 유지한다.
• 결합된 출력은 두 개의 컨볼루션 층을 통해 처리되어 최종적으로 디코더 초기 단계의 출력이 된다.

 

---

3. 스킵 연결(Skip Connections)

U-Net의 핵심 요소 중 하나는 인코더와 디코더 사이의 스킵 연결이다. 스킵 연결은 인코더의 각 단계에서 추출한 특징을 디코더의 대응하는 단계에 전달하여 정보 손실을 줄이고 더 나은 분할 성능을 가능하게 한다. 이는 특히 작은 객체를 분할할 때 중요한 역할을 한다.

스킵 연결은 인코더에서 추출한 고해상도 정보를 디코더에서 업샘플링된 특징 맵과 결합하여 복원 과정에서 세부 정보를 유지한다. 예를 들어, 인코더의 특정 레이어에서 추출된 특징 맵은 디코더의 대응하는 레이어로 전달되어, 업샘플링된 특징과 함께 결합된다. 이를 통해 디코더는 업샘플링 과정에서 손실될 수 있는 중요한 공간적 세부 정보를 유지하며, 더 나은 결과를 얻을 수 있게 된다.

 

아래의 코드는 스킵 연결이 실제로 어떻게 사용되는지를 보여준다.

import torch
import torch.nn as nn
import torch.nn.functional as F

class SkipConnectionExample(nn.Module):
    def __init__(self):
        super(SkipConnectionExample, self).__init__()
        # 인코더와 디코더 정의
        self.encoder = nn.Conv2d(1, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.decoder = nn.ConvTranspose2d(64, 1, kernel_size=2, stride=2)
        self.conv = nn.Conv2d(64 + 1, 1, kernel_size=3, padding=1)
    
    def forward(self, x):
        # 인코더 부분
        encoded_features = self.encoder(x)  # (1, 64, H, W)
        pooled_features = self.pool(encoded_features)  # (1, 64, H/2, W/2)
        
        # 디코더 부분
        upsampled_features = self.decoder(pooled_features)  # 업샘플링: (1, 64, H, W)
        
        # 스킵 연결 - 인코더 출력(encoded_features)와 디코더 출력(upsampled_features)을 결합
        concatenated_features = torch.cat((upsampled_features, x), dim=1)  # 채널 방향으로 결합: (1, 64 + 1, H, W)
        output = self.conv(concatenated_features)  # 최종 출력: (1, 1, H, W)
        return output

# 모델 테스트
x = torch.randn(1, 1, 128, 128)  # 입력 이미지 (1, 채널, 높이, 너비)
skip_model = SkipConnectionExample()
output = skip_model(x)
print("Output Shape:", output.shape)  # (1, 1, 128, 128)

 

인코더에서 추출된 특징 맵(encoded_features)을 디코더의 출력(upsampled_features)과 결합하여 최종적으로 더 정교한 출력 결과를 만들어낸다. 이처럼 스킵 연결은 디코더가 업샘플링 과정에서 손실될 수 있는 세부 정보를 보완하고, 복원을 더 정확하게 만드는 데 중요한 역할을 한다.

 

---

 

오버랩 타일 전략 (Overlap-tile Strategy)

U-Net은 큰 이미지를 처리할 때 오버랩 타일 전략(Overlap-tile strategy)을 사용한다. 이 전략은 특히 GPU 메모리 제한으로 인해 한 번에 전체 이미지를 처리할 수 없을 때 유용하다. 이 전략을 통해 큰 이미지를 작은 타일로 나누어 처리하면서 경계 부분에서 발생할 수 있는 정보 손실을 줄이고, 전체 이미지를 매끄럽게 분할할 수 있다.

 

오버랩 타일 전략의 개념

• 큰 이미지 처리: 대형 이미지를 분할하는 경우, 메모리 제한으로 인해 한 번에 전체 이미지를 처리하는 것이 어려울 수 있다. 이를 해결하기 위해 이미지를 작은 타일(tile)로 나누어 처리한다.
• 오버랩(Overlap): 타일을 서로 겹치도록 설정하여 경계 영역에서 발생할 수 있는 정보 손실을 줄인다. 겹치는 영역을 통해 경계에서의 불연속성을 최소화하고 매끄러운 분할 결과를 얻을 수 있다.
• 미러링(Extrapolation by Mirroring): 경계 영역에서 입력 이미지가 부족한 경우, 경계를 넘어선 부분을 미러링(mirroring)하여 추가적인 정보를 제공한다. 이를 통해 모델은 경계 근처에서도 충분한 컨텍스트 정보를 사용할 수 있게 된다.

 

 

 

왼쪽 이미지는 큰 이미지를 작은 타일로 나눈 예시이다. 파란색 상자는 오버랩된 타일 영역을 나타내며, 노란색 상자는 해당 타일 내에서 모델이 예측을 수행할 부분을 의미한다.

오른쪽 이미지는 분할 결과를 보여준다. 타일의 예측 결과를 결합하여 전체 이미지를 재구성하게 되며, 이때 경계에서의 정보 손실이 최소화되도록 타일들이 겹쳐진 상태로 처리된다.

 

 

오버랩 타일 전략은 특히 경계 부분에서의 정보 손실을 줄이기 위한 중요한 방법이다. 이는 모델이 경계에서 일관된 예측을 수행하도록 돕고, 결과적으로 전체 이미지에 대한 매끄러운 분할을 가능하게 한다

 

왼쪽 그림에서 큰 이미지를 작은 타일로 나누는 모습을 볼 수 있다.
여기서 파란색 상자는 오버랩되는 부분을 나타내고,
노란색 상자는 실제로 모델이 예측하는 부분이다.
타일을 겹치게 잘라서 사용하는 이유는 경계 부분에서 정보를 잃지 않기 위해서이다. ==> 오버랩
이렇게 겹치는 부분이 있으면 타일의 경계에서 발생할 수 있는 부정확한 분할을 줄일 수 있다.

 

오른쪽 그림은 여러 타일의 예측 결과를 결합한 최종 분할 결과를 보여준다.
이때 타일의 예측 결과들을 다시 하나로 합치기 때문에, 경계 부분에서도 매끄러운 결과를 얻을 수 있다.

 

쉽게 말해서, 오버랩 타일 전략은 큰 이미지를 잘게 나누되, 조금씩 겹쳐서 나누는 방법이다. 이렇게 하면 각 타일이 서로의 경계 정보를 조금씩 공유하기 때문에, 경계에서도 일관된 결과를 낼 수 있다. 이 과정 덕분에 전체 이미지를 다시 합쳤을 때 연결이 자연스럽고 깔끔한 결과를 얻을 수 있다.

 

---

 

데이터 증강(Data Augmentation)과 U-Net의 활용

 

U-Net은 이미지 분할 작업에서 높은 성능을 발휘하기 위해 데이터 증강(Data Augmentation)을 사용한다. 데이터 증강은 모델의 일반화 성능을 높이기 위해 원본 데이터에 다양한 변형을 가하여 인공적으로 새로운 학습 데이터를 생성하는 방법이다. 특히 의료 영상과 같이 데이터가 적고 수집이 어려운 경우, 데이터 증강은 매우 중요한 역할을 한다.

탄성 변형(Elastic Deformation)

특히 U-Net에서 중요한 데이터 증강 기법 중 하나는 탄성 변형(Elastic Deformation)이다. 탄성 변형은 이미지의 모양을 탄력적으로 왜곡시키는 방법으로, 다음과 같은 이유로 매우 효과적이다:

1. 다양한 형태의 왜곡: 의료 영상에서 조직이나 세포는 다양한 형태로 존재한다. 탄성 변형은 이미지의 형태를 여러 가지로 왜곡하여 모델이 이러한 다양한 변형에 대해 학습할 수 있도록 한다.
2. 모델의 불변성 학습: 탄성 변형을 사용하면 모델은 특정 형태나 크기에 구애받지 않고 형태의 변형에 불변적인 특성을 학습하게 된다. 즉, 실제로 조직이나 세포의 형태가 다르게 나타나더라도 이를 올바르게 인식할 수 있는 능력을 갖추게 된다.

탄성 변형은 U-Net이 생물학적 조직의 자연스러운 형태 변화를 효과적으로 학습할 수 있게 하며, 이로 인해 적은 양의 데이터로도 강력한 성능을 발휘하게 된다.

 

import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
import numpy as np
import torch

# 데이터 증강 파이프라인 정의
train_transform = A.Compose([
    A.RandomRotate90(),  # 이미지 무작위 90도 회전
    A.Flip(),            # 무작위 뒤집기 (수평 또는 수직)
    A.ElasticTransform(alpha=1, sigma=50, alpha_affine=50, p=0.5),  # 탄성 변형
    A.RandomBrightnessContrast(p=0.2),  # 밝기 및 대비 조절
    A.Normalize(mean=(0.5,), std=(0.5,)),  # 이미지 정규화
    ToTensorV2()          # 텐서로 변환
])

# 샘플 이미지에 데이터 증강 적용
image = cv2.imread("sample_image.png", cv2.IMREAD_GRAYSCALE)
augmented = train_transform(image=image)
augmented_image = augmented["image"]

# 결과 확인
print(f"Original Shape: {image.shape}, Augmented Shape: {augmented_image.shape}")

 

---

 

U-Net 훈련 과정 (Training)

 

1. 이미지 타일 및 훈련 방법

훈련은 Caffe 프레임워크를 사용하여 확률적 경사 하강법(Stochastic Gradient Descent, SGD)으로 수행된다. U-Net에서 사용되는 unpaded 컨볼루션으로 인해 출력 이미지는 입력 이미지보다 일정한 경계 폭만큼 작아진다. 이 때문에 GPU 메모리를 최대한 효율적으로 사용하기 위해서는 큰 배치(batch) 대신 큰 입력 타일을 사용하는 것이 바람직하며, 결과적으로 배치 크기는 단일 이미지로 줄어든다.

• 큰 타일을 사용하는 이유는 GPU 메모리의 제약 때문에 큰 배치를 사용할 수 없기 때문이다.
• 모멘텀 값은 0.99로 높게 설정되는데, 이는 과거에 본 많은 훈련 샘플들이 현재 최적화 단계의 업데이트에 크게 기여하도록 하기 위함이다. 이렇게 하면 배치 크기가 작아도 안정적인 경사하강을 유지할 수 있다.

2. 손실 함수 정의

손실 함수는 최종 특징 맵에 대한 픽셀 단위
소프트맥스(soft-max)와 교차 엔트로피 손실(cross entropy loss)을 사용하여 계산된다.

 

3. 가중치 맵(weight map)

각 픽셀의 중요도를 반영하기 위해 가중치 맵(weight map)을 사전 계산하여 사용한다. 이는 데이터셋에서 특정 클래스의 픽셀 빈도가 다를 때 발생하는 불균형을 보정하고, 서로 붙어 있는 세포 간의 경계를 더 잘 학습하도록 유도하기 위해서이다.

예를 들어, 경계 픽셀은 잘 분리되지 않으면 객체를 정확하게 구분하는 데 어려움이 있으므로, 경계 부분의 픽셀 가중치를 높여 모델이 이를 명확히 학습하도록 한다.

 

4. 네트워크 가중치 초기화

깊은 네트워크에서는 많은 컨볼루션 층과 서로 다른 경로들이 있기 때문에 가중치 초기화가 매우 중요하다. 초기화가 적절하지 않으면 일부 네트워크 경로는 과도하게 활성화되고, 다른 부분은 학습에 기여하지 못하게 된다. 따라서 각 특징 맵이 대략적으로 단위 분산(unit variance)을 가지도록 가중치를 초기화하는 것이 중요하다.

 

아래 이미지를 보면 직관적으로 이해할 수 있다.

 

 

(a) 원본 이미지: DIC 현미경으로 촬영된 HeLa 세포 이미지이다. 이 이미지는 모델의 입력으로 사용된다.

(b) 정답 분할(Ground Truth Segmentation): 원본 이미지에서 각 세포를 다른 색상으로 구분한 정답 데이터이다. 모델이 학습할 때 이 데이터를 기준으로 손실을 계산하여 성능을 평가한다.

(c) 분할 마스크(Segmentation Mask): U-Net 모델이 생성한 결과이다. 세포를 분할하고 배경과 세포 영역을 구분한 결과로, 흰색은 세포, 검은색은 배경을 나타낸다.

(d) 가중치 맵(Weight Map): 모델이 훈련 중 경계 부분을 더 잘 학습하도록 경계 영역에 높은 가중치를 부여한 맵이다. 이를 통해 모델이 세포 간의 경계를 더 명확하게 학습할 수 있다.

 

---

 

U-Net의 응용 분야

U-Net은 주로 의료 영상 분할에 사용되며, 세포나 병변을 정확하게 구분하는 데 탁월한 성능을 보인다. 또한, 위성 이미지 분석, 자율 주행 차량의 도로 표지판 인식, 얼굴 마스크 분할 등 다양한 이미지 분할 작업에 응용되고 있다. U-Net의 간단하면서도 효율적인 구조 덕분에 다른 많은 분야에서도 성공적으로 적용되고 있다.

결론

U-Net은 이미지 분할 문제를 해결하기 위해 개발된 강력한 도구이다. 그 구조의 단순함과 강력함 덕분에 다양한 분야에서 성공적으로 활용되고 있으며, 특히 의료 영상 처리에서 중요한 역할을 하고 있다. 이를 통해 작은 객체나 경계 영역을 정확하게 분할할 수 있어 많은 연구자들이 U-Net을 기반으로 더 발전된 모델을 제안하고 있다. 앞으로도 이미지 분할 분야에서 중요한 역할을 할 것으로 기대된다.

 

 

여기까지 입니다. 끝.

긴 글 읽어주셔서 감사합니다.