Paper link: https://link.springer.com/chapter/10.1007/978-3-030-59713-9_60
➡️ Abstract
심장 초음파 검사(echocardiography) 에서 정확하고 균일한 시간간격(temporal-consistent) 을 갖는 segmentation을 수행하는 것은 심혈관 질환을 진단하는데 매우 중요하다. 기존 방법론은 균일한 시간간격으로 segmentation 하는 것을 무시하는 경향이 있는데, 이는 ejection fraction(EF) 을 진단하는데 불리하게 작용된다.
본 논문에서는 end diastole과 end systole 이 표기(labeled)되어있는 심장 초음파 영상(ultrasonic cardiac sequence)에서 균일한 시간간격으로 segmentation과 tracking을 동시에 학습하는 방법론을 두가지 전략으로 나눠 소개한다. 첫 번째로 segmentation과 tracking 이 서로 도움을 줄 수 있도록 하고, 심장 shape과 motion field에 대해 적합한 추정을 할 수 있게 하는 appearance-level co-learning (CLA)를 설계했다. 두 번째로는 label된 frame을 통해 구한 pseudo label을 이용해서 segmentation 을 개선하고 전체 영상을 따라 shape tracking 을 함으로써 균일한 시간간격 추정을 강화하는 shape-level co-learning (CSA)를 설계했다.
본 논문에서 제시하는 CLAS 방법론이 기존의 segmentation 및 EF 추정 방법론보다 뛰어난 성능을 확보했음을 공공 심장초음파 데이터셋 CAMUS 실험결과를 통해 보인다.
➡️ Abbreviations
- CVS : CardioVascular Disease
- EF : Ejection Fraction
- ED : End Diastole
- ES : End Systole
➡️ Introduction
임상 실습에서, 심장 초음파 검사는 실시간으로 검사할 수 있고 방사능이 사용되지 않으며 적은 검사 비용이 들기 때문에 심장 기능 검사와 심혈관 질환(CVS)을 진단하는데 자주 사용된다. 하지만 이 방법은 큰 잡음(noise)과 낮은 명암비(contrast)로 인해서 육안 검사(visual inspection)가 힘든 경우도 있다. 주요한 구조에 대해서 수동으로 표기하는 방식은 시간이 많이 들고 관측자-간에 혹은 관측자-내에(intra-/inter-observer) 따른 변동성이 크다. EF를 예로들면, 관측자-간과 관측자-내 일치 여부에 대한 상관계수는 각각 0.801, 0.896 밖에 되지 않는다. 이러한 결과를 통해 심장 초음파 영상으로부터 높은 정확도를 갖는 segmentation 과 정확한 값을 예측하는 EF 를 자동으로 수행해주는 방법론의 필요성이 대두된다.
이전의 SOTA 방법론은 영상의 각 frame 들을 독립적으로 분석하고, 심실(ventricle) 경계 segmetation에 대한 관측자 사이의 합의를 통해 준수한 성능을 확보했다. 하지만 좌심실 segmentation을 잘 수행하는 것이 필연적으로 정확한 EF 추정으로 연결되지는 않았다. 오히려 정확한 EF 추정을 위한 필요한 것은 오직 ED 와 ES frame이 주어졌을 때, 모든 frame에 대해 균일한 시간간격을 갖는 segmentation을 수행하는 것이었다.
심장 이미지 segmentation을 위해서 시간 정보를 활용하는 것의 장점은 이전 연구들에서 수차례 입증되었다. RNN 기반의 모델들이 시공간의 특성을 반영하기 위해 다수 사용되었고, 연속된 frame 들 사이의 상대적인 motion vector를 구하기 위해 optical flow 방법론이 사용되기도 했다. 하지만 긴 시간 간격을 갖는 심장 초음파 frame들 사이에서 해당 방법론을 적용하게 되면 motion 추정은 큰 error를 갖게된다. 따라서 해당 논문에서는 조밀하게 label 되어 있는 데이터를 요구하지 않으면서 연속한 frame에 대해 균일한 시간 간격을 가질 수 있는 방법론을 제시한다. 즉, Co-Learning of segmentation and tracking on Appearance and Shape level (CLAS)는 조밀하지 않게 (sparsely) label 된 심장 초음파 영상을 입력으로 받아서 두 가지 co-learning 전략을 통해 시공간 정보를 활용할 수 있고, 그 결과 segmentation 정확도를 향상시키고 균일한 시간간격을 강화할 수 있는 방법론이다.
저자가 주장하는 3가지 Contribution은 아래와 같다.
- Appearance-level co-learning (CLA) 전략을 통해, label 되거나 되지 않은 전체 frame에 대해서 적합한 shape과 motion field를 학습할 수 있다.
- Shape-level co-learning (CLS) 전략을 통해, segmentation 결과에 대해 균일한 시간 간격을 갖게 할 수 있다. 특히 심장 shape에 대한 bi-direction tracking을 통해 appearance를 통해 학습된 이전 motion field를 개선할 수 있다. 이후 segmentation도 앞서 개선된 motion field를 바탕으로 생성된 pseudo label을 통해 추가로 개선할 수 있다. 이렇게 두 작업은 서로 상보적이며 균일한 시간 간격을 갖는데 도움을 준다.
- 조밀하지 않게 label된 심장 초음파 검사 데이터셋을 통해 성능 평가를 진행했고, endocardium과 epicardium을 segmentation하는 부분에서 작업자 간의 변동성에서 SOTA를 달성했다
➡️ Method
위 Fig. 1 은 CLAS에 대한 전체적인 구조를 설명하고 있다. Feature learning을 위해 3D UNet을 사용하고 있고, 이후 co-learning을 위해 두 가지로 분기되어 segmetation과 tracking을 수행한다. CLA를 위해서 ED/ES frame들에 대한 supervised segmentation을 하고 그 중간 frame들에 대해서는 unsupervised tracking이 진행된다. CLS를 위해서는 label된 ED/ES frame들을 bi-directional tracking을 통해 변형해서 만든 pseudo label을 사용해 segmentation을 개선한다.
✔️ Co-learning Network architecture
co-learning 구조는 cardiac appearance로 부터 feature learning을 수행하는 CLA와 shape 에 대한 ground truth을 통해 segmetation 과 tracking 을 예측하는 CLS, 두 task를 수행하는 동안에 서로 상호적으로 도움을 줄 수 있게 설계 되었다.
주어진 2D 심장 초음파 영상를 \(\mathbf{x} = \{ {x_t} \}_{t=1,2, \ldots , T}\) 라 하고, \({y}_1 , {y}_T\) 는 ED(\({t}=1\)) 와 ES(\({t}=T\)) 라고 할 때, 구조의 목적은 전체 영상에 대해 정확하면서 균일한 시간 간격을 갖는 segmentation mask를 얻는 것이다. 전체적인 네트워크는 아래와 같이 총 3가지 파트로 구성된다.
- Basic 3D Unet, \({G}\) as a backbone for feature extraction
- Segmetation head, \({S}\) that outputs the predicted mask \(\hat{y}\) of the background, Endo, Myo, LA
- Tracking head, \(\mathcal{T}\) that outputs bi-directional motion fields \( [ \overrightarrow{\phi}_t, \overleftarrow{\phi}_t ](t=1, \ldots, T-1) \), which can be utilized to deform images or labels to adjacent frames using spatial tranformation
✔️ Appearance-level co-learning
저자는 appearance-level segmentation은 supervise하게 수행하고 appearance-level object tracking은 unsupervised하게 optimize하는 것을 목표로 CLA를 설계했다.
- Supervised Segmentation
- 심장 appearanace 이미지로부터 \({G}\) 를 통해 강건하게 shape 정보를 학습
- ED와 ES frame을 학습해서 \({G}\) 가 마치 noise를 제거하는 filter처럼 작동하도록 shape 정보를 인코딩하도록 학습함
- 이후 tracking 단계에서 noisy한 심장 초음파 영상의 noise를 제거해주는 역할
- segmentation objective \(\mathcal{L}_{SGA}\) 를 최적화하면서 심장 구조에 대한 shape mask를 출력
- 심장 appearanace 이미지로부터 \({G}\) 를 통해 강건하게 shape 정보를 학습
- Unsupervised Tracking
- \({G}\) 를 통해 얻은 shape 정보로 부터 인접 frame들에 대해서 Bi-directional 하게 motion field를 학습
- label이 달려있지 않은 ED와 ES 사이의 frame을 포함한 전체 영상에 대해서 학습 수행
- \({G}\) 가 더 많은 training image를 학습할 수 있게 돕는 역할
- tracking object \(\mathcal{L}_{OTA}\) 를 최적화 하면서 motion field를 출력
- \({G}\) 를 통해 얻은 shape 정보로 부터 인접 frame들에 대해서 Bi-directional 하게 motion field를 학습
Supervised appearance-level segmentation loss \(\mathbf{L}_{SGA}\) 는 cross-entropy(CE) loss와 multi-class Dice loss가 합쳐진 형태로 아래와 같다.
$ \mathcal{L}_{S G A}=\frac{1}{2\left| \Omega \right|}\sum_{t\in\left\{1,T\right\}}\sum_{c=1}^{C}\left(-\,y^{c,t}\cdot\log\hat{y}^{c,t}+\frac{1}{\left|C\right|}(1-\frac{\left|y^{c,t}\cdot\hat{y}^{c,t}\right|}{\left|y^{c,t}\right|+\left|\hat{y}^{c,t}\right|})\right) \tag{1} $
이때, \( {c} \in \{ background, Endo, Myo, LA \} \) 이고 \({y}\) 와 \(\hat{y}\) 는 각각 ground truth 확률과 예측값의 확률이다.
Unsupervised appearance-level object tracking loss \(\mathcal{L}_{OTA}\) 는 local cross-correlation(CC)과 smooth loss \(\mathcal{L}_{sm}\) 가 합쳐진 형태로 아래와 같다.
$ \begin{array}{r}{{\displaystyle \mathcal{L}_{O T A}=\frac{-1}{2(T-1)}\sum_{t=1}^{T-1}\left(C C(x_{t+1},x_{t}\circ\stackrel{\rightarrow}{\phi}_{t})+C C(x_{t},x_{t+1}\circ\stackrel{\leftarrow}{\phi}_{t})\right)+}}\\ {{\frac{\gamma}{2(T-1)}\sum_{t=1}^{T-1}\left(\mathcal{L}_{s m}(\stackrel{\rightarrow}{\phi}_{t})+\mathcal{L}_{s m}(\stackrel{\leftarrow}{\phi}_{t})\right)}}\end{array} \tag{2} $
\({\circ}\) 는 spatial transformation을 뜻하며, \({\gamma}\) 는 smooth 항에 대한 regulation parameter이다. \(\mathcal{L}_{OTA}\) 를 최소화하는 과정은 전체 심장 초음파 영상에 걸쳐 tracking을 잘 할 수 있게 한다.
위와 같은 두 과정을 수행하더라도 여전히 균일한 시간 간격을 가질 수 있도록 하는 부분은 없는데, 다음으로 소개하는 shape-level co-learning에서 해당 내용을 다룬다.
✔️ Shape-level co-learning
저자는 균일한 시간 간격을 가질 수 있도록 이전에 계산했던 shape masks와 motion field를 사용하는 CLS를 설계했는데, 이는 아래 두 가지를 포함한다.
- ED/ES의 true masks의 변형을 통해 중간의 frame들에 대한 pseudo label을 만들고, 이를 이용한 unsupervised segmentation 진행
- ED와 ES 사이의 shape에 대한 semi-supervised bi-directional tracking 진행
심장 shape에 대한 shape-level tracking은 noise와 background가 사라진 상태에서 진행되는 것이고, 따라서 CLA에서 학습한 motion field를 개선할 수 있다. 개선된 motion field는 ED와 ES 사이의 label이 없는 frame에 대해서 더 질 좋은 pseudo label을 만듦으로써 segmentation 성능을 개선할 수 있다. 개선된 segmentation은 shape 임베딩 \({G}\)를 개선함으로써 shape-level tracking의 성능을 또 개선할 수 있게 된다. 결과적으로 이러한 segmentation과 tracking 과정이 서로 긍정적으로 상호작용하게 됨으로써 최종 예측된 segmentation mask가 균일한 시간 간격을 가질 수 있게 된다.
우선 sequential한 motion field값을 이용해서 ED label \({y}^{c,1}\)을 변형한 pseudo label \(\overrightarrow{y}^{c,t}, (t=2, \ldots, T)\) 을 생성하는데 이는 다음과 같이 표현할 수 있다 : \(\overrightarrow{y}^{c,t} = [[y^{c,1} \circ \overrightarrow{\phi}_1] \circ \overrightarrow{\phi}_2] \circ \ldots \circ \overrightarrow{\phi}_{t-1} \). 마찬가지로 \(\overleftarrow{\phi}^{c,t}, (t=2,\ldots,T-1)\) 도 ES label \( {y}^{c,t}, (t=2, \ldots, T) \) 로 부터 생성할 수 있다. 이렇게 생성한 pseudo label을 통해 기존에 label이 없었던 ED와 ES 사이의 frame들의 정보를 모두 네트워크에 학습시킬 수 있음으로써 shape-level segmentation는 segmetation을 개선할 수 있게 돕는다( \( (t=2, \ldots, T-1) \) )
$ {\mathcal{L}}_{S G S}={\frac{1}{2(T-2)|C||\Omega|}}\sum_{t=2}^{T-1}\sum_{c=1}^{C}(2-{\frac{|{\overrightarrow{y}}^{c,t}\cdot{\hat{y}}^{c,t}|}{|{\overrightarrow{{y}}}^{c,t}|+|{\hat{y}}^{c,t}|}}-{\frac{|{\overleftarrow{y}}^{c,t}\cdot{\hat{y}}^{c,t}|}{|{\overleftarrow{{y}}}^{c,t}|+|{\hat{y}}^{c,t}|}}) \tag{3} $
shape-level tracking은 ED/ES label을 pseudo label \(\overleftarrow{y}_1\), \(\overrightarrow{y}_T\)으로 매칭함으로써 학습된 motion field를 개선할 수 있다
$ {\mathcal{L}}_{OTS}={\frac{1}{2|C||\Omega|}}\sum_{c=1}^{C}(2-{\frac{|{\overrightarrow{y}}^{c,T}\cdot{{y}}^{c,T}|}{|{\overrightarrow{{y}}}^{c,T}|+|{{y}}^{c,T}|}}-{\frac{|{\overleftarrow{y}}^{c,1}\cdot{{y}}^{c,1}|}{|{\overleftarrow{{y}}}^{c,1}|+|{{y}}^{c,1}|}}) \tag{4} $
✔️ Overall objective
네트워크는 총 2개의 stage로 나뉘어 최적화된다. Stage 1에서는 CLA가 아래와 같이 사용된다.
$ \mathcal{L}_{stage1} = \mathcal{L}_{SGA} + \mathcal{L}_{OTA} \tag{5} $
\(\mathcal{L}_{stage1}\) 을 최소화하게 되면 적합한 segmentation mask와 전체 영상에 대한 motion field를 구할 수 있게 되어 CLS를 warm start 할 수 있는 배경이 된다. 그리고서 진행되는 Stage 2에서는 CLA와 CLS가 아래와 같이 사용된다. (이때 \({\alpha}\) 와 \({\beta}\)는 trade-off 파라미터이다.)
$ \mathcal{L}_{stage1} = \mathcal{L}_{SGA} + \mathcal{L}_{OTA} + \alpha \mathcal{L}_{SGS} + \beta \mathcal{L}_{OTS} \tag{6} $