2022년 Journal of Imaging에 실린 "Weakly Supervised Polyp Segmentation in Colonoscopy Images Using Deep Neural Networks" 논문입니다. 이전에 리뷰했던 image-level label 데이터를 가지고 한건 아니고, Bounding Box 정도의 weak annotation을 이용해서 그 안에서 segmentation을 진행했습니다.
1. Introduction
- Introduction에서는 암으로 발전 가능한 대장 용종을 전문의가 찾아내지 못함으로 인한 위험과 사망률 등을 통계 수치로 제시하고 있습니다. 이에 따라 대장내시경 검사의 효율을 높이기 위해서 고화질, NBI (Narrow-band Imaging), 확대 내시경, 딥러닝 기반의 진단보조 기술 등이 출현했다고 합니다.
- CNN 기반의 detection, segmentation, classification 등의 기법들이 전문의들의 진단을 보조하고, 때로는 사람보다 높은 정확도를 가질만큼 고도의 연구가 진행되고 있다는 점을 언급합니다.
- 대장내시경의 진단 보조를 위한 딥러닝에는 총 세가지 카테고리가 있다고 말합니다. a) polyp classification, b) polyp detection, c) polyp segmentation입니다. 각 분야에서 어떤 연구들이 어떤 성능을 보였는지 제시하고 있는데 생략하겠습니다.
- 지금까지 공개된 segmentation 분야의 연구들은 real-time에 쓰기엔 부적합하고, pixel-level annotation을 해야하는 것이 굉장히 힘들다고 합니다. 따라서 본 연구에서는 U-Net model과 bounding box 정보만 가지고 real-time WSSS를 진행했습니다.
- main contribtuion은 다음과 같습니다.
- bounding box만을 이용한 segmentation 수행
- 45 FPS 정도의 real-time 성능
2. Materials and Methods
2.1. Deep Learning Architecture
- up-sampling block, down-sampling block이 각각 5개인 U-Net을 사용했고, VGG16 + ImageNet weight를 첫 weight로 사용했습니다.
2.2 Iterative Weakly Supervised Training
bounding box를 기반으로 한 initial segmentation mask (seed area)가 생성이 되어야했고, 3가지 방법으로 시도를 했습니다.
- bounding box의 중심원을 seed area로 설정. 이후 object를 채워나가되, bbox의 경계를 넘어가진 못함
- 공개 데이터셋에 pre-train 된 segmentation 모델을 이용해 prediction 후 bbox를 넘어가는 것은 잘라냄
- 1번과 2번을 조합한 방법: 2번으로 생성한 마스크가 bbox의 30% 이상 차지하지 못한다면, 1번 방법으로 생성
이 후, seed mask를 모델을 이용해 refine -> refine 한 mask를 재학습 -> 다시 prediction 후 refine의 반복 학습을 거칩니다.
2.3. Maksed Loss Function
아래 알고리즘으로 도출한 Mask M과 DNN의 prediction P, Label Y를 이용해 loss를 계산합니다.
2.4. Dataset and Preprocessing
총 네개의 데이터셋을 사용했습니다.
- ImageNet (VGG16 pre-train 전용)
- CVC-ClinicDB: 612 images with GT mask
- Kvasir-SEG: 1000 images with GT mask & bbox
- Polyp-Box-Seg (연구진만의 데이터): 4070 images with bbox
사용한 데이터 증강 기법은 다음과 같습니다.
- random translation
- zoom in and out
- random mirroring
- color shift
2.5. Model Training
위 데이터를 가지고 연구진은 여러 가지 모델을 학습했습니다.
- Full-Sup-1-VGG: CVC-ClinicDB에 대해 학습한 fully-supervised 모델
- Full-Sup-2: Polyp-Box-Seg에 대해 학습한 fully-supervised 모델
- Weak-Sup-Box-CI: 1번방법으로 학습시킨 weakly-supervised 모델
- Weak-Sup-Box-PI: 2번방법으로 학습시킨 weakly-supervised 모델
- Weak-Sup-Box-HI: 3번방법으로 학습시킨 weakly-supervised 모델
- Weak-Sup-Mix: bounding box를 통해 생성한 마스크와 pre-trained segmentation mask를 둘 다 학습한 모델
다양한 실험 전제와 각 데이터 셋에 대해 어떻게 학습을 했는지 (하이퍼 파라미터 설정) 설명했으나 생략하겠습니다.
3. Results
연구진은 실제 inference 영상을 게시하기도 했습니다.
https://www.igb.uci.edu/colonoscopy-ai-for-gi2/
AI for GI2 – UCI Institute for Genomics & Bioinformatics
Real Time Polyp Contour Segmentation in Colonoscopy using Deep Learning A Video Demonstration of Real-Time Polyp Segmentation in Colonoscopy using Weakly Supervised Neural Networks Examples of AI polyp contour segmentation in colonoscopy videos The video i
www.igb.uci.edu
4. Discussion & Conclusion
- 실험 결과 Bounding Box annotation을 이용한 WSSS 학습법이 직접 labeling 한 fully-supervised model의 성능과 큰 차이 없다는 점을 시사했습니다. 본 연구 결과가 annotating 과정에서 드는 노력/비용을 크게 줄일 수 있다고 말합니다.
- 연구진들은 이 연구 결과를 이용해 polyp을 탐지할 뿐만아니라, automated report 생성까지 진행할 수 있을 것 같다고 말합니다.
