[Review] 3D Context Enhanced Region-based Convolutional Neural Network for End-to-End Lesion Detection (3DCE)

2018.07.29에 발표된 후 의료영상처리학회인 MICCAI 2018에 accepted된 논문 중 하나인

'3D Context Enhanced Region-based Convolutional Neural Network for End-to-End Lesion Detection'를

리뷰하도록 하겠습니다.

※ Paper Link : [Link] / Github Link : [Github]

Abstract

- CT scan에서 병변(lesion)을 찾는 것은 중요한 일이지만 병변이 아닌 부분(non-lesion)과 실제 병변(true lesion)이 유사하게 보이기 때문에 어려운 일입니다. 이 두 가지를 구분하는데 중요한 역할을 하는 것이 3D 정보(context)입니다.

- 이 논문에서는 2D 이미지로부터 얻은 feature maps을 합치는 방법을 통해 3D 정보를 활용하는 3D context enhanced region-based CNN (3DCE)를 제안합니다. 3DCE는 학습과 inference 모두 end-to-end로 가능하고 모든 종류의 병변(lesion)을 찾을 수 있는 universal lesion detector입니다. [논문 링크와 저자가 공개한 코드는 이 글의 상단에 있습니다.]

1. Introduction

- Abstract에서도 언급했듯이 3D 정보(context)는 병변과 병변이 아닌 영역을 구분하는 것에 중요한 역할을 합니다. 따라서 기존의 방법은 단순히 2D 이미지가 아닌 3D volume을 input으로 주었습니다. 하지만 직접적으로 3D volume을 input으로 주게 되면 메모리가 많이 소모된다는 문제(주로 작은 3D patches를 만드는 방법으로 해결)와 3D bounding box는 annotate하기 어렵다는 문제가 있습니다.

- 3D patch 또한 2D slice에 비하면 메모리를 많이 차지하기 때문에 연속된 3개의 단면을 네트워크의 input으로 준 방법도 제시되었습니다. 하지만 해당 논문에서 제시한 방법은 병변 후보(lesion candidates)들이 병변인지 최종적으로 결정하기 위해서 독립적인 CNN 분류 모델을 사용했기 때문에 two stages model로 end-to-end하지 않다는 단점이 있습니다.

- 이 논문에서는 위 방법들에서 제기된 문제점들을 모두 해결한 3D context enhanced region-base CNNs (3DCE)를 제시합니다. 2D 네트워크에 여러 개의 이웃한 단면들이 주어지면 각각의 feature maps이 만들어지고 이 feature map들은 최종 결정을 위해 합쳐집니다. 그리고 R-FCN(region-based fully convolutional network)을 향상시켰습니다. 이러한 방법들을 거친 3DCE는 아래의 장점들이 있습니다.
  1) 학습과 inference를 end-to-end로 할 수 있다.

  2) VGG16과 같은 2D CNN backbone과 pretrained weights를 사용할 수 있다.
  3. 3D보다 annotation하는 것이 쉬운 2D bounding box를 요구합니다.

- 이 논문에서는 규모가 매우 크고 다양한 종류의 병변(lesion) 데이터셋인 DeepLesion을 사용합니다. DeepLesion 데이터셋을 이용하여 여러 종류의 병변을 찾을 수 있는 universal lesion detection algorithm을 제시합니다.

2. Method

- 3개의 단면을 그룹으로 묶어 3-channel 이미지를 만든 후, 각각의 이미지에서 feature map을 뽑아냅니다. 그 후 3D 정보를 얻기 위해서 여러 개의 이웃한 이미지에서 얻은 feature maps을 합쳐줍니다. 이렇게 얻어낸 fused feature maps에서 bounding box에 대한 정보를 얻어냅니다.

- 3DCE에서는 널리 쓰이는 Faster R-CNN 대신에 더욱 빠르고 정확하며 메모리 효율적인 R-FCN을 이용했습니다. R-FCN은 물체의 클래스 정보와 위치 정보를 담고 있는 position-sensitive score maps로 구성되고 이러한 score maps는 이후의 position-sensitive region of interest (PSROI) pooling layer에서 summarize됩니다. 최종 bounding box에 대한 결과는 다른 pooling 연산에 의해 얻어집니다.

- 3DCE는 PSROI pooling 뒤에 3개의 새로운 레이어를 추가함으로써 R-FCN을 향상시켰습니다. 추가한 3개의 레이어는 다음과 같습니다: 2048D FC layer, ReLU 레이어, 그리고 classification과 bounding-box regression을 위한 두 개의 레이어. 3개의 추가적인 레이어덕분에 0.7%만큼 성능을 향상할 수 있었다고 합니다.

- Feature map의 resolution을 향상시켜 작고 spare한 병변들을 예측하기 위해 VGG16의 마지막 두개의 pooling 레이어 (pool4와 pool5)를 제거했습니다. 이를 통해 PSROI pooling 뒤에 레이어를 3개 추가해주었음에도 불구하고 기존의 R-FCN과 크게 차이나지 않는 속도를 보입니다.

- Fig 1.에 나타나있듯이 M개의 3-channels 이미지가 그룹화됩니다. 학습할 때, 중간의 단면이 ground-truth bounding-box를 가지고 다른 단면들은 3D 정보를 제공합니다. M개의 3-channels 이미지가 (VGG16의 Conv1-5의) convolution 연산을 거쳐 M개의 feature maps을 만듭니다. Lesion proposals를 생성하기 위해서 오직 가운데의 단면에서 얻은 feature map만 region proposal network(RPN)로 전달됩니다. 나머지 feature maps들은 모두 추가적인 convolution layer (conv6)를 거친 후 concatenate되어 S2DM-channel feature map을 형성합니다. (S = pooled feature map의 크기 / D = 2D feature maps의 개수 / 논문에서는 S = 7, D = 10으로 설정)

- 네트워크에 인풋으로 주어지는 3-channels 이미지의 개수를 뜻하는 M이 클수록 많은 정보가 네트워크에 주어지지만 메모리 소모가 크고 오버피팅(over-fitting)의 위험이 있습니다. 논문에서는 M을 3에서 9 사이의 값으로 설정하는 것이 바람직하다고 합니다.

- Implementation Details
  1) Deep Learning Framework로 MXNet 사용

  2) Conv1-Conv5의 weights로 ImageNet pretrained VGG-16의 weights를 사용

  3) 5가지 Anchor box scales (16, 24, 32, 48, 96)를 사용했고 3개의 anchor ratios (1:2, 1:1, 2:1) 사용

  4) Improved R-FCN에서 bounding-box regression의 loss weights을 10으로 설정

  5) M < 7일 때는 batch-size를 2로 설정하였으며 M >= 7일 때는 1로 설정 (각각의 sample은 M 3-channel images로 구성)

  6) SGD optimizer 사용 (초기 learning rate = 0.001 → 4번째와 5번째 epochs 이후에는 10배씩 감소)

3. Experiments

- DeepLesion Dataset에 대한 전처리는 아래와 같습니다.
  1) -1024 ~ 3071 HU 값을 갖는 이미지의 픽셀값들을 [0, 255]의 범위 안의 값으로 rescaling

  2) 각각의 픽셀이 0.8mm에 해당하도록 resizing

  3) 각 단면 간의 거리가 2mm가 되도록 설정

  4) 데이터셋 분할 : Training (70%) - Validation (15%) - Test (15%)

  5) Data augmentation을 사용하지 않음

4. Conclusion

- 3DCE는 3D patch가 아닌 2D 단면을 input으로 받으면서도 각각의 2D 단면에서 얻어낸 feature maps을 concat해주는 방법을 통해 3D 정보를 이용했다는 특징이 있습니다. 이러한 방법을 통해 기존의 3D 모델과는 달리 메모리적으로 효율적이고 빠르다는 장점을 보입니다. 그리고 모든 과정이 학습과 inference 과정에서 end-to-end로 진행된다는 장점도 있습니다.