[논문 리뷰] OR-UNet: an Optimized Robust Residual U-Net for Instrument Segmentation in Endoscopic Images

Fabian Isensee, Klaus Maier‐Hein|arXiv (Cornell University)|2020. 04. 27.

Colorectal Cancer Screening and Detection참고 문헌 6인용 수 17

한 줄 요약

이 논문은 ImageNet 사전 훈련이나 시간적 데이터를 사용하지 않고 Robust-MIS 2019 데이터셋에서 랜덤 초기화 상태에서 훈련한, 인코더에 잔차 블록과 딥 서프비전을 통합한 2D U-Net인 OR-UNet을 제안한다. 5983개의 내 endoscopic 영상에서 8중 교차 검증 시 평균 Dice 점수는 87.41(중위수 94.35)를 기록하며, 사전 훈련 없이도 기구 분할에 뛰어난 성능을 보임을 입증한다.

ABSTRACT

Segmentation of endoscopic images is an essential processing step for computer and robotics-assisted interventions. The Robust-MIS challenge provides the largest dataset of annotated endoscopic images to date, with 5983 manually annotated images. Here we describe OR-UNet, our optimized robust residual 2D U-Net for endoscopic image segmentation. As the name implies, the network makes use of residual connections in the encoder. It is trained with the sum of Dice and cross-entropy loss and deep supervision. During training, extensive data augmentation is used to increase the robustness. In an 8-fold cross-validation on the training images, our model achieved a mean (median) Dice score of 87.41 (94.35). We use the eight models from the cross-validation as an ensemble on the test set.

연구 동기 및 목표

• 로버스트-MIS 2019 데이터셋만을 사용하여 외과 기구의 강건하고 정확한 분할 모델을 개발하는 것.
• 대규모 도메인 전용 데이터셋에서 랜덤 초기화 상태로 훈련하는 것이 ImageNet 사전 훈련된 인코더를 사용하는 모델보다 성능이 뛰어나거나 유사한지 평가하는 것.
• 딥 서프비전과 광범위한 데이터 증강을 통해 기울기 흐름과 일반화 능력을 향상시키는 것.
• 연결 요소 분석을 통한 후처리를 통해 이진 분할 출력을 탐지 및 인스턴스 분할 작업에 효과적으로 재사용할 수 있는지 평가하는 것.

제안 방법

• 모델은 인코더에 잔차 블록과 디코더에 표준 컨볼루션을 사용한 2D U-Net 아키텍처를 사용하며, 특징 맵의 수는 48에서 512로 증가한다.
• 잔차 블록은 원래의 제안 방식을 따르며, 각 컨볼루션 후 배치 정규화와 ReLU를 적용한다.
• 디코더의 다양한 해상도에서 분할 예측을 생성함으로써 딥 서프비전을 구현하며, 각 해상도에 대응하는 손실 함수를 사용한다.
• 손실 함수로는 최고 해상도에서 소프트 Dice 손실과 교차 엔트로피 손실의 합을 사용하고, 저해상도 예측에는 소프트 정답을 고려한 조정된 Dice 손실과 평균 제곱 오차 손실을 적용한다.
• 모델의 강건성을 향상시키기 위해 광범위한 데이터 증강을 적용하며, 모든 모델은 ImageNet 사전 훈련 없이 무작위 초기화 상태에서 훈련된다.
• 탐지 및 인스턴스 분할 작업을 위해 이진 분할 출력은 연결 요소 분석을 통해 전경 객체에 고유 ID를 할당한다.

실험 결과

연구 질문

• RQ1 잔차 연결과 딥 서프비전을 갖춘 U-Net이 ImageNet 사전 훈련된 가중치 없이도 내 endoscopic 기구 분할에서 높은 성능을 달성할 수 있는가?
• RQ2 대규모 도메인 전용 데이터셋에서 랜덤 초기화 상태로 훈련하는 것과 사전 훈련된 인코더를 사용하는 것 간의 분할 정확도에서의 성능 비교는 어떠한가?
• RQ3 딥 서프비전과 데이터 증강이 내 endoscopic 영상 분할 네트워크의 강건성과 기울기 흐름 향상에 어느 정도 기여하는가?
• RQ4 간단한 후처리인 연결 요소 분석을 통해 이진 분할 결과를 탐지 및 인스턴스 분할 작업에 효과적으로 재사용할 수 있는가?
• RQ5 이전 영상 프레임의 시간 정보를 통합하면 이 설정에서 분할 성능 향상에 기여하는가?

주요 결과

• 훈련 세트에서 8중 교차 검증 시 평균 Dice 점수는 87.41(중위수 94.35)를 기록하였으며, 중위수 점수가 대부분의 영상에서 높은 성능을 나타낸다.
• Dice 점수 분포는 점수 0을 기록하는 영상이 다수 존재하는 것으로 나타났으며, 주로 가짜 양성 또는 빠진 탐지로 인해 평균 점수가 크게 낮아졌다.
• 94.35의 중위수 Dice 점수는 대부분의 케이스, 특히 기구가 명확하게 보이는 케이스에서 모델이 매우 우수한 성능을 보임을 시사한다.
• 테스트 세트에서의 성능은 훈련 세트의 교차 검증 점수보다 높을 것으로 예상되며, 이는 테스트 평가에서 정답이 없는 영상가 제외되기 때문이다.
• 입력 채널에 프레임을 스택하여 시간 정보를 통합한 실험은 성능 향상을 가져오지 못했으며, 이는 현재 모델이 단일 쿼리 영상에 의존하는 것으로도 충분하다는 것을 시사한다.
• 이진 예측에 연결 요소 분석을 적용한 결과, 특히 기구가 겹치는 빈도가 낮은 점을 감안할 때 탐지 및 인스턴스 분할 작업에 실용적인 기준이 되는 결과를 제공하였다.