[01 X-ray Hand] - SegNet의 아키텍처

SegNet은 이미지 분할 문제를 해결하기 위해 개발된 합성곱 신경망(CNN) 기반의 모델이다. SegNet은 주로 실시간 세그멘테이션에 초점을 맞추어 설계되었으며, U-Net과 같은 목적을 지니고 있지만 더 가볍고 효율적인 구조를 가지고 있다. 이 모델은 Enet 및 FCN과 같은 다른 세그멘테이션 모델들과 비교하여, 인코더-디코더 구조를 채택하고 있으며, 이를 통해 공간 정보를 유지하면서 효율적으로 이미지 분할을 수행할 수 있다.

---

SegNet의 아키텍처

SegNet은 크게 인코더(Encoder)와 디코더(Decoder)로 구성된 아키텍처이다. 이 아키텍처는 VGG16에서 영감을 받은 인코더 구조를 사용하지만, 주로 특징을 추출한 후 그 특징을 복원하는 방식으로 설계되었다. 인코더는 이미지를 점점 작게 만들어 특징을 추출하고, 디코더는 이 특징을 사용해 다시 원래의 해상도로 복원한다. 중요한 차이점은 디코더 단계에서 최대 풀링(Max Pooling)의 인덱스 정보를 활용해 원래 이미지의 공간 정보를 되살린다는 것이다.

 

 

SegNet 코드

아래는 SegNet의 파이썬 코드이다. PyTorch를 사용하여 구현되었으며, 각 인코더 및 디코더 단계에 대한 자세한 설명이 포함되어 있다.

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

class SegNet(nn.Module):
    def __init__(self, num_classes=21):
        super(SegNet, self).__init__()
        
        # 인코더 정의 (VGG16 구조에서 가져온 부분)
        self.encoder_block1 = self.make_encoder_block(3, 64)
        self.encoder_block2 = self.make_encoder_block(64, 128)
        self.encoder_block3 = self.make_encoder_block(128, 256, num_layers=3)
        self.encoder_block4 = self.make_encoder_block(256, 512, num_layers=3)
        self.encoder_block5 = self.make_encoder_block(512, 512, num_layers=3)
        
        # 디코더 정의
        self.decoder_block5 = self.make_decoder_block(512, 512, num_layers=3)
        self.decoder_block4 = self.make_decoder_block(512, 256, num_layers=3)
        self.decoder_block3 = self.make_decoder_block(256, 128, num_layers=3)
        self.decoder_block2 = self.make_decoder_block(128, 64)
        self.decoder_block1 = self.make_decoder_block(64, num_classes)

    def make_encoder_block(self, in_channels, out_channels, num_layers=2):
        layers = []
        for _ in range(num_layers):
            layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))
            layers.append(nn.BatchNorm2d(out_channels))
            layers.append(nn.ReLU(inplace=True))
            in_channels = out_channels
        layers.append(nn.MaxPool2d(kernel_size=2, stride=2, return_indices=True))
        return nn.Sequential(*layers)

    def make_decoder_block(self, in_channels, out_channels, num_layers=2):
        layers = []
        for _ in range(num_layers):
            layers.append(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))
            layers.append(nn.BatchNorm2d(out_channels))
            layers.append(nn.ReLU(inplace=True))
            in_channels = out_channels
        layers.append(nn.MaxUnpool2d(kernel_size=2, stride=2))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        # 인코더 부분
        x, indices1 = self.encoder_block1(x)
        x, indices2 = self.encoder_block2(x)
        x, indices3 = self.encoder_block3(x)
        x, indices4 = self.encoder_block4(x)
        x, indices5 = self.encoder_block5(x)

        # 디코더 부분
        x = self.decoder_block5(x, indices5)
        x = self.decoder_block4(x, indices4)
        x = self.decoder_block3(x, indices3)
        x = self.decoder_block2(x, indices2)
        x = self.decoder_block1(x, indices1)

        return x

# 모델 테스트
model = SegNet(num_classes=21)
x = torch.randn(1, 3, 256, 256)  # 입력 이미지 (1, 채널, 높이, 너비)
output = model(x)
print("Output Shape:", output.shape)  # 예상 출력: (1, 21, 256, 256)

 

코드 설명

1. Encoder Block
   • VGG16에서 가져온 구조를 활용하여 인코더를 구성한다. make_encoder_block() 함수는 인코더 블록을 생성하며, 각 블록은 Convolution, Batch Normalization, 그리고 ReLU 활성화 함수로 이루어져 있다.
   • 각 인코더 블록의 마지막에는 Max Pooling이 사용되며, 이때 풀링 연산의 인덱스를 반환한다. 이 인덱스 정보는 디코더에서 Unpooling할 때 사용되어 원래의 공간 정보를 되살리는 데 도움을 준다.
2. Decoder Block
   • 디코더 블록은 인코더 블록의 역순으로, make_decoder_block() 함수로 생성된다. 각 블록은 Convolution , Batch Normalization , 그리고 ReLU 활성화 함수로 구성된다.
   • 디코더에서 중요한 점은 Max Unpooling을 사용하여, 인코더에서 저장된 Index 정보를 기반으로 원래의 해상도로 복원한다는 점이다. 이를 통해 인코더에서 손실된 공간 정보를 복원하고, 분할 결과의 해상도를 유지할 수 있다.
3. Forward Pass
   • 인코더 부분에서는 입력 이미지를 점점 축소하며 특징을 추출하고, 각 단계에서 Polling 인덱스를 저장한다.
   • 디코더 부분에서는 저장된 풀링 인덱스를 사용하여 이미지의 해상도를 복원한다. 이는 원래 이미지의 공간 정보를 최대한 유지하면서 복원하는 데 중요한 역할을 한다.
4. 모델 테스트
   • 마지막으로 모델을 생성하고, 임의의 입력 이미지를 통해 출력의 형태를 확인한다. 입력 크기가 (1, 3, 256, 256)일 때 출력 크기도 동일하게 (1, 21, 256, 256)이 된다. 이는 SegNet이 입력과 동일한 해상도의 출력 분할 맵을 생성한다는 것을 의미한다.

---

SegNet의 장점

1. 효율적인 Incoder-Decoder 구조: SegNet은 VGG16을 기반으로 한 인코더-디코더 구조를 사용하여 특징 추출과 복원을 효과적으로 수행한다. 특히 Max-pooling 인덱스를 사용하여 해상도를 복원하는 방식은 계산 효율성을 높이면서도 공간 정보를 유지할 수 있는 장점이 있다.
2. 가벼운 모델: U-Net과 달리 SegNet은 더 가벼운 구조를 가지고 있어 실시간 처리가 가능하다. 이는 자율 주행 차량이나 드론과 같은 실시간 세그멘테이션이 필요한 애플리케이션에서 유리하다.
3. 공간 정보 유지: SegNet은 인코더에서 풀링 시 저장한 인덱스 정보를 디코더에서 사용하여 언풀링함으로써, 공간 정보를 최대한 유지하면서 이미지를 복원할 수 있다. 이는 분할 정확도를 높이는 데 큰 기여를 한다.

 

SegNet의 응용 분야

SegNet은 도로 표지판 인식, 자율 주행, 위성 이미지 분석, 그리고 도시 환경에서의 객체 분할 등 다양한 분야에서 사용된다. 특히 실시간으로 높은 정확도의 세그멘테이션이 필요한 응용 분야에서 많이 사용되며, 상대적으로 경량화된 구조 덕분에 하드웨어 제약이 있는 환경에서도 활용이 가능하다.