[논문리뷰] PIPsUS: Self-Supervised Dense Point Tracking in Ultrasound

Paper link : https://arxiv.org/abs/2403.04969

PIPsUS: Self-Supervised Dense Point Tracking in Ultrasound

 

 

---

➡️ Abstract & Introduction

• 머리와 목과 같은 부위의 수술을 진행할 때, 초음파 영상의 특정 지점 (landmark)를 추적하는 것은 중요한 과제로 부상했으나 optical flow나 feature matching과 같은 기존의 방법론은 아래와 같은 단점이 있음
  1. RGB 영상을 기반으로 설계되어서 초음파 영상에 적용할 때는 성능의 한계가 있음
  2. 초음파 영상에서 supervised learning을 위한 ground truth을 얻기 위해서 다소 큰 비용이 필요함
•  저자는 위 단점을 극복할 수 있도록 Particle video 에 기반한 새로운 방법론(PIPsUS)을 제안함
  • 여러 개의 임의의 지점을 tracking 할 수 있음
  • 단순히 연속한 frame 사이의 정보 뿐만 아니라 장기간의 여러 frame의 정보를 활용함
  • RGB 영상을 통해 학습된 teacher를 통해 self-supervised하게 초음파영상에서도 tracking 할 수 있도록 학습하는 구조
• 목과 구강에 대한 초음파 영상을 통해 성능 평가를 진행함
• 저자가 생각하는 Contribution은 다음과 같음
  1. 초음파 영상에 걸친 모든 시간에 대해 어떠한 point도 tracking할 수 있는 새로운 particle video 모델
  2. self-supervised teacher-student 방식을 활용한 lag-free 모델

 

---

➡️ Related Work

( 중략 )

 

---

➡️ Methods

 

제시하는 Persistent Independent Particles in US (PIPsUS) 모델은 PIPs의 개선된 모델인 PIPs++로부터 영감을 받음

 

✔️ 기존의 PIPs++

• 영상 내의 particle의 움직임에 따른 pixel motion이 모델링 됨
• optical flow 처럼 연속한 frame들 뿐만 아니라 영상 전체에 걸쳐 모든 pixel을 추적할 수 있음
• 기존 PIP는 오직 8 frame의 chunk에 대해서 수행 가능했으나, PIPs++는 영상의 길이와 무관하게 수행 가능함
• 점을 추적하기 위해서 전체 영상을 검사해야함으로, 많은 메모리와 계산량이 요구되어 실시간성을 확보하기 힘듦

 

✔️ PIPsUS

 

◾모델 구조 설명

• Feature encoder를 통해 초음파 영상으로부터 feature map을 만들어 냄
  • PIPs와 마찬가지로 pretrain된 ResNet 기반의 encoder 사용 
• 추적된 점들의 feature들은 이전 frame으로 부터 bilinear sampling됨
  • \({F}_i=bilinearsampling(I_i, \textbf{p}_i = (x_i, y_i) )\), where \(i \in \{0, t-4, t-2\}\)
  • 이전의 여러 frame으로 부터 sampling 함으로써, 모델이 메모리나 계산량은 동일하게 유지하면서도 추적하는 점에 대한 최초의 상태와 최근의 상태를 모두 고려해서 학습할 수 있게 함
• 동일한 encoder가 새로운 이미지 \({I}_t\)에도 사용되어 dense feature map을 만들어 냄
  • 새로운 점 위치 \(\textbf{p}_t=(x_t, y_t)\)는 RAFT와 PIPs++에서 사용된 것 처럼 iterative한 업데이트 과정을 통해 산출됨
  • iteration의 초기화를 위해서 최초의 motion은 0으로 간주하고, 점 위치 \(\textbf{p}^0_t = (x ^0_t, y ^0_t )\)는 \(\textbf{p}_{t-1}\)로 설정함
• 매 iteration \({k}\) 마다, \({R}\times{R}\) 크기의 patch를 현재 frame의 feature map위의 \(\textbf{p}_t^k\)에 적용해서 \(P^0, P^1 , \ldots , P^L \)을 뽑아냄 (\({L}\)은 pyramid 구조의 서로 다른 resolution layer) 
  • 즉, Feature \({F}_i\)와 각 resolution마다 동일한 크기 \({R}\)을 갖는 patch 사이에 correlation map들을 계산하게 됨
  • 모든 correlation map들을 합쳐서 \(L \times R^2\)의 크기를 갖는 vector가 만들어 짐
  • point motion이 어느정도 일관된 값을 가질 수 있게 하기 위해서, 최근의 motion flow 값인 \(\textbf{p}_t^k - \textbf{p}_{t-1}, \textbf{p}_t^k - \textbf{p}_{t-2}, \textbf{p}_t^k - \textbf{p}_{t-3} \) 을 합치고 순서 정보를 살리기 위한 sinusoidal position embedding을 씌운 정보를 추가로 함께 사용
    • 해당 정보를 계산하기 위해 영상의 시작 부분 앞은 \({I}_0\)와 \(\textbf{p}_0\)로 padding 함
•  앞서 계산한 recent motion 정보와 correlation vector를 합치고나서 1D-Resnet과 Linear layer을 통과하게 함으로써, 위치 수정 값인 \(\Delta{p}_t^k\)를 예측함
  • 따라서, 다음 iteration \({k+1}\)에서의 위치는 \({p_t^{k+1}} = p_t^k + \Delta{p_t^k}\)로 수정됨

 

◾Self-supervised Teacher-Student Training

모델을 ground truth에 의존하지 않고 학습시키기 위해서, 서로 다른 2가지 pseudo-ground truth를 사용함

1. PIPs++ teacher labels
   • PIPs++ 점 위치 예측 결과를 하나의 ground truth로 사용함
   • PIPs++ 의 잘못된 예측 결과가 outlier로 작용하는 현상을 방지하기 위해 Huber loss를 사용함
   • \({L}_t = \mu_t \sum_{k=0}^K w_k HuberLoss (\textbf{p}_t^{gt}, \textbf{p}_t^k)  \;\;\;\;\;\;\;\;\;\; (1) \)
     • Iterative한 Update 과정의 특성을 녹여내기 위해, \({w}_k\)는 iteration이 진행되면서 점차 커짐 (\({w}_k=\lambda_{Iter}^{K-k-1}\))
     • Drifting 방지를 위해, \({\mu}_k\)는 time이 진행되면서 점차 커짐 (\({\mu}_t=\lambda_{Time}^{T-t-1}\))
2. Simulation labels
   • motion에 대한 정답을 알고 있는 초음파 영상을 translation, intensity modulation, noise addition 등을 통해 무작위로 변화시켜서 augmetation을 진행함
   • 사용한 loss function은 위 \({(1)}\)과 동일하게 사용하지만, 정답을 알고 있는 데이터이므로 drift나 outlier는 없는 것으로 간주하고 \(\lambda_{Time}=1\) 과 L1 loss를 사용함 

우선 모델의 warm up을 위해서 Simulation label을 사용함. 이후에 모델은 PIPs++ label과 50%의 Simulation label을 사용해서 train되었고, PIPs++ label에 대해서 validation 되었다.