Segmentation of the lungs on a Chest X-ray image

안녕하세요! M_AI 입니다!

이번에는 Segmentation인데, U-Net을 이용한 기본적인 segmentation을 하는 방법을 소개하고자 합니다!

---

이 글을 읽으면 좋을 것 같은 사람들!

- 기본적인 분류(classfication) 모델을 Tensorflow 2.x 버전으로 구현할 줄 아는 사람
- Clssification을 넘어서 Semantic segmentation을 하길 원하는데 오픈 소스를 봐도 이해가 가질 않는 사람.
- 오픈 소스를 이해하더라도 입맛대로 수정하는 방법을 모르는 사람.

본 게시글의 목적

※ 본 게시글은 딥러닝 모델의 상세한 이론적 설명보다는 모델 구조를 바탕으로 코드 분석하고 이해하는 데에 목적에 있습니다!
※ 상세한 이론을 원하면 다른 곳에서 찾아보시길 바랍니다!

본 게시글의 특징

① Framework : Tensorflow 2.x
② 학습 환경
- Google Colab Pro
③ Batch size = 1
④ Pulmonary Chest X-Ray Abnormalities 데이터셋을 이용한 폐(lungs)를 분할
( 데이터셋 출처 : https://www.kaggle.com/kmader/pulmonary-chest-xray-abnormalities )

Pulmonary Chest X-Ray Abnormalities

⑤ 게시글에서는 반말

---

1. Data and Preprocessing

본 게시글에서 사용되는 데이터는 Kaggle에 오픈된 Pulmonary Chest X-Ray Abnormalities 데이터셋으로, Input Image로는 Chest X-ray와 Output Image는 left lung mask, right lung mask로 구성되어 있다.

Figure 1. A Sample of Dataset

전처리로는 크게 네 단계로 나뉘는데 이는 다음과 같다.

1.1. Integrate left lung mask and right lung mask

Figure 1과 같이 좌우로 나누어진 lung mask를 Figure 2처럼 하나로 통합한다.

Figure 2. Integration into one mask

1.2. Resize to (256, 256) - Grayscale

4000이 넘는 큰 이미지를 (256, 256)으로 리사이즈 해준다. 또한, grayscale로 진행한다.

1.3. Dataset Split

해당 데이터셋의 크기는 138개로, 모델 검증을 위해서 훈련, 검증, 테스트 데이터셋을 각각 90, 20, 28로 랜덤하게 나눈다.

FIgure 3. Dataset Split

1.4. Data Augmentation

적은 데이터 수로 인해 데이터 증강은 필수로 진행한다. 여기서 사용된 방식은 좌우 반전, 상하 반전, 90도 회전, 270도 회전을 한다.

Figure 4. Data Augmentation

2. Training

A. Model : U-Net
Input Shape : (256, 256, 1)
Output Shape : (256, 256, 1)

Figure 5. U-Net Architecture

B. Optimizer : Adam (Learning Rate = 0.001)

C. Performance Evaluation : Dice Coefficient

D. Loss Function : Dice Loss + Binary Cross Entropy
Dice Loss

Binary Cross Entropy

3. Postprocessing

위의 방식대로 학습을 진행한 후 테스트 데이터셋으로 검증할 때는 Test-Time Augmentation(TTA)를 진행한다.

이는 테스트 할 때 테스트 데이터셋을 증강하여 모델을 평가하는 방식으로, 증강한 데이터로 나온 예측값들을 평균화하여 예측하는 방식이다.

Figure 5. Test-Time Augmentation (출처 :  https://www.youtube.com/watch?v=xlC385lEzIY&t=963s)

본 게시글에서는 전처리 때 사용한 좌우반전, 상하반전, 90도 회전, 270도 회전을 한 후에 평균을 낸 뒤에 픽셀값이 0.5를 기준으로 이진화를 진행한다.

아래 Figure 6, 7은 TTA를 진행한 영상이다. 맨 좌측 영상에서 Red circle은 원본 영상에서 제대로 segmenation하지 못하였지만, 맨 우측 영상에서 red circle은 TTA를 통해 분할하지 못한 부분을 메꿔지는 효과를 시각적으로 확인 할 수 있다.

FIgure 6. TTA 1

Figure 7. TTA 2

4. Result and Conclusion

4.1. Result - Performance

Before TTA : 0.96349
After TTA : 0.96679

4.2. Conclusion

- TTA를 하면 미세하지만 성능이 증가한다.
- 적은 데이터셋으로 인한 충분하지 않은 학습이 일어난다.

FIgure 8. 예측이 잘 안되는 데이터

- 또한, 테스트한 성능에 대해 확실한 보장을 할 수가 없다.
- 그로므로 더 많은 데이터셋이 필요하다.


Code : https://github.com/yhu0409/-Segmentation-of-the-lungs-on-a-Chest-X-ray-image/tree/main

GitHub - yhu0409/-Segmentation-of-the-lungs-on-a-Chest-X-ray-image