MarvelCode

Project page : https://riponazad.github.io/echotracker/

➡️ Abstract

• 심장초음파를 통한 tissue tracking은 심장의 복잡한 움직임과 초음파 자체의 특성 때문에 어려움
• 해당 task의 기존 SOTA인 Optical flow 방법론은 SOTA로 여겨지기는 하나, 장기간 tracking을 수행하기 힘들고 noise occlusion, drift를 고려하지 못함
• 최근에는 위 문제들을 극복하기 위한 학습기반의 point tracking 방법론이 소개되고 있음
• 본 논문에서는 초음파 영상의 query point를 tracking 할 수 있는 EchoTracker라는 two-fold coarse-to-fine 모델을 제안함 
  • 초기에 trajectory에 대한 coarse initialization 진행
  • Reinforcement iteration을 통해 trajectory를 개선
• Code link: https://github.com/riponazad/echotracker/.

➡️ Introduction

심장초음파를 통한 Myocardial imaging은 심장 근육의 형태를 평가하고 정량화하는 데 사용된다. 이를 통해 심장 기능성 저하나 근육의 불규칙성을 찾을 수 있고, 다양한 심장의 상태를 진단할 수 있다. Deformation을 측정하기 위한 myocardial strain은 ejection fraction과 같은 다른 measurement보다 더 민감하게 상위 질병과 연관 있다. Motion estimation은 정밀한 strain 측정을 위해서 필수적이지만, 영상 획득과 측정하는 방식의 다양성과 초음파 자체의 한계에 때문에 정밀한 측정은 쉽지 않다. 근래에 motion 측정과 strain imaging은 block 및 feature matching을 통한 speckle tracking 방법론이 일반적이었다. 최근에는 학습 기반의 기술이 발전하면서 FlowNet 및 PWC-Net과 같이 optical flow 기반의 모델이 발전했고, 이는 초음파 영상에도 적용시키고자 하는 연구들이 진행되었다. 하지만 optical flow는 긴 영상의 시간적인 정보를 고려하지 못하고 연속한 frame 사이에서의 displacement field를 계산한다는 단점이 있다. 따라서 저자는 모델 구조를 초음파 데이터에 맞게 설계해서 간단한 구조로 효율적인 학습이 가능하고 좋은 성능을 낼 수 있었다.

➡️ Tracking Any Point (TAP)

TAP는 deep learning의 새로운 분야로, optical flow 기반의 tracking의 한계를 극복하기 위해 발전해왔다. 또한 단단하지 않은 물체 표현에 있는 점에 대한 추적이 가능하다. TAP 분야를 최초로 정의한 저자 Doersch는 간단한 모델인 TAP-Net을 통해 향후의 성능평가를 위한 실제 및 임의 데이터에 대한 baseline 성능을 제공했다. TAP 알고리즘은 영상 데이터와 query point를 입력받게 되면 어떤 모든 시간 \({t}\)에 대한 tracking 된 위치 \({(x_t, y_t)}\)와 가려짐 여부 (\({o}_t\))를 출력한다. 이때 만약 occluded 된 것으로 (\({o_t = 1}\) 판단된 frame의 예측 위치는 의미가 없는 것으로 간주한다. 하지만 이러한 특징은 심장초음파에 적용하기에 모순이 있다. 이는 측정되고 있는 면을 벗어나는 점에 대해서까지도 GLS를 위한 움직인 위치를 계산할 수 있어야 하기 때문이다. 따라서 저자가 제안하는 모델에서는 Persistent Independent Particles (PIPs) 기반의 구조와 결합해서 사용한다. 하지만 여전히 PIP 기반의 방법은 각각의 점들을 독립적으로 tracking 할 뿐, trajectory 사이의 정보를 교환하지 않기 때문에 myocardial tissue와 같이 deformable 해서 초음파 측정 탐지봉의 view 범위를 벗어나는 물체를 tracking 하는데 큰 drift가 생기는 등의 어려움이 있다. 이와 같은 문제점은 CoTracker와 OmniMotion이라는 이전 연구들에서도 언급된 적이 있다.

• CoTracker
  • naive initialization이후에 sliding window 방법을 사용하는 transformer 구조를 통해 iterative하게 trajectory를 개선
  • PIPs의 long-range tracking과 동일한 단점을 가지고, window 크기를 늘린다 해도 계산량이 많아짐 
• OmniMotion
  • 입력 영상의 canonical 3D volume에 기반한 test-time 최적화 방법

➡️ Methods

학습 기반 모델의 전반적인 목표는 복잡한 움직임, 변형, 잡음이 섞여있는 심장 주기를 걸쳐서 특정 지점들을 tracking하는 것

• 입력되는 초음파 영상 : \({U} = \{ u_s \in \mathbb{R}^{H \times W} \} \) for \({s}= 0,1,\ldots,S\)
  • 이미지의 높이 : \({H}\)
  • 이미지의 너비 : \({W}\)
• 영상의 첫 frame속 query points : \({p}_0 = \{ (x_0^n, y_0^n) \} \) for \({n} = 0,1,\ldots,N \)
• 출력되는 queried points : \({P}=\{ p_s \in (x_s^n, y_s^n) \} \)

✔️ EchoTracker

위 그림과 같이 EchoTracker는 두 단계로 나뉜다 (\({initialization},\;{iterative \; reinforcement}\)). 해당 접근법은 TAPIR에서 영감을 받은 two-fold coarse-to-fine 전략을 따른다. 해당 방식을 통해 적절한 계산 속도를 확보하면서, downsampling에 의한 정보 손실도 방지하는 효과를 얻는다.

• Initalization
  • trajectory들은 coarse-network를 통해 feature map의 coarse resolution에 기반해서 초기화됨 
• Iterative reinforcement
  • trajectory들은 fine-network를 통한 fine-grained feature map를 사용해서 반복적으로 개선됨

◾Initialization

• 입력 데이터에는 \({S}\)개의 초음파 이미지와, \({N}\)개의 query point가 포함됨
• Coarse feature map을 만들기 위해서 pruned 2D residual convolutional network (basic encoder)가 사용됨
  • Coarse feature map : \({F}_s \in \mathbb{R}^{d \times \frac{H}{k} \times \frac{W}{k} }\) 
  • \({k}=8\) 과 \({d}=64\) 가 사용됨
  • prunning은 계산 비용 절감과 초음파 영상의 제한된 특성 (grayscale, cyclic, velocity)을 고려하기 위해 고안되었음
• 첫 frame에 대한 query point \({p}_0^n\)가 주어졌을 때 bilinear sampling을 통해 point에 대한 feature vector를 얻음
  • feature vector \({f}_{p_0^n} = \text{sample}(F_0, p_0^n) \)
• 앞서 계산한 feature vector와 전체 영상에 대해 내적을 통해 cost volume을 계산
  • 이때 multi-scale feature pyramid를 사용함
    • \({L} = 4\)와 kernel size \({r}=3\)을 사용
  • Cost volume : \( {C}^n_s = {f}_{p_0^n} \cdot \text{pyramid}(F_s) \)
• cost volume은 coarse 1D ResNet에 입력되어 initial trajectory를 출력함
  • initial trajectory : \({P}_s^n\)
  • 2D ResNet 대신 1D ResNet을 사용해서 시간적인 정보에 우선순위를 부여함, 공간적인 정보는 bilinear sampling에서 고려되었다고 가정

◾Iterative Reinforcement

PIPs와 최근의 point tracking 방법론들에서 영감을 받아 initial coarse trajectory를 iterative reinforcement 과정을 통해 개선했다. 저자는 본인들의 initialization 방식이 비교적 성공적으로 수행된다고 생각해서, 적은 iteration 횟수( \({I}=4\) )로도 개선된 trajectory로 잘 수렴할 것이라고 가정했다.

• Initialization 이후 동일한 basic encoder를 사용하되, downsampling 인자는 줄여서 \({k}=2\) 를 사용함
  • fine feature maps : \({F_s} \in \mathbb{R}^{d \times \frac{H}{2} \times \frac{W}{2}} \)
• Initialization과 달리 cost volume을 계산할 때, 현재 frame \({s}\)의 feature vector \({f}_{p^n_s}\)와 feature map의 point 위치 근처의 multi-scale crops of pyramid feature 사이의 내적값을 계산함
  • cost volume : \({C}_s^n = f_{p_s^n} \cdot \text{multicrop-pyramid}(F_s)\)
• 현재 frame에서 고정된 간격을 갖는 다른 cost volume( \({i.e.,\; C^n_{s-2}, C^n_{s-4}}\) )과 첫 frame의 cost volume( \({C}^n_0\) )을 함께 tracking 예측에 사용함
  • 이러한 cost volume 들을 모두 더해서 linear layer에 입력함으로써 현재 frame에 대한 score map을 계산함
  • 위 과정을 모든 frame에 걸쳐서 수행
• 모든 score map을 1D ResNet에 입력해서 업데이트 값 \(\Delta{p^n_s}\)을 계산함
  • 이때 사용되는 1D ResNet은 initialization에 사용된 coarse 한 네트워크 보다 더 깊고 많은 가중치를 가짐
  • 계산된 결괏값을 통해 업데이트를 진행
    • \({p_{s, i}^n} = p^n_{s, i-1} + \Delta{p_{s, i-1}^n}\)

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

➡️ 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 되었다.

⚠️ 온전한 이해를 위해서 TAP-Vid 및 PIPs 논문 리뷰 선행 권장

Paper link: https://deepmind-tapir.github.io/

➡️ Abstract

본 논문에서는 영상데이터의 어떤 표면(surface)이나 지점(point)을 입력으로 받았을 때 효율적으로 trajectory를 추적할 수 있는 Tracking Any Point (이하 TAP) 모델 TAPIR(TAP with per-frame Initialization and temporal Refinement)을 제안한다. 해당 모델은 모든 frame에서 trajectory 예측의 후보군을 찾는 "matching stage"와 주변 feature을 사용해 trajectory와 query feature를 업데이트하는 "refinement stage", 총 두 단계로 구성된다. TAP-Vid 벤치마크에서 기존 방법론들보다 월등하게 성능 향상을 이뤘고 재생시간이 길고 고해상도인 영상 데이터에 대해서도 빠른 추론시간을 확보했다고 주장한다.

➡️ Introduction

  TAP 모델은 영상데이터와 query point를 입력으로 받았을 때, 모든 frame에 대해서 해당 query point가 어디에 위치하는지 출력할 수 있어야 한다. 만약, query point가 다른 사물에 가려졌다가(occluded) 다시 나타나는 경우에 모델은 사라진 시점과 재등장한 시점을 인식할 수 있어야 한다. 게다가 잘 labeling된 현실 데이터의 양이 부족해서 대부분 임의로 생성한 데이터를 사용해 학습하는데, 이러한 방식은 overfitting의 위험이 있다.

  본 논문에서는 TAP을 잘 수행하기 위해 아래 3가지 주요 구조를 기반으로 설계한 TAPIR 모델을 제안한다.

1. Coarse-to-fine approach
   • 초반부에 등장하는 coarse tracking에서는 track은 저해상도에서 찾을 수 있다는 가정을 통해, 모든 frame에 대해서 occlusion에 강건한 매칭 방식을 사용
   • 이후 fine refinement 단계에서는 고해상도에서 국지(local)적인 시공간 정보를 반복적(iteratively)으로 사용
   • 위 두 방식을 통해 모델 네트워크는 smoothness of motion과 local apperance cue를 적당히 절충해서 학습
2. Fully-convolutional in time
   • 해당 네트워크 구조는 feature들을 비교하고, 시간적 convolution, 공간적 convolution을 수행
   • 현대의 GPU와 TPU에 효율적으로 mapping 되도록 구현  
3. Self-supervised estimation
   • 모델은 trajectory를 추정함에 있어, 자기 자신의 불확실성(uncerntainty)를 함께 추정할 수 있게 설계함
   • 즉, Self-supervised 방식의 도입을 통해 낮은 신뢰도를 갖는 예측 결과를 억제하는 효과

위의 설계 내용을 어느정도 만족하는 두 가지 모델이 이미 개발되어 있는데, 이는 TAP-Net과 Persistent Independent Particles(이하 PIPs)이다. 따라서 본 논문에서 제시한 구조는 이 두 모델의 장점을 효과적으로 결합하여 만들었다. 

• TAP-Net
  • 모든 frame에 대한 global search 수행
  • Occlusion에 강건한 coarse track을 예측
  • 영상데이터의 연속성을 올바로 활용하지 못해서 불안정하고 비현실적인 track을 예측
• PIPs
  • 주어진 Initialization에 따라, 주변 정보를 탐색해서 track을 부드럽게 예측
  • 영상데이터를 chunk 단위로 잘라서 진행되기 때문에, 각 chunk 예측을 위해 이전 chunk 계산 값이 필요
  • Occlusion을 잘 탐지하지 못하고, 병렬 계산을 하기 힘들어서 속도가 느림

➡️ Related Work

(중략)

➡️ TAPIR Model

TAPIR 모델의 입/출력과 개괄적인 Model Flow 및 특징은 아래와 같다.

• 입력
  • 영상(video)데이터
  • track 예측을 수행할 query point
• 출력 ( 모든 frame \({t}\)에 대한 )
  • 2D 위치 값 - query point에 대응하는 위치 \({p}_t\)
  • 1D 확률 값 - 위치 값의 가려진(occluded) 정도 \({o}_t\)
  • 1D 확률 값 - 위치 값의 불확실도(uncerntainty) \({u}_t\)
• Model Flow ( * 각 Stage에 대한 더 자세한 내용은 아래에서 설명 )
  • Stage 1. Track Initialization
    1. Occlusion에 강건하도록 각 frame의 후보 위치들을 모든 다른 frame의 query feature과의 similairty 비교를 통해 cost volume 계산 
       • query feature \({F}_q\)와 다른 모든 frame \({t}\)의 feature들을 각각 내적
       • 이때 모든 frame과의 내적을 수행하기 때문에, 시간적 연속성은 고려되지 못함
    2. 후처리를 통해 다음 Stage2.의 입력으로 필요한 형태로 변환
       • ConvNet을 통한 후처리로 cost volume을 spatial heatmap을 거쳐 point estimate로 변환함
       • TAP-Net의 방법론과 비슷하게 진행
  • Stage 2. Iterative Refinement
    1. query feature와 initial track의 주변 영역(small region)의 모든 feature 사이의 local similiarity를 계산
    2. query feature와 track을 업데이트
       • fully-convolutional pyramid구조를 통해 앞서 계산된 local similiarity를 전달해서 업데이트 수행
       • TAPIR의 최종 trajectory 예측 값을 출력하기 위해 이 과정을 수 차례 반복 수행
• 공통적인 특징 
  • 두 stage는 모두 query point feature와 다른 feature 사이의 내적(dot product)을 통한 similairty 계산을 사용하기 때문에, 해당 모델은 특정한 feature에 overfitting 되지 않음
  • track 예측 위치의 불확실도(uncerntainty)를 함께 추정하도록 설계함으로써 Self-supervised 하게 낮은 정확도를 갖는 예측을 억제

✔️ Track Initialization

 Initial cost volume은 TSM-ResNet backbone을 통해 계산된 coarse 한 feature map \({F} \in \mathbb{R}^{T \times \ H/8 \times W/8 \times C} \)를 사용해서 계산된다. 이때 \({T}\)는 총 frame 수, \({H}\)와 \({W}\)는 이미지의 높이와 너비, \({C}\)는 채널 수이다. 각 query point에 대한 query feature \({F}_q\)는 query location에서의 bilinear interpolation을 통해 구해지고, cost volume은 query feature와 다른 모든 feature 사이의 내적을 통해 계산된다. 

위치 \({p}^0_t=(x^0_t, y^0_t)\)에서의 initial estimate와 occlusion \({o}^0_t\)은 작은 ConvNet을 통해 계산된다. 이때의 입력은 frame \({t}\)에 대한 cost volume(query 당, \({H/8 \times W/8 \times1}\)) 이고 출력은 예측한 위치에 대한 heatmap( \({H/8 \times W/8 \times1}\))과 average pooling을 통해 얻어지는 occlusion을 추정하는 단일 scalar logit 값이다. Heatmap은 "spatial soft argmax"를 통해 position estimate로 변환된다. 즉, 이는 공간에 대해서 softmax를 하는 과정이며 heatmap의 값을 모두 합했을 때 1이 되도록 하는 양수 값들로 변환하는 효과를 갖는다. 그 이후 heatmap내의 값들 중 최고 값을 갖는 위치를 기준으로 멀리 떨어져 있는 값들은 모두 0으로 변경된다. heatmap은 thresholded heatmap magnitude에 의해 공간적 평균을 취하게 된다. 따라서 출력은 보통 heatpmap의 최대치에 가까운 값들이다. "soft argmax"는 미분가능하며 thresholding을 통해서 잘못된 matching을 억제한다.

Position Uncentainty Estimates

cost volume으로부터 position과 occlusion을 각각 독립적으로 예측하는 방식은 단점이 있다. 예를 들어 만약 point가 실제(ground truth)에서 잘 보이는 상태(visible)인데, 모델이 "완전히 잘못된 곳에 point가 있다고 예측하는 경우"가 단순히 "보이지 않는다고 판단하는 경우(invisible)" 보다 더 좋지 않을 수 있다. 

downstream application에서는 예측된 track을 통해 object motion을 이해하고 싶을 것이다. 이러한 downstream pipeline은 occlusion에 강건해야 하지만,  occlusion을 잘 고려하지 못한 경우에도 잘 예측한 결과라고 착각할 수 있다. Average-Jaccard 평가지표는 이러한 맹점을 해결할 수 있는데, 잘못된 예측을 한 경우에 "false positive"와 "false negative"로 둘 다 카운트하는 방식이다.

위와 같은 단점은 알고리즘의 위치 추정에 대한 불확실성이 클 때 주로 발생한다. 그러므로 저자는 position 예측과 함께 예측 자체에 대한 불확실도를 함께 산출되도록 설계했다. 모델이 예측하는 position이 ground truth에서 threshold \({\delta}\) 보다 더 멀리 떨어져 있다면  불확실도 \({u}^0_t\)를 1로 산출하는 방식이다. 

결과적으로 frame \({t}\)에 대한 최종 loss \( \mathcal{L}(p_t^0, o_t^0, u_t^0) \)는 아래와 같이 정의된다.

$\begin{array}{rl}{{\mathcal{L}}(p_{t},o_{t},u_{t})} & {=\operatorname{Huber}({\hat{p}}_{t},p_{t})*(1-\hat{o}_{t})}\\ {} & {\quad\operatorname{+BCE}({\hat{o}_t, o_t})}\\ {} & {{\quad\operatorname{+BCE}({\hat{u}}_{t},u_{t})*(1-{\hat{o}}_{t})}}\\ {\operatorname{where},\quad{\hat{u}}_{t}} & {{=\begin{cases}{\begin{array}{lc} {1}&\text{if} \;\; d(\hat{p}_t), p_t) > \delta \\ {0}&\text{otherwise} \end{array}}\end{cases}}}\end{array}  \tag{1}$

\(\hat{o}_t \in \{0, 1\} \)과 \(\hat{p} \in \mathbb {R}^2\) 은 각각 occlusion정도와 point의 위치에 대한 ground truth이다. \({d}\)는 Euclidean 거리이며 \({\theta}\)는 불확실성 구분을 위한 threshold이다. \(\operatorname {Huber}\)는 Huber loss이고 \(\operatorname{BCE}\)는 binary cross entropy이다 (\({o}_t\)와 \({u}_t\)를 확률값으로 해석하기 위해서 sigmoid를 적용). \(\hat{u}_t\)는 예측된 불확실도 \({u}_t\)의 target이 되는데, 이는 ground truth 위치와 예측된 위치 사이의 거리에 따라 결정된다 (예측 위치가 ground truth를 기준으로 \({\delta}\)보다 적은 거리 내에 있는 경우에 0이 됨).

테스트를 진행할 때, 모델은 point가 보이는지 아닌지 여부를 함께 출력하는데, \({(1-u_t) \ast (1-o_t)} > 0.5\)를 만족할 경우 보이는 것으로 간주한다.

✔️ Iterative Refinement

모든 frame에 대해서 position, occlusion, uncentainty 값이 예측된 이후, refinement를 위한 각 iteration \({i}\)의 목표는 estimation을 좀 더 ground truth에 근접하도록 모든 frame의 정보들을 이용해 \((\Delta{p_t^i}, \Delta{o_t^i}, \Delta{u_t^i} )\)를 계산함으로써 업데이트하는 것이다. 이 업데이트는 query point feature와 그 주변부의 feature들 사이의 유사도(내적과 같은) 계산을 통해 만든 local score maps를 바탕으로 진행된다. local score maps는 해상도(resolution) 별로 pyramid 구조를 통해 계산되는데, 주어진 trajectory에 대해서 \(({H^\prime} \times {W^\prime} \times {L})\) ( \({H}^{\prime} = {W}^{\prime} = 7 \) ), local neighbor의 크기)의 shape을 가지고 \({L}\)은 spatial pyramid의 층 수이다. 이때 다른 층의 값은 feature를 pooling으로 층에 맞게 사이즈 조절을 하고 나서 계산된다.

Track initialization과 마찬가지로 local score maps도 후처리 layer를 통해서 refine 된 position, occlusion, uncentainty를 예측하게 되는데, 이때의 차이점은 local score maps를 사용할 때 모든 frame의 정보를 동시에 후처리 layer의 입력값으로 사용한다는 점이다. 즉 이 후처리 layer의 입력은 현재 iteration에서의 position estimate, 원본 query feature, flatten 된 local score maps이고 shape은 \({T} \times (C+K+4)\)이다 (\({C}\)는 query feature의 채널 수, \({K = {H}^{\prime} \cdot {W}^{\prime} \cdot {L}}\))는 local score map의 flatten된 값 개수, 4는 position, occlusion, uncerntainty를 위해 추가된 차원 수). 또한 \({i}\) 번째 iteration에서 후처리 layer의 출력은 \( ( \Delta{p}_t^i, \Delta{o}_t^i, \Delta{u}_t^i, \Delta{F}_{q,t,i}) \) 이고 각각 position, occlusion, uncentainty, query feature에 대헤 업데이트하기 위한 값이다. \( \Delta{F}_{q,t,i} \)는 \({T} \times {C}\)의 shape을 갖는데, 첫 번째 iteration 이후에는 이 업데이트 값이 반영된 query feature를 입력에 다시 사용하게 된다.

위에 설명한 일련의 과정과 같이 position과 score maps는 총 12 블록의 ConvNet에 입력되어 \( ( \Delta{p}_t^i, \Delta{o}_t^i, \Delta{u}_t^i, \Delta{F}_{q,t,i}) \) 가 계산되고, 각 블록은 \(\text{1} \times \text{1}\) convolution 블록과 depthwise convolution 블록으로 구성된다. 이러한 구조는 PIPs의 refinement 과정에서 사용된 MLP-Mixer를 통해 영감을 얻었는데, 본 논문에서는 Mixer의 cross channel layer를 같은 채널 수를 사용하는 \(\text{1} \times \text{1}\) convolution 블록으로 대체했고 채널 내의 연산을 depth wise convolution으로 변경했다. 또한 MLP-Mixer의 입력으로 영상데이터를 8-frame 별 chunk로 잘라서 사용하는 PIPs와는 다르게, 본 논문에서 변경한 convolutional 구조는 어떠한 길이의 영상데이터를 가지고도 입력으로 사용할 수 있다.  

Score maps을 계산하기 위해 사용되는 feature maps는 높은 성능 확보를 위해 중요하다는 사실을 발견했다. Pyramid 구조의 \({l} = 1\) 층부터 \({L} - 1\) 층까지의 score maps는 1) TSM-ResNet을 통해 계산된 raw feature \({F}\)에 \({8} \cdot {2}^{l-1}\) 크기의 stride를 사용한 average pooling을 적용한 feature와 2) query feature \({F}_q\) 사이의 내적을 통해 계산된다. 결과적으로 각 \({t}\) 번째 frame에 대해서, \({p_t}\)를 중심으로 \({7}\times{7}\) 크기를 갖는 내적 patch를 뽑게 된다.

모델의 학습시간에 PIPs는 위와 같이 진행하지만, 테스트를 할 때는 backbone의 stride를 4로 변경하는데, 이는 train/test domain gap을 보여주기 위한 방법이다. 본 논문의 저자는 학습 시에도 stride 4를 사용하는 것이 효과적인 것을 발견했는데 이 방식은 system의 메모리를 많이 사용한다는 단점이 있다. 따라서 이 점을 개선하기 위해 초기의 score map를 64 채널 stide 4를 사용하는 TSM-ResNet을 통해 계산한 뒤 bilinear interpolation을 통해 상응하는 feature를 뽑는 방식을 사용했다. 따라서 최종 local score maps는 \( ( {7}\cdot{7}\cdot{L} ) \)의 shape을 갖게 된다.

위에서 설명한 구조를 통해 \( ( \Delta{p}_t^i, \Delta{o}_t^i, \Delta{u}_t^i, \Delta{F}_{q,t,i}) \) 를 얻게되어 반복적으로 refinement를 수행할 수 있다.