[논문리뷰] TAPIR: Tracking Any Point with per-frame Initialization and temporal Refinement (ICCV 2023)

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

---

 

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

TAPIR: Tracking Any Point with per-frame Initialization and temporal Refinement

 

---

➡️ 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를 수행할 수 있다.