SKU-Net: Improved U-Net using Selective Kernel Convolution for Retinal Vessel Segmentation

Abstract

In this paper, we propose a deep learning-based retinal vessel segmentation model for handling multi-scale information of fundus images. we integrate the selective kernel convolution into U-Net-based convolutional neural network. The proposed model extracts and segment features information with various shapes and sizes of retinal blood vessels, which is important information for diagnosing eye-related diseases from fundus images. The proposed model consists of standard convolutions and selective kernel convolutions. While the standard convolutional layer extracts information through the same size kernel size, The selective kernel convolution extracts information from branches with various kernel sizes and combines them by adaptively adjusting them through split-attention. To evaluate the performance of the proposed model, we used the DRIVE and CHASE DB1 datasets and the proposed model showed F1 score of 82.91% and 81.71% on both datasets respectively, confirming that the proposed model is effective in segmenting retinal blood vessels.

본 논문에서는 안저영상의 다중 스케일 정보를 다루기 위한 딥러닝 기반의 망막 혈관 분할 모델을 제안한다. 제안 모델은 이미지 분할 딥러닝 모델인 U-Net과 선택적 커널 합성곱을 통합한 합성곱 신경망으로 안저영상에서 눈과 관련된 질병을 진단하는데 중요한 정보가 되는 망막 혈관의 다양한 모양과 크기를 갖는 특징 정보를 추출하고 분할한다. 제안 모델은 일반적인 합성곱과 선택적 커널 합성곱으로 구성된다. 일반적인 합성곱 층은 같은 크기 커널 크기를 통해 정보를 추출하는 반면, 선택적 커널 합성곱은 다양한 커널 크기를 갖는 브랜치들에서 정보를 추출하고 이를 분할 주의집중을 통해 적응적으로 조정하여 결합한다. 제안 모델의 성능 평가를 위해 안저영상 데이터인 DRIVE와 CHASE DB1 데이터셋을 사용하였으며 제안 모델은 두 데이터셋에 대하여 F1 점수 기준 82.91%, 81.71%의 성능을 보여 망막 혈관 분할에 효과적임을 확인하였다.

I. Introduction

망막 혈관의 구조는 고혈압, 동맥 경화증, 당뇨병의 미세혈관 합병증인 당뇨병 망막증과 같은 망막 질환을 발견하는데 중요한 정보가 된다. 안구 검사를 위해서는 신경과 전문의, 심장 전문의, 안과 전문의 및 기타 망막 혈관 질환 전문가의 참여와 작업이 필요하지만, 끊임없이 증가하는 망막 이미지를 처리하는 데 한계가 있다. 과거 많은 연구는 이러한 한계를 극복하기 위해 망막 혈관 구조를 검출하는 자동화 기술을 제안하였다.[1, 2, 3]. 그러나 망막 혈관이 미세하고 안저영상의 배경과 큰 차이가 없어 망막 혈관 검출 자동화 기술은 많은 어려움이 있다.

최근 발전한 딥러닝은 컴퓨터비전 분야에서 뛰어난 성능을 보였다. 특히 이미지넷 챌린지 데이터셋에서 우승한 AlexNet[4]이 등장하면서 많은 CNN 모델들이 제안되었다. 이에 따라 망막 혈관 검출 또한 딥러닝을 이용하는 방법이 제안되었다[5, 6, 7]. 그리고 이미지 분할 모델인 U–Net[9] 기반의 방법들은 각각 안저영상 40장과 28장으로 구성된 DRIVE, CHASE DB1과 같은 적은 수의 데이터에도 뛰어난 망막 혈관 분할 성능을 보였다. 하지만 U-Net 기반의 방법들은 대부분 3×3 합성곱 층을 사용한다. 3×3 합성곱 층으로 이루어진 인코더, 디코더는 같은 수용 영역 크기를 가지고 있어 고정적인 스케일의 공간 정보를 다룬다. 본 논문에서 다중 합성곱 브랜치를 사용하여 동적인 수용 영역을 통해 정보를 추출하고 추출된 다중 스케일의 공간 정보들을 각각 다른 가중치를 적용하여 결합하는 U-Net 기반 모델을 제안한다.

본 논문의 구성은 다음과 같다. 2장에서 망막 혈관 검출 및 분할에 관한 이전 연구를 소개한다. 3장에서 제안모델에서 사용한 선택적 커널 방법을 소개하고 제안모델의 자세한 내용을 설명한다. 4장에서는 DRIVE와 CHASE DB1 데이터셋을 이용한 실험을 통하여 제안 모델의 우수성을 입증하고 마지막 5장에서 결론을 맺는다.

II. Related works

1. Previous Works

망막 혈관 분할을 목적으로 하는 초기 방법에는 픽셀의 밝기나 지식기반으로 설계된 특징들을 이용한 방법들이 있다. Akram 등[2]은 웨이블릿 변환을 기반으로 이미지에서 혈관의 특징 향상 기술을 제안하였다. Onkaew 등[3]은 이미지의 그레디언트 방향을 사용하여 망막 혈관을 분할하는 방법을 제안하였다. 그리고 지도 학습 기반의 방법에서는 픽셀의 특징들을 추출하고 레이블 된 데이터를 이용하여 특징을 분류한다. Marin 등[1]은 픽셀 분류를 위해 신경망을 사용하고 픽셀 표현을 위해 그레이 레벨 및 모멘트 기반 특징으로 구성된 7D 벡터를 계산하였다. Aslani 등[8]은 여러 특징 추출기를 통해 계산된 하이브리드 특징 벡터로 픽셀을 특성화하고 랜덤 포레스트 분류기를 이용하여 이진 분류하는 분할 방법을 제안하였다.

2. U-Net based Retinal Vessel Segmentation

최근 딥러닝이 발달하면서 이미지 분할 분야에서 완전 합성곱 신경망(fully convolutional network)[8] 기반의 모델들이 등장하고 높은 성능을 내고 있다. Ronneberger 등이 제안한 U-Net[9]은 U자 모양의 FCN 모델로 다양한 이미지 인식 분야에 사용되는 모델이다[10,11]. U-Net은 공간 차원을 줄이는 인코더와 공간 차원을 늘리는 디코더로 구성된다. 또한 U-Net의 디코더에서는 같은 공간 차원을 갖는 낮은 수준(low-level)의 특징 맵과 높은 수준(high-level)의 특징 맵을 결합한다. 최근 망막 혈관 분할을 비롯한 의료 영상 분야에서 U-Net 기반의 다양한 모델들이 제안되었다. Residual U-Net[12]은 깊은 모델의 학습과 특징 맵의 정보 전파를 위해 사전 활성화 잔차(pre–activation residual)구조로 이루어진 모델이다. R2U–Net[13]은 반복 잔차(recurrent residual)구조를 활용하여 U-Net과 동일한 파라미터에서 향상된 성능을 내는 모델이다. 그리고 DenseBlockU-Net[14]은 U–Net의 합성곱 블록을 밀집 블록(dense block)으로 변형하여 학습 과정에서 발생하는 기울기 소실 및 폭발 문제를 완화하고 계층 간의 특징 전파를 강화한 모델이다. LadderNet[15]은 체인구조의 다중 U-Net으로 이루어진 모델로 다중 인코더-디코더 쌍을 이용하여 성능을 높였다. IterNet[16]은 U-Net과 여러 mini-UNet의 반복적인 구조로 분할 결과를 반복적으로 분할하여 분할 결과에 가려진 혈관의 미세한 부분을 찾아낸다.

이러한 많은 U-Net 기반의 모델들은 많은 발전을 이루었지만, 대부분의 계층이 3 × 3 커널 크기의 합성곱 층으로 이루어져 있다. 이는 모델의 수용 영역의 크기를 제한하고 안저영상의 공간 정보를 고정적으로 다루어 크기와 모양이 다양한 망막 혈관을 검출하는 데 한계가 있다. DU–Net[17]은 이러한 한계를 해결하기 위해 변형 가능한 합성곱(deformable convolution)과 변형 가능한 관심 영역 풀링(deformable ROI pooling)을 이용하여 안저영상에서 실제 객체의 크기에 따라 수용 영역을 조정하였다.

본 논문은 DUNet[17]에서 더 나아가 다른 커널 크기의 브랜치들을 통해 안저영상의 다중 스케일 공간 정보를 추출하고 다중 스케일 공간 정보를 분할 주의집중 모듈을 통해 적응적으로 조정하여 결합하는 방법을 사용한다.

III. The Proposed Scheme

이번 장에서는 제안 모델의 핵심인 Selective kernel networks(SK-Net)[18]에서 제안한 선택적 커널 합성곱(selective kernel convolution)을 소개하고 제안모델에서 사용한 선택적 커널 합성곱 블록을 설명한다. 또한 U-Net[9]과 선택적 커널 합성곱 블록을 통합한 제안 모델 Selective kernel U-Net(SKU-Net)의 구조를 설명한다.

1. Selective kernel convolution

선택적 커널 합성곱은 Selective kernel networks[18] 에서 제안한 방법으로 동적으로 2i+1의 커널 크기를 갖는 다중 브랜치와 분할 주의집중(split-attention)으로 이루어진 모듈이다. Fig 1은 제안 모델에서 사용된 선택적 커널 합성곱 블록 구조이다. 제안 모델의 선택적 커널 합성 곱은 먼저 3×3커널 크기의 합성곱 층이 있고, [18] 과유사하게 두 개의 브랜치와 분할 주의집중(split– attention) 모듈로 구성된다. 두 개의 브랜치는 3×3 커널 크기의 합성곱 층과 5×5 커널 크기의 합성곱 층으로 이루어진다. 또한 선택적 커널 합성곱 블록은 ResNet[19] 의 잔차(residual) 구조를 따른다. 선택적 커널 합성곱 블록에는 합성곱 층 다음으로 드롭블록(Dropblock)[20], 배치 정규화(Batch normalization)[21], ReLU 활성화 함수를 추가한다. 드롭블록(Dropblock)[20]은 무작위로 활성화 단위를 선택하고 그 주변 영역의 활성화 단위들을 같이 생략하여 모델의 과적합을 완화하는 효과가 있다.

Fig. 1. Selective Kernel Convolution Block

입력 특징 맵 \(X \in \mathbb{R}^{H \times W \times C^{\prime}}\)가 주어졌을 때, 3×3 커널 크기의 합성곱 층에서 특징 맵 \(\bar{X} \in \mathbb{R}^{H \times W \times C}\)를 추출한다. 다음으로 3×3 커널 크기의 합성곱 층과 5×5 커널 크기의 합성곱 층에서 각각 특징 맵 \(\widetilde{U} \in \mathbb{R}^{H \times W \times C}\) 와 \(\hat{U} \in \mathbb{R}^{H \times W \times C}\)을 추출한다. 수식은 다음과 같다.

\(\begin{aligned} &\bar{X}=\bar{F}(X), \\ &\tilde{U}=\tilde{F}(\bar{X}), \\ &\widehat{U}=\hat{F}(\bar{X}), \end{aligned}\)       (1)

식 (1)에서 \(\bar{F}(\cdot)\)와 \(\tilde{F}(\cdot)\)는 3×3 커널 크기의 합성곱 연산이고 \(\tilde{F}(\cdot)\)는 5×5 커널 크기의 합성곱 연산이다. 다중 브랜치에서 추출한 특징 맵은 분할 주의집중 모듈을 통해 브랜치 별로 가중치를 구한다. 제안 모델에서 사용한 분할 주의집중 모듈은 [18]의 방법과 유사한 방법으로 다른 스케일의 정보를 적응적으로 조정하는 역할을 한다.

Fig 2는 제안 모델에서 사용한 분할 주의집중 모듈이다. 분할 주의집중 연산을 위해 먼저 다중 경로에서 생성된 특징 맵들을 원소별 합을 통하여 특징 맵 \(U \in \mathbb{R}^{H \times W \times C}\)를 구한다. 수식은 다음과 같다.

Fig. 2. Split-Attention

\(U=\tilde{U}+\check{U}.\)       (2)

다음으로 대표 특징 맵 U를 전역 평균 풀링(global average pooling) 연산을 통해 채널별로 공간 차원 정보를 압축하여 특징 벡터 \(s\in \mathbb{R}^{d\times1}\)를 구한다.

\(s_{c}=\frac{1}{H \times W} \sum_{i=1}^{H} \sum_{j=1}^{W} U_{c}(i, j)\)       (3)

식 (3)에서 U_c는 U의 c번째 채널이다. s_c는 c번째 채널에서 계산된 하나의 벡터로 특징 맵의 공간 차원 정보를 압축한 정보이다. 채널별 대표 정보 벡터 s는 더 작은 채널 차원인 벡터 \(z\in \mathbb{R}^{d\times1}\)로 변형한다.

z=Ws,       (4)

식 (4)에서 \(W\in \mathbb{R}^{d\times C}\)는 완전 연결 층(fully connected layer)의 가중치이다. 는 압축을 위한 매개변수로 다음식을 통해 구한다. 본 논문에서는 r=8, L=32를 사용하였다.

d=max(C/r, L). (5)

압축 벡터 z를 다시 원래의 차원으로 변형한 후 소프트맥스 연산을 하여 커널 크기별로 다른 가중치를 부여한다. c번째 요소에 대한 계산식은 다음과 같다.

\(\begin{aligned} &\widetilde{a}_{c}=\frac{\exp \left(\widetilde{A_{c}} z\right)}{\exp \left(\widetilde{A}_{c} z\right)+\exp \left(\widehat{A}_{c} z\right)}, \\ &\hat{a}_{c}=\frac{\exp \left(\widehat{A}_{c} z\right)}{\exp \left(\widetilde{A}_{c} z\right)+\exp \left(\widehat{A}_{c} z\right)}, \end{aligned}\)       (6)

식 (6)에서 \(\tilde{A},\hat{A} \in \mathbb{R} ^{C \times d}\)는 특징 벡터 z에 선형 결합하는 완전 연결 층(fully connected layer)의 가중치이다. \(\tilde{a}_c\)와 \(\hat{a}_c\)는 소프트 맥스를 통해 계산된 벡터로 브랜치에 따라 다른 가중치를 갖는다. \(\tilde{a}\), \(\hat{a}\)의 합은 1이며 각각 \(\tilde{U}_c\)와 \(\hat{U}_c\)에 결합 후 원소별로 합하여 특징 맵 \(V \in \mathbb{R}^{H \times W \times C}\) 를 구한다. 번째 요소에 대한 계산식은 다음과 같다.

\(V_c=\tilde{a}_c \tilde{U}_c + \widehat{a}_c\widehat{U}_c ,\)       (7)

식 (7)에서 \(C_c \in \mathbb{R}^{H \times W}\)는 c번째 요소에 대한 연산 결과이다. 특징 맵 V는 선택적 커널 합성곱 블록의 입력 X와 더 한 후 ReLU 활성화 함수를 거쳐 최종 특징 맵을 얻는다. 수식은 다음과 같다.

\(\begin{aligned} O &=V+X, \\ Y &=\sigma(O), \end{aligned}\)       (8)

식 (8)에서 \(\sigma\)는 ReLU 활성화 함수이다. 최종적으로 구한 \(Y \in \mathbb{R}^{H \times W \times C}\)는 공간 차원을 조절한 후 다음 인코더나 디코더의 입력이 된다.

2. Network architecture

Fig 3은 제안 모델의 구조이다. 제안 모델은 다양한 수용 영역을 위해 선택적 커널 합성곱 블록(SK Block)이 추가된 인코더와 디코더로 구성되어있다. Fig 4의 왼쪽은 U-Net, 오른쪽은 제안모델의 합성곱 블록이다. U-Net은 3×3 합성곱 층과 ReLU 활성화 함수로 이루어진 블록을 사용한다. 제안 모델의 합성곱 블록은 U-Net에서 사용하는 합성곱 블록에서 드롭블록(Dropblock)[20]과 배치 정규화(BN)[21] 가 추가된다.

Fig. 3. Network Architecture

Fig. 4. Convolution block for U-Net(left), Convolution block for SKU-Net(right)

제안 모델의 인코더와 디코더는 드롭블록이 추가된 합성 곱 블록과 선택적 커널 합성곱 블록으로 구성된다. 인코더에서 공간 차원을 줄이기 위한 방법으로 2×2 최대 풀링(max pooling) 연산을 사용한다. 제안 모델의 디코더는 인코더와 동일한 구조이다. 디코더에서 공간 차원을 증가하기 위한 방법으로 2×2전치 합성곱(transpose conv– olution)을 사용한다. 전치 합성곱은 공간 차원을 증가시키고 특징 맵의 채널을 반으로 줄이는 역할을 한다. 줄어든 특징 맵은 인코더에서 같은 해상도인 특징 맵과 접합되어 다음 디코더의 입력이 된다. 제안 모델의 마지막 특징 맵은 1×1 합성곱 층과 시그모이드 활성화 함수를 거쳐 출력된다. 학습에서 사용되는 손실함수로 이진 교차 엔트로피(binary cross-entropy)를 사용한다.

IV. Experimental Results

1. Dataset

제안 모델의 성능 평가를 위해 DRIVE와 CHASE DB1 데이터셋을 사용하였다. 두 데이터셋은 전문가에 의해 망막 혈관을 전경, 나머지를 배경으로 이진 분할되었다. DRIVE 데이터셋은 훈련을 위한 20장과 평가를 위한 20장을 포함한 40장의 안저영상을 포함한다. CHASE DB1 데이터셋은 28장의 안저영상으로 이루어져있다. CHASE DB1은 훈련 데이터셋과 시험 데이터셋을 구분하지 않으므로 다른 방법들[5, 13, 15, 16, 17, 22, 23]과의 비교를 위해 첫 20장은 훈련 데이터셋으로 나머지 8장은 시험 데이터셋으로 나누었다. DRIVE와 CHASE DB1의 해상도는 각각 565×584, 999×960이고 제안모델의 입력에 맞추기 위해 양선형 보간법을 사용하여 각각 592×592, 1008×1008로 크기 조정하였다. 또한 적은 수의 데이터 문제를 해결하기 위해 RGB 이미지에서 21개, 회색조 이미지로 변환 후 3개의 데이터 증강 기법을 사용하여 회색조변환을 포함한 총 25개의 데이터 증강 기법을 DRIVE와 CHASE DB1 훈련 데이터셋에 적용하였다. 데이터 증강기법에 대한 자세한 내용은 Table 1과 같다.

Table 1. Data Augmentations

2. Evaluation metrics

제안 모델의 정량 평가를 위해 안저영상의 모든 픽셀에 대한 실제 값과 분할 결과를 true positive(TP), false positive(FP), false negative(FN), true negative(TN)로 산출하고 정확도(Accuracy)와 F1 점수를 계산하였다. 정확도(Accuracy)와 F1 점수(F1-score)의 계산식은 다음과 같다.

\(\text { Accuracy }=\frac{T P+T N}{T P+F N+T N+F P}\)       (9)

\(\text { Precision }=\frac{T P}{T P+F P},\)       (10)

\(\text { Recall }=\frac{T P}{T P+F N}\)       (11)

\(F 1-\text { Score }=2 \times \frac{\text { Precision } \times \text { Recall }}{\text { Precision }+\text { Recall }},\)       (12)

모델의 성능 평가하기 위해 추가로 ROC 곡선(AUC)을 계산하였다. AUC 값이 1이면 완벽한 분할을 의미한다.

3. Experiments details

제안 모델 SKU-Net의 성능 평가를 위해 U-Net과 성능을 비교하였다. 또한 드롭블록의 유효성을 검증하기 위해 드롭 블록이 추가된 U-Net을 추가하여 총 세 개의 모델을 실험하였다. 세 개의 모델들은 스크래치에서 DRIVE와 CHASE DB1 데이터셋 모두 50 epoch 동안 학습하였다. 학습을 위한 파라미터로 학습률은 1×10^-3, 배치 크기는 DRIVE와 CHASE DB1 각각 8과 4로 적용하였다. 손실함수에 대한 최적화 알고리즘으로는 아담 최적화 방식 (Adam optimizer)을 사용하고 학습률을 줄이는 스케줄러로 ReduceLROnPlateau를 사용하였다. 본 논문에서 사용한 드롭블록의 파라미터로 블록 크기 ×을 적용하고 keep rate는 DRIVE에서 0.85, CHASE DB1에서 0.80으로 지정하였다. 또한 훈련 데이터셋의 10%는 검증 데이터셋으로 사용하였다. 제안 모델의 프레임워크는 Pytorch로 구현 하였으며, 자세한 실험환경은 Table 2와 같다.

Table 2. Experiment Environment

4. Results

Table 3과 4는 각각 DRIVE, CHASE DB1 데이터셋에대한 실험 모델들의 망막 혈관 검출 성능 결과이다. Table 3과 Table 4에서 DRIVE, CHASE DB1 데이터셋 모두 U-Net+Dropblock이 U-Net보다 F1 점수(F1), 정확도 (Accuracy), AUC에서 향상된 성능을 보인다. 이는 정규화를 위한 드롭블록이 망막 혈관 분할의 성능개선에 효과적임을 보여준다. 제안 모델인 SKU-Net은 U-Net+Drop– block과 비교했을 때 두 데이터에 대한 F1 점수, 정확도, AUC 지표에서 모두 향상된 성능을 보였다.

Table 3. DRIVE Dataset Result

Table 4. CHASE DB1 Dataset Result

Fig 5와 6은 DRIVE와 CHASE DB1 두 데이터셋의 몇 가지 이미지, 제안 방법의 분할 결과, 정답 레이블이다.

Fig. 5. DRIVE example original image(left), prediction(middle), ground truth(right)

Fig. 6. CHASE DB1 example original image(left), prediction(middle), ground truth(right)

5. Comparisons with other methods

Table 5와 6은 각각 DRIVE와 CHASE DB1에 대한 제안 모델과 다른 방법들의 F1 점수, 정확도, AUC 지표에 대한 성능 비교를 보여준다. Table 5에서 제안모델을 다른 방법들과 비교하였을 때 F1 점수(F1), 정확도(Accu– racy), AUC 기준 82.91%, 96.99%, 98.74%로 가장 높은 성능을 보인다. 그리고 Table 6에서 제안모델은 F1 점수, 정확도, AUC 기준 81.71%, 97.68%, 99.09%로 가장 높은 성능을 보인다.

Table 5. DRIVE Dataset Comparisons

Table 6. CHASE DB1 Dataset Comparisons

Ⅴ. Conclusions

본 논문에서는 안저영상에서 망막 혈관 분할을 위한 SKU-Net을 제안한다. SKU-Net은 선택적 합성곱을 사용하여 안저영상에서 다중 스케일의 공간 정보를 추출하고 이를 적응적으로 결합하여 수용 영역을 조정한다. 실험을 통해 선택적 커널 합성곱 블록의 유효성을 확인하였고, 실험 결과 제안 모델은 DRIVE와 CHASE DB1 데이터에서 F1 점수 기준 각각 82.91%, 81.71%의 높은 성능을 보였다. 이는 같은 실험환경에서 DRIVE와 CHASE DB1 데이터에 대하여 U-Net보다 F1 점수 기준 1.71%p, 2.44%p의성능 향상을 보임으로써 제안 모델 SKU-Net이 망막 혈관 분할에 효과적임을 증명하였다.