현재 기준 (2022년 10월 24일) paperswithcode Medical Image Segmentation 분야에서 SOTA를 기록한 알고리즘입니다. 최근에 내시경 영상을 이용한 병변 detection & segmentation을 연구하고 있어 읽어보았습니다. Mix Transformer와 feature-pyramid 모델을 결합해 real-time segmentation이 가능한 모델을 제시했습니다.
ABSTRACT
- 폐암은 세계적으로 가장 치명적인 질병이기때문에 조기에 발견하는 것이 중요함
- 기관지 내시경(bronchoscopy)으로 이를 탐지하는 게 가장 효율적임
- 전문의들은 white-light bronchoscopy (WLB), autofluorescence bronchoscopy (AFB), narrow-band imaging (NBI) 세 가지 영상을 봐가면서 판독을 해야 하지만, 현실적으로 힘듦
- 따라서 본 논문에서는 기관지 내시경 비디오를 동기화해 멀티 모달 분석을 가능하게 하는 프레임 워크를 소개함
- 여러 public dataset에서 학습, 일반화 평가에서 SOTA를 달성함
1. INTRODUCTION
- 폐 용종은 폐 암으로 발전할 가능성이 높아 조기에 발견해야 함
- 그런 용종을 발견하기 위한 방법은 기관지 내시경으로 점막을 관찰하는 방법밖에 없음
- Autofluorescenes(형광) 영상으로 관찰하는 것이 가장 효과적
- 그러나 저 AFB 영상으로 직접 눈으로 병변을 검출하는 것은 피곤하고 실수가 많을 작업임
- 몇몇 AFB 영상을 이용한 computer-based 병변 분석이 연구되었지만, 한계점이 세 가지 있음
1. 전처리 과정이 복잡함
2. 정확성 부족, 매 프레임 segmentation 하는 것이 불가
3. AFB 영상을 real-time으로 입력받기가 힘들어 실시간 진단에서 쓰이기 어려움
- 따라서 본 논문에서 AFB 영상을 이용한 real-time detection 및 segmentation을 위해 딥러닝 기반 아키텍처를 구축함 (별다른 이미지 enhancement 없이)
- pretrained Mix Transformer 인코더를 백본으로, stage-wise 한 피쳐 피라미드를 디코더로 함
2. METHODS
2.1 Backbone
- Mix transformer encoder (MiT)는 vision Transformer 네트워크의 아이디어를 차용한 것임. 4개의 중첩된 self-attention 네트워크를 사용함
- 이 중첩 네트워크는 high-resolution 피쳐와 low-resolution 피쳐를 둘 다 제공함
- MiT 인코더는 ImageNet 데이터베이스를 학습시킨 모델
2.2 Efficient Stage-wise Feature Pyramid (ESFP)
- U-net 같은 지난 SOTA 알고리즘들은 Multi-level의 피쳐를 통합하는 데에 의존함. 이는 네트워크가 필요 없는 local feature들을 학습하게 해 좋지 않음
- 이를 해결하기 위해 SSFormer (피쳐 전처리를 위해 컨볼루션 레이어를 결합한 모델)라는 모델도 나왔으나 real-time segmentation에는 부적합함
- ESFP는 가볍고 channel-wise 한 피쳐 피라미드 네트워크인 CfpNet에서 영감을 받음
- ESFP는 linear prediction으로 시작해서 global에서 local까지 linear fuse가 진행됨. fuse 된 피쳐들은 마지막 segmentation에서 함께 이용됨
3. VALIDATION EXPERIMENTS
3.1 Implementation details
- MiT-B0은 RTX 3090 환경에서, MiT-B1~5는 TESLA A100에서 학습시킴
- 이미지를 352 x 352로 리사이즈함
- random flipping, rotation, brightness 적용
- weighted IOU loss와 weighted binary cross-entropy를 결합한 loss function
- AdamW optimizer, learning rate 0.0001, epoch 200
3.2 Dataset
3.2.1 AFB dataset
- 직접 수집한 150 AFB 프레임, 10명의 환자를 대상으로 함. 전문가가 레이블링 해줌
- 6:2:2로 데이터셋 분할
3.2.2 Polyp segmentation datasets
- 모델의 learning ability와 robust generalization capability를 연구하기 위해 3가지 테스트를 진행함
- 여기서 언급되는 데이터셋은 모두 대장 용종 데이터
1. Learning ability experiment
Kvasir와 CVC-ClinicDB를 이용, 8:1:1로 나누어 결과 관찰
2. Generalization capability experiment
Kvasir와 CVC-ClinicDB의 90%를 학습하고, CVC-ColonDB와 ETIS-LaribPolypDB 데이터 모두 테스트에 사용
3. Power balance experiment
2번과 같은 방법으로 학습시킨 뒤, Kvasir, CVC-ClinicDB의 남은 10%와, CVC-ColonDB, CVC-T, ETIS-LaribPolypDB의 모든 데이터를 테스트에 사용
4. RESULT
- MiT 인코더의 scale 에따라, ESFPNet-T, ESFPNet-S, ESFPNet-L를 각각 제시함. (MiT-B0, -B2, -B4) 사용
4.1 Result analysis on AFB dataset
- mean Dice와 IOU를 이용해 성능 평가. FLOPs도 계산해 연산을 얼마나 하는지 검사함
- 더 적은 GFLOPs임에도 SSFormer와 CARANet보다 Mean Dice가 높음
- 평균적으로 27 FPS가 나옴
4.2 Quantitative comparison on Polyp datset
- CVC-ClinicDB에서는 바로 SOTA를 찍어버림
- 좋은 Generalization capability를 보임
5. CONCLUSION
- ESFPNet이 AFB 영상에서 좋은 segmentation 성능을 보였고, 실시간 기관지 내시경 검사에도 쓰일 수 있는 결과를 보였음
- 연구진들이 아는 한, AFB 영상에서의 automatic real-time segmenation 연구는 이 연구가 처음임
- 다른 public 데이터셋에 대한 결과를 통해 medical image segmentation 분야에서 strong capability를 보임
https://arxiv.org/abs/2207.07759
https://paperswithcode.com/paper/esfpnet-efficient-deep-learning-architecture