본문 바로가기

AI/Paper - Theory

[UNETR 논문 리뷰] - UNETR: Transformers for 3D Medical Image Segmentation

반응형

*UNETR 논문 리뷰를 위한 글이고, 질문이 있으시다면 언제든지 댓글로 남겨주세요!

 

UNETR paper: [2103.10504] UNETR: Transformers for 3D Medical Image Segmentation (arxiv.org)

 

UNETR: Transformers for 3D Medical Image Segmentation

Fully Convolutional Neural Networks (FCNNs) with contracting and expanding paths have shown prominence for the majority of medical image segmentation applications since the past decade. In FCNNs, the encoder plays an integral role by learning both global a

arxiv.org

 

UNETR github: https://github.com/Project-MONAI/MONAI

 

GitHub - Project-MONAI/MONAI: AI Toolkit for Healthcare Imaging

AI Toolkit for Healthcare Imaging. Contribute to Project-MONAI/MONAI development by creating an account on GitHub.

github.com


Contents

1. Simple Introduction

2. Background Knowledge: 3D U-Net, Vision Transformer

3. Method

4. Furthermore


Simple Introduction

Segmentation Result

UNETR은 Vision Transformer와 3D U-Net을 구조 결합해서 segmentation을 진행하였다!

2D 형태의 input이 아닌 3D input으로 Vision transformer에 적용하는 개념도 매우 흥미롭다.

 

UNETR로 3D Medical image segmentation은 본격적으로 Transformer를 도입하기 시작했고, 성능도 향상되었다.

 

한번 모델의 구조를 같이 알아봅시다!


Background Knowledge: 3D U-Net, Vision Transformer

(3D)U-Net 논문 리뷰: https://kyujinpy.tistory.com/9

 

[(3D) U-Net 논문 리뷰] - 3D U-Net: Learning Dense VolumetricSegmentation from Sparse Annotation

3D U-Net paper: [1606.06650] 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation (arxiv.org) 3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation This paper introduces a network for volumetric segmentation that learns

kyujinpy.tistory.com

 

Vision Transformer 논문 리뷰: https://kyujinpy.tistory.com/3

 

[Vision Transformer 논문 리뷰] - AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

Vision Transformer paper: https://arxiv.org/abs/2010.11929 An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale While the Transformer architecture has become the de-facto standard for natural language processing tasks, its application

kyujinpy.tistory.com

 

위의 개념을 반드시 알고 오셔야 모델의 구조를 이해할 수 있습니다!

반드시 공부하는 것을 추천드립니다!

*해당 논문 리뷰는 Vision Transformer와 3D U-Net에 대한 개념이 있다는 가정하에 진행됩니다!


Method

UNETR Structure

모델의 간단한 Overview structure이다.

일단 3D 형태의 dataset을 3D patches로 쪼갠 후1D로 flatten 시킨다.

그리고 Transformer Encoder에 넣은 후, Decoder에 넣어서 segmentation을 진행한다.

 

좀 더 자세히 들어가보자!

UNETR

모델 구조는 다음과 같다.

1. Input dataset을 3D patch로 만든 후 linear projection 진행한다.

2. Normalization과 Multi-Head-Attention을 적용(Transformer block이라고 한다.)한 후, residual network의 구조를 통하여서 합쳐진 데이터를 MLP에 다시 넣는다.

3. 위의 과정, 즉 Transformer block을 총 N=12번 반복한다.

4. Transformer 연산이 끝난 후, 나온 embedding을 reshape 과정을 통하여 3D voxel(데이터)형식으로 reconstruction한다.

5. Upsampling을 수행 하면서, N=9, N=6, N=3, N=0(original)에서 나온 feature들과 skip-connection을 진행한다.

6. 마지막 layer에서 최종 dimension과 channel을 조절한 후 segmentaiton result를 만든다.

 


Transformer block

+) MLP 구조는 두개의 linear와 GELU activation function을 이용한다.


Furthermore

UNETR은 Medical segmentation 분야에서 널리 이용될 가능성이 크다고 생각한다.

이 외에도 Transformer와 3D U-Net을 결합한 모델이 계속해서 나오고, medical field에서 segmentaiton을 하는 논문들이 계속해서 나오니, 이 분야로 관심이 있다면 논문 구현도 해보길 추천한다!


 

- 2022.12.30 kyujinpy 작성.

 

반응형