본문 바로가기
  • 책상 밖 세상을 경험할 수 있는 Playground를 제공하고, 수동적 학습에서 창조의 삶으로의 전환을 위한 새로운 라이프 스타일을 제시합니다.
Computer Vision

[2025-1] 황영희 - U-Net: Convolutional Networks for Biomedical Image Segmentation

by 앵히 2025. 2. 13.

https://arxiv.org/abs/1505.04597

 

U-Net: Convolutional Networks for Biomedical Image Segmentation

There is large consent that successful training of deep networks requires many thousand annotated training samples. In this paper, we present a network and training strategy that relies on the strong use of data augmentation to use the available annotated

arxiv.org


1.  U-Net 이란?

이미지 세그멘테이션(Image Segmentation)은 이미지의 각 픽셀이 특정 카테고리(예: 자동차, 사람, 도로 등)에 속하는지를 분류하는 작업을 의미한다. 이는 이미지 전체를 단일 카테고리로 분류하는 이미지 분류(Image Classification)과는 달리, 픽셀 단위로 세밀한 분석을 수행해야 하므로 더 어려운 문제로 여겨진다.

대표적으로 알려진 LeNet, AlexNet, VGG, GoogLeNet 등의 모델은 주로 이미지 분류 문제를 다루는 데 사용되며, 이미지 세그멘테이션과는 차이가 있다.

이미지 세그멘테이션의 주요 유형

  • Semantic Segmentation: 이미지 내 객체를 카테고리별로 분할(예: 모든 자동차를 하나의 범주로 분류).
  • Instance Segmentation: 같은 카테고리에 속하더라도 서로 다른 객체를 구분하여 더 정밀하게 분할.

이미지 세그멘테이션은 의료 이미지 분석(종양 경계 추출), 자율주행 차량(도로와 보행자 감지), 증강현실(객체 감지 및 상호작용) 등 다양한 분야에서 활용되고 있다. 딥러닝 기술의 발전으로 이미지 세그멘테이션 분야에서 수백 가지 알고리즘이 제안되었으며, 그중 의료 이미지에 최적화된 구조로 잘 알려진 것이 U-Net이다.

이와 같은 의료 데이터를 다룰 때는 다음과 같은 제약이 있다.

  1. 데이터 희소성: 의료 데이터는 수집이 어렵고, 의료법에 의해 보호받아 공개된 데이터셋이 극히 적다.
  2. 라벨링의 어려움: 데이터 라벨링은 전문적인 지식이 요구되며, 의사가 데이터를 직접 라벨링하기에는 많은 시간과 노력이 소요된다.

논문에서는 이러한 제약을 극복하기 위해 다음과 같은 설정을 제시했다.

  • 작은 데이터셋: 단 30개의 이미지 데이터로 학습 진행.
  • 데이터 크기: 약 30MB 수준으로, 대규모 데이터셋이 아닌 환경에서 성능 검증.
  • 데이터 증강(Augmentation): 제한된 데이터를 다양하게 변환하여 학습에 활용.

이러한 접근법은 의료 데이터를 다룰 때 발생하는 문제를 효과적으로 해결하며, U-Net이 가진 강점과 특성을 보여준다.

2.  U-Net 구조

U-Net은 인코더-디코더(encoder-decoder) 기반의 딥러닝 모델로, 이미지의 특징을 추출하고 복원하는 과정을 통해 정확한 세그멘테이션을 수행한다. U-Net의 구조는 인코더와 디코더를 중심으로 설계되었으며, 이를 연결하는 스킵 연결(skip connection)이 특징적이다. 

2.1. 기본 구조

U-Net의 인코더는 입력 이미지의 특징을 추출하며, 디코더는 이를 복원하여 원본 이미지와 같은 해상도로 변환한다.

  • 인코더: 입력 이미지를 작은 크기로 압축하면서 중요한 특징을 추출한다. 컨볼루션을 통해 의미 있는 패턴을 학습하고, 맥스 풀링을 사용해 차원을 줄이며 불필요한 정보를 제거한다. 이 과정에서 채널 수를 증가시켜 다양한 특징을 포착할 수 있도록 한다. 결과적으로, 입력 이미지의 전체적인 구조를 단순화하면서도 중요한 정보를 유지하는 역할을 한다.  이는 오토인코더의 축소 경로(contracting path)와 유사한 역할을 한다.
  • 디코더: 인코더에서 압축된 정보를 이용해 원본 크기로 복원한다. 업샘플링과 전치 컨볼루션을 사용하여 점진적으로 해상도를 늘리고, 인코더에서 얻은 특징을 스킵 커넥션을 통해 보완하여 세밀한 정보를 유지한다. 이를 통해 단순히 크기를 키우는 것이 아니라, 중요한 특징을 반영한 고해상도 출력을 생성하는 역할을 한다. 이는 확장 경로(expanding path)로 불린다.

일반적인 인코더-디코더 모델에서는 디코딩 단계에서 업샘플링된 저해상도 특징만을 사용하기 때문에, 인코딩 과정에서 손실된 세부 정보나 위치 정보를 복원하는 데 한계가 있다.

 

U-Net의 차별점 - 스킵연결

U-Net은 인코더에서 얻어진 고해상도의 공간 정보를 디코더로 직접 전달하는 스킵 연결(skip connection) 을 도입했다.

  • 스킵 연결: 인코더에서 추출된 특징 맵을 디코더의 대응되는 레이어에 병합(concatenation)하여, 업샘플링된 저해상도 정보뿐만 아니라, 원래의 고해상도 공간 정보도 함께 활용한다. 
  • 이를 통해 세밀한 위치 정보전체 맥락 정보를 동시에 보존하여 더 정확한 세그멘테이션 결과를 도출할 수 있다.

U-Net의 신경망 구조는 스킵 연결을 평행하게 배치하고, 좌우 대칭이 되도록 설계되었다. 이 대칭적 구조는 네트워크의 형태가 U자와 같다고 하여 U-Net이라는 이름을 얻게 된 배경이기도 하다.

 

 

U-Net은 다음 세 가지 주요 부분으로 구성된다.

 

  • 인코더(Encoder): 이미지의 특징을 추출하며, 컨볼루션과 풀링을 통해 점진적으로 해상도를 줄이고 채널 수를 늘린다.
  • 브릿지(Bridge): 인코더와 디코더를 연결하는 중간 레이어로, 가장 추상적인 특징 정보를 포함한다.
  • 디코더(Decoder): 업샘플링을 통해 해상도를 점진적으로 복원하며, 스킵 연결을 활용해 인코더에서 전달된 고차원 정보와 병합한다.

이러한 구조적 특징 덕분에 U-Net은 적은 데이터셋에서도 높은 성능을 발휘하며, 특히 세밀한 경계와 위치 정보가 중요한 이미지 세그멘테이션 문제에서 강력한 성능을 보인다.

 

2.2. Contracting Path

U-Net의 Contracting Path(축소 경로)는 인코더의 역할을 수행하며, 입력 이미지의 특징을 점진적으로 추출하는 단계다. 이 과정에서 이미지의 차원은 축소되고, 채널 수는 증가한다.

구조와 역할

  • 맵(map)의 차원과 채널 수
    • 세로 방향 숫자는 맵의 공간적 차원(가로 x 세로)을, 가로 방향 숫자는 맵의 채널 수를 나타낸다.
    • 예를 들어, 세로 숫자 256x256과 가로 숫자 128은 해당 레이어가 256x256 크기의 맵128개의 채널을 가지는 것을 의미한다.
  • 입력 데이터:
    • 입력 이미지는 512x512x3 크기를 가지며, 이는 512x512 해상도를 가진 RGB 3채널 이미지이다.
  • 처리 과정:
    • 각 레이어에서는 2개의 3x3 컨볼루션 연산을 수행하며, 이를 통해 이미지의 채널 수를 늘린다.
    • 이어서, 2x2 맥스 풀링(max pooling) 연산으로 차원을 절반으로 축소한다.
    • 이 과정은 입력 이미지의 공간적 정보를 압축하고, 더 복잡한 특징을 추출하는 역할을 한다.

특징

  • Contracting Path의 단계가 진행될수록
    1. 차원 축소: 공간적 차원(이미지 크기)이 절반씩 줄어든다.
    2. 채널 증가: 특징 맵의 채널 수가 증가하여 더 많은 정보를 담을 수 있다.

이 축소 경로는 전역적인 맥락 정보(global context)를 학습하는 데 중점을 두며, 이후 디코딩 단계에서 세부 정보를 복원하기 위한 기반을 제공한다.

2.2.1. ConvBlock

ConvBlock은 U-Net의 인코더(Contracting Path)에서 각 단계마다 반복적으로 나타나는 핵심 연산 단위로, 특징 추출을 위한 연산들이 포함된 구성 요소다.

 

구조

ConvBlock은 다음의 연산들이 순차적으로 이루어진 두 개의 연속된 연산 묶음으로 구성된다.

  1. 3x3 Convolution: 작은 3x3 필터를 사용해 이미지를 스캔하며 특징을 추출한다.
    • 이 연산은 입력 이미지의 해상도를 좌우, 상하로 각각 1씩 줄인다.
    • 단, padding 옵션을 사용하면 입력과 출력의 해상도를 동일하게 유지할 수 있다.
  2. Batch Normalization: 학습 속도를 향상시키고, 과적합(overfitting)을 방지하기 위해 각 층의 출력을 정규화한다.
  3. ReLU 활성화 함수: 비선형성을 추가하여 신경망이 더 복잡한 패턴을 학습할 수 있도록 한다.

특징

  • 각 ConvBlock은 2개의 파란색 박스(3x3 Convolution → Batch Normalization → ReLU)를 하나로 묶어서 구현한다.
  • 이 ConvBlock을 사용하면 레이어를 효율적으로 관리할 수 있으며, 코드 재사용성이 높아진다.
  • 해상도 유지가 필요할 경우, ConvBlock의 컨볼루션 연산에 padding 옵션을 적용할 수 있다.

이 ConvBlock은 인코더에서 특징을 추출하고, 이후 디코더에서도 동일하게 사용되어 네트워크의 일관성을 유지한다.

ConvBlock을 코드로 구현하면 다음과 같다.

""" ConvBlock: U-Net의 기본 연산 블록 """
import tensorflow as tf
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation

class ConvBlock(tf.keras.layers.Layer):
    def __init__(self, n_filters):
        super(ConvBlock, self).__init__()
        # 두 개의 3x3 Conv2D 레이어 정의 (padding='same'으로 해상도 유지)
        self.conv1 = Conv2D(n_filters, kernel_size=3, padding='same')
        self.conv2 = Conv2D(n_filters, kernel_size=3, padding='same')
        
        # 두 개의 BatchNormalization 레이어 정의
        self.bn1 = BatchNormalization()
        self.bn2 = BatchNormalization()
        
        # 활성화 함수로 ReLU 사용
        self.activation = Activation('relu')

    def call(self, inputs):
        # 첫 번째 Conv → BatchNorm → ReLU
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.activation(x)
        
        # 두 번째 Conv → BatchNorm → ReLU
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.activation(x)
        
        return x

2.2.2. EncoderBlock

EncoderBlock은 U-Net의 인코더에서 사용되는 구성 요소로, ConvBlock다운샘플링(down sampling)을 결합하여 각 인코더 단계에서 반복적으로 나타나는 블록이다.

 

구조

  • ConvBlock: 특징을 추출하는 역할을 담당하며, 하나의 ConvBlock이 포함된다.
  • 출력
    1. 첫 번째 출력은 디코더로 복사하기 위한 연결선(녹색 화살표)이다.
    2. 두 번째 출력은 2x2 Max Pooling을 통해 다운샘플링된 결과로, 인코더의 다음 단계로 전달(빨간색 화살표)된다.

특징

  • EncoderBlock은 ConvBlock과 Max Pooling을 하나의 블록으로 묶어 구현하면 코드의 재사용성과 가독성을 높일 수 있다.
  • 이 블록은 디코더 단계에서 사용할 스킵 연결(skip connection)의 출력을 생성하면서, 다음 단계로 전달할 다운샘플링된 특징도 생성한다.

EncoderBlock을 코드로 구현하면 다음과 같다.

""" EncoderBlock: U-Net 인코더의 기본 구성 요소 """
class EncoderBlock(tf.keras.layers.Layer):
    def __init__(self, n_filters):
        super(EncoderBlock, self).__init__()
        # ConvBlock 정의
        self.conv_blk = ConvBlock(n_filters)
        # 2x2 Max Pooling 정의
        self.pool = MaxPooling2D((2,2))

    def call(self, inputs):
        # ConvBlock으로 특징 추출
        x = self.conv_blk(inputs)
        # Max Pooling으로 다운샘플링
        p = self.pool(x)
        return x, p

2.3. Bridge

인코더와 디코더를 이어주는 정보를 담고 있는 브릿지는 전체적으로 1개의 ConvBlock으로 표현할 수 있는 구조를 가지고 있다.

2.4. Expanding Path

2.4.1. DecoderBlock

DecoderBlock은 U-Net의 디코더(Expanding Path)에서 사용되는 핵심 구성 요소로, 인코더에서 추출된 저해상도 특징을 업샘플링하여 해상도를 점진적으로 복원하는 역할을 한다.
또한, 스킵 연결을 통해 인코더의 대응되는 고해상도 특징 맵을 결합하여 세부적인 공간 정보를 보존하고, 더 정확한 복원이 가능하도록 한다.

구조

  1. 연한 노란색 박스 (Transposed Convolution)
    • 입력: Bridge 또는 이전 디코더 레이어의 출력(feature map).
    • 작동: Transposed Convolution Layer(전치 컨볼루션)를 사용하여 해상도를 2배로 늘리고, Conv 연산을 통해 채널 수는 절반으로 줄인다.
      • 전치 컨볼루션은 해상도를 증가(업샘플링)하는 컨볼루션 연산
      • 일반적인 Conv2D는 해상도를 줄이는 데 사용, 전치 컨볼루션은 해상도를 키움
    • 목적: 업샘플링을 통해 공간 정보를 복원하고, 다음 단계의 입력으로 활용 가능한 해상도를 생성한다.
  2. 녹색 박스 (Skip Connection)
    • 입력: 인코더의 대칭되는 레이어에서 전달된 feature map.
    • 스킵 연결을 통해 고차원 특징 정보를 디코더에 전달하여, 저차원 정보만 사용하는 한계를 극복한다.
  3. 결합 (Concatenation)
    • Transposed Convolution 결과(연한 노란색 박스)와 스킵 연결 결과(녹색 박스)를 concatenation으로 결합한다.
    • 이를 통해 저차원과 고차원 정보를 모두 포함한 feature map을 생성한다.
  4. ConvBlock
    • 결합된 feature map을 ConvBlock으로 처리하여 채널 수를 다시 절반으로 줄인다.
    • 목적: 디코더의 출력 정보를 가공하여 다음 단계로 전달.

위 그림에서 반복적으로 나타나는 회색 박스를 한 개의 레이어 블록으로 구현하여 사용하면 편리하다. 이 블록 이름을 DecoderBlock이라고 가정하자.

DecorderBlock

다음은 DecorderBlock을 구현한 코드이다.

""" DecoderBlock: U-Net 디코더의 레이어 블록 """
import tensorflow as tf
from tensorflow.keras.layers import Conv2DTranspose, Concatenate

class DecoderBlock(tf.keras.layers.Layer):
    def __init__(self, n_filters):
        super(DecoderBlock, self).__init__()
        # Transposed Convolution: 해상도 2배 증가, 채널 수 감소
        self.up = Conv2DTranspose(n_filters, kernel_size=(2, 2), strides=2, padding='same')
        # ConvBlock 정의
        self.conv_blk = ConvBlock(n_filters)

    def call(self, inputs, skip):
        # Transposed Convolution으로 업샘플링
        x = self.up(inputs)
        # 업샘플링된 출력과 스킵 연결된 feature map 병합
        x = Concatenate()([x, skip])
        # ConvBlock으로 병합된 맵 처리
        x = self.conv_blk(x)
        return x

 

  • (2,2) : 2×2 크기의 필터를 사용해 업샘플링

디코더 그림의 맨 상단 오른쪽 부분은 U-Net의 최종 출력으로, 1×1 컨볼루션을 통해 특징 맵을 처리하여 입력 이미지의 각 픽셀을 특정 클래스로 분류하는 세그멘테이션 맵을 생성하는 단계이다. 이때 1×1 컨볼루션의 역할은 디코더에서 얻은 다채널 특징 맵을 클래스 개수에 맞춰 변환하는 것으로, 컨볼루션 필터의 개수는 예측할 카테고리 개수와 동일하게 설정된다. 출력층의 활성 함수로는 이진 분류(Binary Segmentation)일 경우 sigmoid 함수를 사용하여 픽셀별로 0~1 사이의 값을 출력하며, 다중 클래스 분류(Multi-class Segmentation)일 경우 softmax 함수를 사용하여 각 픽셀에 대해 클래스별 확률 분포를 예측한다.

 

EncoderBlock과 DecoderBlock을 사용하면 다음과 같이 U-Net을 코드로 간단히 구현할 수 있다.

""" U-Net Model """
class UNET(tf.keras.Model):
    def __init__(self, n_classes):
        super(UNET, self).__init__()

        # Encoder
        self.e1 = EncoderBlock(64)
        self.e2 = EncoderBlock(128)
        self.e3 = EncoderBlock(256)
        self.e4 = EncoderBlock(512)

        # Bridge
        self.b = ConvBlock(1024)

        # Decoder
        self.d1 = DecoderBlock(512)
        self.d2 = DecoderBlock(256)
        self.d3 = DecoderBlock(128)
        self.d4 = DecoderBlock(64)

        # Outputs
        if n_classes == 1:
            activation = 'sigmoid'
        else:
            activation = 'softmax'

        self.outputs = Conv2D(n_classes, 1, padding='same', activation=activation)

    def call(self, inputs):
        s1, p1 = self.e1(inputs)
        s2, p2 = self.e2(p1)
        s3, p3 = self.e3(p2)
        s4, p4 = self.e4(p3)

        b = self.b(p4)

        d1 = self.d1(b, s4)
        d2 = self.d2(d1, s3)
        d3 = self.d3(d2, s2)
        d4 = self.d4(d3, s1)

        outputs = self.outputs(d4)

        return outputs

 

2.5. 전체 네트워크 구조

2.6. 전체 코드

다음은 U-Net을 함수 기반으로 구현한 전체 코드로, EncoderBlockDecoderBlock 클래스를 사용하지 않고 함수형 구조로 작성하였다. 이 모델은 100×75×3 크기의 이미지를 입력받아 100×75×1 크기의 세그멘테이션 이미지를 출력한다. 해상도가 홀수이므로 디코딩 과정에서 해상도를 맞추기 위해 약간의 수정이 이루어졌다. 또한, 3×3 컨볼루션을 사용할 때 padding 옵션을 적용하여, 컨볼루션 후에도 해상도가 줄어들지 않고 입력 이미지의 해상도를 그대로 유지하도록 구현하였다.

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Conv2D, Conv2DTranspose, MaxPool2D, BatchNormalization, Activation, Concatenate, Input

# Convolution Block: 3x3 Conv2D 두 개를 사용하여 특징 추출 (Padding='same'으로 해상도 유지)
def ConvBlock(n_filter, inputs):
    x = Conv2D(n_filter, 3, padding='same')(inputs)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    x = Conv2D(n_filter, 3, padding='same')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    return x

# U-Net 모델 정의 (함수형 구조)
def unet_like():
    # 입력 레이어 (이미지 크기: 100x75, 채널 수: 3)
    inputs = Input(shape=(100, 75, 3))

    # ----- Encoder (Downsampling) -----
    c1 = ConvBlock(64, inputs)  # ConvBlock 적용 (100x75x64)
    p1 = MaxPool2D(2)(c1)       # MaxPooling 적용 (50x37x64)

    c2 = ConvBlock(128, p1)     # ConvBlock 적용 (50x37x128)
    p2 = MaxPool2D(2)(c2)       # MaxPooling 적용 (25x18x128)

    c3 = ConvBlock(256, p2)     # ConvBlock 적용 (25x18x256)
    p3 = MaxPool2D(2)(c3)       # MaxPooling 적용 (12x9x256)

    c4 = ConvBlock(512, p3)     # ConvBlock 적용 (12x9x512)
    p4 = MaxPool2D(2)(c4)       # MaxPooling 적용 (6x4x512)

    # ----- Bridge (Bottleneck) -----
    b = ConvBlock(1024, p4)     # Bridge Layer (6x4x1024)

    # ----- Decoder (Upsampling) -----
    d1 = Conv2DTranspose(512, (2,2), strides=2, output_padding=(0,1))(b)  # 업샘플링 (12x9x512)
    d1 = Concatenate()([c4, d1])  # 스킵 연결 (12x9x512 + 12x9x512)
    d1 = ConvBlock(512, d1)       # ConvBlock 적용 (12x9x512)

    d2 = Conv2DTranspose(256, (2,2), strides=2, output_padding=(1,0))(d1)  # 업샘플링 (25x18x256)
    d2 = Concatenate()([c3, d2])  # 스킵 연결 (25x18x256 + 25x18x256)
    d2 = ConvBlock(256, d2)       # ConvBlock 적용 (25x18x256)

    d3 = Conv2DTranspose(128, (2,2), strides=2, output_padding=(0,1))(d2)  # 업샘플링 (50x37x128)
    d3 = Concatenate()([c2, d3])  # 스킵 연결 (50x37x128 + 50x37x128)
    d3 = ConvBlock(128, d3)       # ConvBlock 적용 (50x37x128)

    d4 = Conv2DTranspose(64, (2,2), strides=2, output_padding=(0,1))(d3)   # 업샘플링 (100x75x64)
    d4 = Concatenate()([c1, d4])  # 스킵 연결 (100x75x64 + 100x75x64)
    d4 = ConvBlock(64, d4)        # ConvBlock 적용 (100x75x64)

    # ----- Output Layer -----
    outputs = Conv2D(1, 1, padding='same', activation='sigmoid')(d4)  # 최종 출력 (100x75x1, 픽셀별 확률 값)

    # 모델 정의
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model

output_padding=(0,1)은 Conv2DTranspose에서 업샘플링 후 너비 방향으로 1픽셀을 추가하여 해상도를 맞추는 역할을 한다. 이는 홀수 해상도에서 업샘플링할 때 정확한 크기를 유지하기 위해 사용된다.