VoxelMorph: A Learning Framework for Deformable Medical Image Registration
We present VoxelMorph, a fast learning-based framework for deformable, pairwise medical image registration. Traditional registration methods optimize an objective function for each pair of images, which can be time-consuming for large datasets or rich defo
arxiv.org
0. Abstract
VoxelMorph은 변형 가능한 pair별 의료 image registration을 위한 빠른 learning-based framework이다.
Image Registration (영상 정합)
공간좌표가 서로 다른 영상에서 연관관계를 쉽게 파악하기 위하여 동일 좌표계로 정렬하는 것으로, 서로 맵핑되는 기하학적 변환을 찾는 과정이라 할 수 있다.
Voxel-based registration (복셀 기반 정합)
전체 영역을 전처리 과정 없이 두 영상의 복셀 간 상관도를 측정한다. 상대적으로 정확도가 높은 대신 모든 복셀들을 가지고 유사도 평가를 하기 때문에 수행시간이 오래 걸린다. 다시 명암도 정보를 대표적인 스칼라값이나 방향 등으로 축소하여 적용하는 중심축 및 모멘트 정합 (principal axes and moments based method)과 정합 과정에서 전체 영상 정보를 모두 사용하는 명암도 기반 정합으로 나뉜다. 명암도 기반 정합에는 동일 모달리티 영상 간 정합을 위한 유사도 측정방법으로 명암도 차이 측정(SSD, Sum of Squares of Intensity Differences)과 상관계수(correlation) 측정이 있다. 서로 다른 모달리티 영상 간 정합을 위한 유사도 측정방법으로 조인트 엔트로피 (joint entropy) 측정 및 상호정보 (mutual information) 측정 방법 등이 있다.
기존의 image registration은 각 paired images에 대해 objective function를 최적화하기 때문에 대규모 dataset나 변형 가능한 model에서 시간이 많이 소요됐다. 최근의 학습 기반 방법으로, registration을 하나의 함수로 보고, input paired image를 넣으면 align된 image가 나오게 된다. CNN을 통해 함수를 parameter화하고 이미지 세트에서 신경망의 parameter를 최적화한다.
두 가지 훈련 방법을 제안한다.
(i) unsupervised learning으로 image간의 matching objective function을 최적화
(ii) auxiliary segmentation map을 학습에 활용
제안 된 unsupervised model은 SOTA 방법들과 비교했을 때 경쟁력 있는 결과를 보여주며 연산 속도는 매우 빠르다. 또한 auxiliary data를 사용했을 때 registration 성능이 올라갔다.
1. Introduction
Volume
의료 영상 처리 및 컴퓨터 비전 분야에서 volume은 일반적으로 3차원 이미지이다. 이는 단일 2D 이미지가 아닌, 여러 개의 2D 단면(슬라이스)을 겹쳐서 형성된 3차원 데이터 구조이다. 예를 들어, CT(컴퓨터 단층촬영)나 MRI(자기공명영상) 스캔에서 얻어진 데이터는 수십에서 수백 장의 2D 슬라이스가 합쳐져 3D 볼륨을 구성하게 된다.
이 volume은 각 복셀(voxel, volume element)로 구성되며, 이는 3D 공간 내의 작은 단위이다. 2D 이미지가 픽셀(pixel)로 구성된 것과 같은 원리이다. 각 복셀은 그 위치의 밀도, 강도 또는 기타 속성을 나타내며, 이를 통해 3차원 구조 내에서 인체의 장기, 조직, 병변 등을 시각화하고 분석할 수 있다.
따라서, 의료 영상에서 volume은 3차원 데이터를 의미하며, 이를 다루는 알고리즘은 복셀 간의 상관관계를 이용해 정합, 분할 및 분석 등의 작업을 수행할 수 있다.
Deformable registration(변형 정합)은 다양한 의료 영상 연구에서 필수적인 작업이며, 수십 년간 활발한 연구 주제였다. 두 개의 3D MRI 뇌 스캔과 같은 이미지 쌍 사이의 밀도 있고 비선형적인 대응을 설정하는 작업이다. 기존의 registration은 유사한 외형을 가진 voxel을 정렬하고 registration mapping에 대한 제약을 부과하면서 각 paired image에 대해 최적화 문제를 해결한다. 그러나 이 과정은 시간 소모가 커서 실제로 느리게 작동한다.
VoxelMorph는 이러한 한계를 극복하기 위해 learning-based의 registration으로, volume(이미지) 집합으로부터 매개변수화된 registration 함수를 학습한다. 이 함수는 CNN을 사용하여 구현되며, n-D input paired volume을 받아 한 volume의 모든 voxel을 다른 volume에 매핑하는 함수를 구한다.
VoxelMorph의 접근법은 개별 테스트 이미지쌍에 대한 cost가 큰 최적화를 전역적인 함수 최적화로 대체하여, 새로운 테스트 스캔 쌍의 정합이 함수를 단순히 평가함으로써 빠르게 수행되도록 한다. 이 방법은 GPU뿐만 아니라 CPU에서도 신속한 정합을 가능하게 하며, registration 프레임워크로서 다양한 loss function를 채택할 수 있다. 이 논문에서는 두 가지 loss function를 제안하는데, 하나는 unsupervised laerning based이고, 다른 하나는 보조적으로 사용할 수 있는 해부학적 세분화를 활용한 방법이다.
2. Background
기존 Volume Registration에서, 하나의 이미지(이동 또는 소스 이미지)는 두 번째 이미지(고정 또는 타겟 이미지)에 맞게 변형됩니다. 논문에 포함된 그림은 3D MRI volume의 2D coronal slice를 보여주며, 해부학적 구조의 경계선을 나타낸다. 피험자 간에 자연스러운 뇌 해부학의 차이 및 건강 상태의 차이로 인해 상당한 변동성이 존재한다. 변형 정합은 스캔 간 구조 비교를 가능하게 하며, 이는 집단 간의 차이나 질병을 앓고 있는 개인의 뇌 해부학의 변화 추적에 유용하다.
이러한 변형 정합 전략은 일반적으로 두 단계로 이루어진다. 첫 번째 단계는 전역 정렬을 위한 affine transformation이며, 두 번째 단계는 더 많은 자유도를 가진 느린 deformable transformation이다. 저자는 deformable transformation에 집중하여 모든 voxel에 대해 dense non-linear correspondence(밀집된 비선형 일치)을 계산한다.
대부분의 기존 변형 정합 알고리즘은 에너지 함수에 기반하여 변환을 반복적으로 최적화한다. 이때, 고정 이미지(f)와 이동 이미지(m), 그리고 f의 좌표를 m의 좌표에 mapping하는 registration field(φ)가 사용된다.
위 수식은 이미지 유사성을 측정하는 함수$L_{sim}$과 변형 필드의 정규화(regularization)를 부과하는 $L_{smooth}$로 구성된 전체 손실 함수를 포함한다.
3. Related Work
A. Medical Image Registration (Non-learning-based)
기존 다수의 연구는 displacement vector field를 최적화하는 elastic type model, statistical parametric mapping, b-splines, discrete methods, demons와 같은 방법들이 있다. 또한, 위상 보존을 위한 diffeomorphic transforms(다형 변환)과 LDDMM(대규모 변형 거리 지표 매핑)과 같은 방법도 존재한다. 하지만 이러한 접근법은 각 image pair에 대해 에너지 함수를 최적화하므로 처리 속도가 느린 단점이 있다. 일부 GPU 기반 알고리즘이 성능을 개선했으나, 여전히 GPU가 필요하다.
B. Medical Image Registration (Learning-based)
Warp Field
이미지 정합 또는 변형에서 사용하는 변환 필드로, 하나의 이미지(주로 이동 이미지)를 다른 이미지(주로 고정 이미지)에 맞춰 정렬하기 위해 픽셀 또는 복셀 단위의 위치를 조정하는 역할을 합니다. 각 픽셀 또는 복셀의 좌표가 이동될 위치로 매핑되며, 이러한 변형 정보의 전체가 warp field로 불립니다.
의료 이미지 정합에서 warp field는 두 스캔 사이의 구조적 차이를 보정하고, 각 픽셀이나 복셀이 고정 이미지의 대응 위치에 정확히 맞도록 변형을 적용하는 데 사용됩니다. VoxelMorph 같은 모델은 이 warp field를 예측하기 위해 학습된 함수 $\phi$를 사용하며, 이를 통해 이동 이미지 $m$을 고정 이미지 $f$와 일치시키도록 변형합니다.
수식으로는 $\phi(x) = x + u(x)$와 같은 형태로 표현되며, 여기서 $u(x)$는 각 좌표 $x$에 대해 적용할 변위(displacement)를 나타냅니다. 이를 통해 이동 이미지가 변형되어 고정 이미지에 맞춰지는 것입니다.
Learning-based medical image registration에 대한 최근 연구들은 신경망을 활용하여 image registration을 학습하는 방법들이 있다. 대부분의 연구는 실제 warp field에 대한 supervised learning 기반(ground truth warp field에 의존한다는 단점을 가짐)으로 하며, 기존 registration 방법을 통해 얻은 warp field를 학습 데이터로 사용한다. 이러한 방법은 데이터 수집이 복잡할 수 있으며, 학습 가능한 변형 유형에 제한을 줄 수 있다. 이에 비해 VoxelMorph는 unsupervised 방식으로, 필요 시 해부학적 분할 데이터를 추가 학습 정보로 활용할 수 있다. 초기의 unsupervised learning based 방법은 3D 서브영역이나 2D 슬라이스만 다루며, 작은 변형만 지원했다.
C. 2D Image Alignment
Optical flow algorithm(광학 흐름 알고리즘)은 두 개의 2D image 간의 작은 displacement를 나타내는 dense displacement vector field를 반환한다. 기존에는 variational method(변분법 : 데이터를 간단하게 표현하면서도 원래 데이터의 중요한 특징을 잘 유지하도록 최적화하는 과정)를 사용하여 최적화 문제를 해결하며, 대규모 displacement를 다루기 위한 appearance include feature based matching(기능 기반 매칭)도 존재한다. 최근에는 신경망을 이용해 2D image registration을 학습하는 방법이 등장했고, 일부는 supervised learning을 필요로 하지만, unsupervised learning 방법도 개발됐다. VoxelMorph는 위 방법을 참고하여, n-D volume registration에 학습 시 분할 데이터를 활용해 정확도를 높이는 방법도 제공한다.
4. Method
A. VoxelMorph CNN Architecture
U-Net
U를 닮은 모델의 구조는 크게 3가지로 나눌 수 있습니다.
- 점진적으로 넓은 범위의 이미지 픽셀을 보며 의미정보(Context Information)을 추출하는 수축 경로(Contracting Path)
- 의미정보를 픽셀 위치정보와 결합(Localization)하여 각 픽셀마다 어떤 객체에 속하는지를 구분하는 확장 경로(Expanding Path)
- 수축 경로에서 확장 경로로 전환되는 전환 구간(Bottle Neck)
모델의 Input은 이미지의 픽셀별 RGB 데이터이고 모델의 Output은 이미지의 각 픽셀별 객체 구분 정보(Class)입니다. Convolution 연산과정에서 패딩을 사용하지 않으므로 모델의 Output Size는 Input Size보다 작습니다. 예를들어 572×572×(RGB) 크기의 이미지를 Input으로 사용하면 Output으로 388×388×(Class) 이미지가 생성됩니다.
Input(Width×Height×RGB)−>Model−>Output(Width×Height×Class)
UNet 기반 합성곱 신경망(CNN)을 기반으로 하며, 입력 volume pair $(f, m)$을 결합하여, 2채널 3D 이미지를 구성한다. 네트워크는 encoder와 decoder로 구성되어 있으며, encoder는 입력 이미지의 차원을 절반으로 줄이기 위해 3D stride convolution을 사용하고, decoder는 upsampling을 통해 원래 크기로 복원한다. encoder와 decoder사이에는 skip connection이 있어, encoder의 특징을 decoder에 전달하여 registration 성능을 향상시킨다.
B. Spatial Transformation Function
Warp field를 학습하기 위해 공간 변환 함수를 사용한다. 각 voxel $p$에 대해 위치 $p'$를 계산하고, 8개의 이웃 voxel의 선형 보간을 통해 $m \circ \phi$를 구한다. 이 과정은 미분 가능하도록 구성되어 있어, 신경망이 역전파를 통해 손실 함수를 최적화할 수 있다. $Z(p')$는 $p'$의 인접 voxel이다.
C. Loss Functions
이 두 손실 함수를 통해 VoxelMorph는 정확한 image registration을 달성하며, 필요한 경우 보조 데이터를 사용하여 registration 성능을 향상시킨다.
(1) Unsupervised Loss Function
비지도 손실 함수 $L_{us}$는 이미지 유사성 손실 함수 $L_{sim}$과 변형 필드의 매끄러움(보간법을 이용하여 빈 픽셀을 매꾸는 것)을 위한 정규화 항 $L_{smooth}$로 구성된다.
- $f$ : 고정 이미지(fixed image)
- $m$ : 이동 이미지(moving image)
- $\phi$ : warp field
- $\lambda$ : 매끄러움 제어를 위한 정규화 매개변수
$L_{sim}$
고정 이미지와 변형된 이동 이미지 간의 유사성을 극대화하는 데 사용된다.
Mean Squared Error (MSE)
고정 이미지 \( f \)와 변형된 이동 이미지 \( m \circ \phi \) 사이의 픽셀 단위 차이를 제곱하여 계산하는 방식이다. 여기서 $\Omega$는 전체 이미지 영역이다.
Local Cross-Correlation (CC)
이 방식은 이미지 간의 국소 영역 내 상관관계를 측정하여, 서로 다른 강도를 가진 이미지들 간의 유사성을 잘 평가할 수 있다. CC 계산은 이미지가 정렬되었을 때 더 높은 값을 가지므로, 손실 함수에서는 음수로 사용하여 이 값을 최소화한다.
$L_{smooth}$
Warp field $\phi$의 공간적 매끄러움을 제어한다. 변형이 너무 급격하게 변화하지 않도록 필드의 공간적 변화율을 최소화한다. $u = \phi - Id$는 변형 필드의 변위(displacement)를 나타내며, $\nabla u(p)$는 해당 변위 필드의 공간적 그래디언트이다.
(2) Auxiliary Data Loss Function
보조 데이터 손실 함수 \( L_{a} \)는 비지도 손실 \( L_{us} \)에 추가적으로 해부학적 분할(특정 구조에 대한 segmentation)을 사용한 분할 손실 함수 \( L_{seg} \)을 포함하여 registration 성능을 향상시킨다.
\[
L_{a}(f, m, s_f, s_m, \phi) = L_{us}(f, m, \phi) + \gamma L_{seg}(s_f, s_m \circ \phi)
\]
- \( s_f \)와 \( s_m \) : 각각 고정 이미지와 이동 이미지에 대한 분할 마스크(segmentation map)
- \( \gamma \) : 분할 손실의 가중치를 조절하는 매개변수
\( L_{seg} \)
분할 손실은 각 해부학적 구조가 정렬되었는지 평가하는 Dice 점수로 정의되며, 각 구조가 정확히 정합되도록 한다. \( K \)는 해부학적 구조의 개수이며, Dice 점수는 각 구조가 얼마나 잘 맞는지 측정합니다. Dice 점수는 두 영역 간의 겹치는 정도를 측정하여 1에 가까울수록 좋은 registration이다.
\[
L_{seg}(s_f, s_m \circ \phi) = -\frac{1}{K} \sum_{k=1}^K \text{Dice}(s_f^k, (s_m \circ \phi)^k)
\]
5. Experiments
A. Experimental Setup
총 3,731개의 T1 가중 뇌 MRI 스캔이 사용되었으며, 모든 스캔은 보간되었다. 이 실험에서는 VoxelMorph와 기존 registration 알고리즘(ANTs, NiftyReg)의 성능을 비교하기 위해 여러 평가 지표가 사용된다. 평가 지표는 Dice 점수를 기반으로 하며, 이는 registration된 해부학적 구조 간의 겹치는 정도를 측정한다.
B. Atlas-based Registration
Atlas-based registration(아틀라스 기반 정합)
특정 기준이 되는 atlas image를 fixed image로 사용하여 다른 대상 이미지들과 정렬하는 정합 방식이다. 여기서 atlas는 여러 사람의 데이터를 평균하여 만든 reference image(참조 이미지)로, 주로 해부학적 구조가 잘 보이는 평균적인 image로 구성된다. 이를 통해 개별 대상 이미지가 atlas image와 정렬되면, 서로 다른 피험자의 뇌 구조를 같은 기준에 맞춰 비교할 수 있게 된다.
- 집단 비교 : 아틀라스를 기준으로 모든 피험자들의 이미지가 정렬되므로, 여러 피험자의 구조적 차이를 비교하기가 용이해진다. 예를 들어, 특정 질병을 앓고 있는 집단과 건강한 집단의 뇌 구조를 비교하는 연구에 자주 사용된다.
- 표준화된 좌표계 제공 : 아틀라스는 모든 정합된 이미지가 같은 좌표계에서 동일한 위치에 정렬되도록 하기 때문에, 정렬 후 각 해부학적 위치가 서로 의미 있는 비교가 가능해진다.
- 효율적인 데이터 분석 : 아틀라스 기반 정합을 통해 개별 이미지의 편차를 줄이고, 같은 위치에서 데이터의 변화를 분석함으로써 해부학적 연구나 임상 연구에서 효율적인 분석을 가능하게 한다.
고정된 이미지(아틀라스)와 각 volume 간의 registration을 수행한다. 이 실험은 인구 기반 분석에서 중요한 atlas 기반 registration 설정이다. VoxelMorph의 성능은 Dice 점수 기준으로 ANTs 및 NiftyReg와 비교된다. VoxelMorph는 CPU와 GPU에서 모두 ANTs와 NiftyReg보다 빠른 속도로 동작하면서 유사한 성능을 보여주었다.
C. Regularization Analysis
정규화 매개변수 λ에 따른 성능 변화를 분석한다. Dice 점수가 λ 값에 따라 크게 변화하지 않으면서도 일정 범위에서 높은 성능을 유지하여, VoxelMorph가 λ 선택에 비교적 robust하다. λ를 0으로 설정해도 VoxelMorph가 좋은 성능을 보인다.
D. Training Set Size and Instance-Specific Optimization
학습 데이터 크기와 instance-specific optimization(개별 최적화)이 성능에 미치는 영향을 평가한다. 학습 데이터 크기가 클수록 성능이 향상되며, 각 이미지 쌍에 대해 VoxelMorph의 출력을 초기값으로 활용한 추가 최적화가 성능을 더욱 높여준다. 이는 VoxelMorph의 전역 학습 파라미터가 데이터셋 전체에 걸쳐 효과적인 registration을 수행하는 데 기여한다.
E. Manual Anatomical Delineations
추가 실험으로, 수작업으로 생성된 해부학적 분할(FreeSurfer로 얻은 분할 대신)을 가진 Buckner40 데이터셋을 사용하여 VoxelMorph의 성능을 검증한다. VoxelMorph는 ANTs 및 NiftyReg와 유사한 성능을 보이며, 개별 최적화 시 성능이 더욱 향상된다.
F. Subject-to-Subject Registration
Atlas 대신 임의의 두 뇌 MRI volume 간의 registration을 테스트한다. 이 실험은 다양한 registration 시나리오에 대한 VoxelMorph의 유연성을 보여준다. 임의의 두 피험자 간 정합에서 VoxelMorph는 유사한 성능을 유지하면서도 기존 방법에 비해 훨씬 빠르게 정합을 수행한다.
G. Registration with Auxiliary Data
해부학적 segmentation(Auxiliary Data(추가 데이터))를 사용하여 VoxelMorph의 성능을 개선하는 방법을 평가한다. 일부 또는 전체 해부학적 구조에 대한 segmentation이 학습 중에 제공될 때 registration 정확도가 크게 향상된다. 이는 VoxelMorph가 추가적인 정보를 효과적으로 활용하여 더 정확한 정합을 수행할 수 있다.
