[AI] 파이토치로 처음부터 U-Net 구현하기 - 1

파이토치로 처음부터 U-Net 구현하기 - 1

 

U-Net의 기본 개념

U-Net은 'U-Net: Convolutional Networks for Biomedical Image Segmentation'이라는 논문에서 제안한 구조이다.

* 논문 주소: https://arxiv.org/pdf/1505.04597.pdf

매우 적은 수의 학습 데이터(위 이미지 형태)만을 사용하여 정확한 시맨틱 세그멘테이션 성능을 보여주어 2015년 세포 추적 대회에서 큰 점수 차로 우승했다.

여기서 시맨틱 세그멘테이션(Semantic Segmentation)이 무엇일까?

시맨틱 세그멘테이션은 한국말로 '의미론적 분할'이라고 하는데, 아래의 예시 이미지와 같이 이미지 내에서 픽셀 단위로 분류하여, 모든 픽셀이 어떤 카테고리에 속하는지 분류하는 것이다. 일반적으로 Computer Vision에서는 Object Detection, Image Classification 등 다양한 분야가 있지만 Semantic Segmentation은 비교적 이러한 문제들보다 난이도가 높은 문제에 속한다.

 

지금까지 설명한 시맨틱 세그멘테이션은 아래의 한 예시처럼 오늘날 의료용 이미지 분석, 자율주행 등 아주 광범위한 분야에서 사용되고 있다.

 

이번 포스팅부터 시작하여 몇 차례로 나눠서 의료용 이미지에서 의미론적 분할(Semantic Segmentation)을 수행하기 위해 많이 사용되는 모델인 U-Net의 구조에 대해 알아보고 파이토치로 직접 구현까지 해보고자 한다.

포스팅 초반부에서 본 논문을 간략하게 소개하면서 매우 적은 학습 데이터를 사용했음에도 높은 정확도를 기록했다고 했는데, 의료용 이미지와 같은 경우 상식적으로 생각하더라도 데이터가 부족할 수 밖에 없다. 실제로 이 논문에 나오는 훈련용 데이터셋 역시 총 30개 밖에 되지 않으며 용량 또한 훈련 및 테스트 데이터셋 전체가 대략 22mb 정도로 매우 작다.(본 데이터는 https://github.com/alexklibisz/isbi-2012 여기서 다운로드 할 수 있다.)

데이터 부족을 커버하기 위해 본 논문에서는 여러 데이터 증강(Data Augmentation) 기법들 또한 소개하고 있다.

 

U-Net의 구조

U-Net은 이름에서 알 수 있듯이 구조가 알파벳 'U' 모양을 띠고 있는 인코더-디코더 기반 모델이다.

 

인코딩 단계에서는 입력 이미지의 특징을 포착하기 위해 차원을 축소해 나가며, 디코딩 단계에서는 저차원으로 인코딩된 정보만 이용하여 채널의 수를 줄이고 차원을 늘려서 고차원의 이미지를 복원한다. 하지만 인코딩 단계에서 차원 축소를 거치면서 이미지 객체에 대한 자세한 위치 정보를 잃게 되고, 디코딩 단계에서도 저차원의 정보만을 이용하기 때문에 위치 정보 손실을 회복하지 못하게 된다.

하지만 U-Net은 인코더(contract path)의 피처맵을 디코더(expansive path)의 피처맵에 연결하여 위치 정보 전달하는 식으로 위치 정보 손실의 문제를 극복했다.

 

U-Net은 아래 이미지와 같이 이미지를 압축하는 수축 경로(contracting path)와 수축된 이미지를 다시 원본 이미지의 크기로 복원하는 확장 경로(expansive path)로 구성되어 있다.

https://glassboxmedicine.com/2020/01/21/segmentation-u-net-mask-r-cnn-and-medical-applications/