U-Net로 하는 의미론적 분할

PyTorch로 배우는 이미지 딥러닝

Michal Oleszak

Machine Learning Engineer

의미론적 분할

동일 클래스 내 개별 객체 구분 없음
의료 영상, 위성 이미지 분석에 유용
인기 아키텍처: U-Net

U-Net 아키텍처

U-Net 아키텍처 다이어그램

인코더:

합성곱·풀링 레이어
다운샘플링: 공간 크기는 줄이고 채널 수는 늘림

U-Net 아키텍처

U-Net 아키텍처 다이어그램

디코더:

인코더와 대칭
전치 합성곱으로 특성 맵 업샘플링

U-Net 아키텍처

U-Net 아키텍처 다이어그램

스킵 연결:

인코더에서 디코더로 연결
다운샘플링으로 잃은 세부 정보를 보존

전치 합성곱

전치 합성곱 다이어그램

디코더에서 특성 맵 업샘플링: 높이·너비는 늘리고 채널 수는 줄임
전치 합성곱 과정:
1. 입력 특성 맵 사이/주변에 0을 삽입
2. 0 패딩된 입력에 일반 합성곱 수행

PyTorch의 전치 합성곱

import torch.nn as nn

upsample = nn.ConvTranspose2d(
    in_channels=in_channels,
    out_channels=out_channels,
    kernel_size=2,
    stride=2,
)

U-Net: 레이어 정의

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


        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)


        self.upconv3 = nn.ConvTranspose2d(512, 256,
                                          kernel_size=2, stride=2)
        self.upconv2 = nn.ConvTranspose2d(256, 128,
                                          kernel_size=2, stride=2)
        self.upconv1 = nn.ConvTranspose2d(128, 64,
                                          kernel_size=2, stride=2)


        self.dec1 = self.conv_block(512, 256)
        self.dec2 = self.conv_block(256, 128)
        self.dec3 = self.conv_block(128, 64)
        self.out = nn.Conv2d(64, out_channels, kernel_size=1)

인코더:

합성곱 블록

def conv_block(self, in_channels, out_channels):
return nn.Sequential(
  nn.Conv2d(in_channels, out_channels),
  nn.ReLU(inplace=True),
  nn.Conv2d(out_channels, out_channels),
  nn.ReLU(inplace=True)
)

풀링 레이어

디코더:
- 전치 합성곱
- 합성곱 블록

U-Net: forward 메서드

def forward(self, x):

    x1 = self.enc1(x)
    x2 = self.enc2(self.pool(x1))
    x3 = self.enc3(self.pool(x2))
    x4 = self.enc4(self.pool(x3))


    x = self.upconv3(x4)

    x = torch.cat([x, x3], dim=1)

    x = self.dec1(x)


    x = self.upconv2(x)
    x = torch.cat([x, x2], dim=1)
    x = self.dec2(x)

    x = self.upconv1(x)
    x = torch.cat([x, x1], dim=1)
    x = self.dec3(x)


    return self.out(x)

입력을 인코더의 합성곱 블록과 풀링 레이어로 통과시킵니다
디코더와 스킵 연결:
- 인코딩된 입력에 전치 합성곱 적용
- 해당 인코더 출력과 연결(concatenate)
- 합성곱 블록 통과
- 모든 디코더 단계에서 반복
마지막 디코더 단계의 출력을 반환

추론 실행

model = UNet()
model.eval()


image = Image.open("car.jpg")
transform = transforms.Compose([transforms.ToTensor()])
image_tensor = transform(image).unsqueeze(0)


with torch.no_grad():
    prediction = model(image_tensor).squeeze(0)


plt.imshow(prediction[1, :, :])
plt.show()

원본 자동차 이미지

자동차 이미지에 겹쳐진 의미 마스크

연습해 봅시다!

PyTorch로 배우는 이미지 딥러닝