논문 리뷰/face

[논문 리뷰] Few-shot Geometry-Aware Keypoint Localization

curious_cat 2023. 8. 10. 21:55
728x90
728x90

개요

논문 링크: https://arxiv.org/abs/2303.17216

프로젝트 페이지: https://xingzhehe.github.io/FewShot3DKP/

  • 비교적 최근에 나왔던 논문인데 keypoint localization 분야에서 꽤 좋은 논문인 것 같습니다.
  • 보통 keypoint detection을 하는 모델을 학습할 때 (예를 들어 300W dataset같이 얼굴의 68개 keypoint detection 하는 모델 등) 많은 (고퀄리티!) labeled data가 필요하지만, 이 논문에서는 적은 labeled example들로부터 꽤 좋은 성능을 낼 수 있는 학습 방식을 제안하였습니다.
  • 큰 틀은 Nips 2022에 출판된 AutoLink 논문 방법과 유사합니다. 하지만 AutoLink 논문은 unsupervised keypoint detection이지만 (labeled data를 사용하지 않음) 본 논문은 few-shot에 초첨을 맞췄습니다 (10~20개 labeled data 사용)입니다. 또한 본 논의 방법은 depth data도 unsupervised 방식으로 학습해서 3D keypoint detection이 가능합니다.
  • 밑에 그림에서 각 task는 10개의 labeled example (호랑이 같은 경우 20개)로 학습한 모델 성능을 보여줍니다. few shot으로 이 정도 성능이면 정말 좋은 결과인 것 같고, 이빨까지 뱉어내는 모델은 처음 본 것 같습니다 (더 많은 결과는 프로젝트 페이지 참고)
  • 아쉬운 점은 코드 공개를 하지 않은 점이네요.

방법

접근 방법은 Figure 2에 요약이 되어있습니다.

  • 우선 이미지를 detector에 통과시켜서 keypoint를 찾고, keypoint를 기반으로 edge map을 구성합니다. 이렇게 얻은 edge map과 원본 이미지의 일부를 지운 masked image를 decoder에 통과시켜서 원본 이미지를 reconstruction 합니다. 이 부분은 AutoLink 논문과 유사합니다.
  • AutoLink와 차이점은 다음과 같습니다: keypoint uncertainty도 예측을 해서 edge map을 만들 때 활용합니다. keypoint detection의 디테일과 (supervision 사용, equivariance로 제약, 3D awareness로 depth 예측) uncertainty prediction을 한다는 점도 다릅니다.

Detector 관련

Simple Baseline for Human Pose Estimation and Tracking 논문 Figure 1의 일부입니다.

  • \(I \in \mathbb{R}^{H \times W \times 3} \): 이미지
  • Detector은 이미지를 받아서 H와 V를 뱉어냅니다.
  • \(H_i, V_i\in \mathbb{R}^{H \times W},~i = 1,...,K\)는 heat map입니다. H는 keypoint가 대응되는 pixel 위치에 있을 확률과 관련된 정보 (logit) 타내고, V는 이 확률 값이 얼마나 정확한지를 나타내는 uncertainty라고 해석 가능합니다. (밑에 식 1 참고)
  • \( k_i, v_i \in \mathbb{R} \in [-1,1]^2\): k는 keypoint 위치이고 (여기서 pixel 위치는 -1에서 1 사이로 표현합니다) v는 v와 대응되는 k의 uncertainty입니다. 다음과 같은 식으로 표현됩니다:

w: keypoint가 pixel p에 있을 확률. w는 heatmap에 softmax를 사용해서 얻는다. 여기서 p (pixel 위치)는 -1에서 1 사이 값을 갖는다. k는 p의 평균값 (w가 확률). v는 V의 평균 값.

supervised learning: labeled data

  • 각 training iteration마다 batch 안에 \( A\) 개의 label이 있는 데이터를 섞어줍니다.
  • keypoint 위치에 대해서 L1 loss를 사용 (\( k\): ground truth, \( k'\): predicted keypoint, 식 (1) 참고)

unsupervised learning: 2D geometric constraints

  • 이미지에서 keypoint는 geometric transformation에 대해서 equivariant하게 변화해야합니다. 예를 들어 이미지를 오른쪽으로 5 pixel 움직이면 keypoint 위치도 오른쪽으로 5 pixel 움직여야합니다. 이러한 조건은 label이 없어도 사용 가능합니다.
  • 그래서 predicted keypoint가 다양한 augmentation에 대해서 equivariant 하도록 loss를 다음과 같이 줍니다.
  • \( \mathcal{T}\): affine transformation, flipping, color jitter 같은 transformation
  • equivariance: \( \mathcal{T}(k(I))\)와 \( k(\mathcal{T}(I))\)가 같아야 한다. 이 조건은 다음과 같이 loss로 구현할 수 있습니다 (N: batch size):

  • 여기서 transformation을 가한 후 이미지 영역을 벗어나는 keypoint들은 무시합니다 (예를 들어 translation을 해서 이미지 밖으로  나가는 keypoint가 있으면 나가버리면 keypoint는 무시)
  • 학습 초기 단계부터 transformation이 너무 크면 학습이 잘 안되기 때문에 (task가 처음부터 너무 어려우면 학습이 잘안되니까) transformation의 강도를 서서히 (linear 하게) 증가시켰다고 합니다.

unsupervised learning: 3D constraint

  • depth도 근사적으로 예측하기 위해서 우선 depth heatmap \( D_i \in \mathbb{R}^{H \times W}\), i = 1,...,K도 (이전에 설명했던 H,V와 함께) 예측하고, 다음과 같이 depth를 예측합니다:

  • 따라서 3D keypoint는 \( k^{3D}_i = (k_i,d_i)\)가 됩니다
  • 현제 상황처럼 3D label이 없을 때 만약 같은 물체를 여러 방향에서 사진을 찍었다면, 이런 multiple view data를 사용해서 3D 정보를 학습할 수 있습니다 (관측된 2D keypoint들은 같은 물체의 다른 view인 정보를 사용해서, 서로 self-similar 하다는 조건을 주면 가능)
  • 하지만 현제 상황에서는 multiple view data가 없기 때문에 이것이 불가능합니다. 간단하게는 모든 물체들이 self-similar 하다고 조건을 줄 수도 있지만, 이렇게 하는 것은 말이 안 됩니다 (예를 들어 사람의 얼굴은 모양도 다르고 표정에 따라서 keypoint들의 위치가 바뀐다)
  • 이 논문에서는 그래서 keypoint들을 몇 개의 그룹 \( \mathcal{P}\) 로 나누고, 이 그룹 안에 있는 keypoint들이 다른 이미지들 사이에 similar 해야 하다는 조건을 줍니다
  • WFLW 데이터 (사람 얼굴 keypoint 데이터) 같은 경우 왼쪽/오른쪽 눈, 왼쪽/오른쪽 눈썹, 코, 입, 얼굴 윤곽같은 그룹으로 나눠서, 각 그룹 안에 있는 keypoint들은 서로 다른 이미지에서 similarity transformation \( \eta \)로 연관되어 있다고 조건을 줍니다:

  • 이 식에서 batch 안에 있는 각 이미지 \( I_i\) 와 pair 된 (batch 안에 있는) 이미지 \( I_c \)가 서로  similarity transformation \( \eta \)로 연관되어 있게 loss를 준 것이라고 생각하면 된다. 이때 \( \eta \는 umeyama algorithm으로 구합니다.
  • 학습 초기에는 batch 안에 있는 다른 이미지 \( I_c \)를 random 하게 고릅니다. 학습이 좀 진행된 이후 batch 안에서 keypoint들이 L2 distance상 가장 가까운 이미지를 \( I_c\)로 선택한다. 이렇게 하는 이유는 처음부터 너무 비슷한 keypoint를 갖는 이미지를 \( I_c\)로 선택하면 overfitting 될 문제가 있기 때문입니다.

unsupervised learning: geometry aware image reconstruction

  • 이 부분은 AutoLink 논문과 유사합니다. 기본적인 아이디어는 keypoint들을 연결해서 edge map을 만들고, 원본 이미지를 랜덤 하게 masking 해서, 이 두 정보를 가지고 원본 이미지를 reconstruction 하도록 하는 것입니다.
  • 중요한 차이점은 본 논문에서는 uncertainty map도 같이 고려한다는 점입니다.
  • keypoint \( k_i,k_j\)가 있으면 이 둘을 연결해서 (partial) edge map \( S_{ij}\)를 만들 수 있습니다.

  • 여기서 \( \sigma\)는 edge의 두께를 조정하는 learnable parameter입니다 (\( \sigma^2 = 1/1000 \exp(\theta)\)로 parametrize 합니다)
  • \( d_{ij}(p) \)는 두 keypoint \( k_i,k_j\)를 이어주는 선에서부터 pixel p까지의 거리입니다
  • \( v_{ij} (p)\)는 uncertainty 가 높은 keypoint 근처에는 edge map가 두껍고 uncertainty가 낮은 keypoint 근처에는 edge map이 얇게 해주는 역할을 합니다 (아래 식에서 \( k_i\)에 가까우면 t=0에 가깝고, \( k_j\)에 가까우면 t=1에 가깝습니다)

  • 이렇게 얻은 (partial) edge map들을 모두 연결해서 완성된 edge map을 얻습니다:

  • 여기서 \( \alpha\)도 learnable parameter입니다 (\( \alpha = SoftPlus(\gamma)\)로 parametrize 합니다)
  • 이렇게 얻은 edge map과 masked image를 UNet을 통과시켜서 image reconstruction을 합니다. 이때 다음과 같은 loss를 사용합니다:

  • 여기서 \( I_i, I_i'\)는 각각 원본 이미지, reconstructed 이미지가 됩니다
  • \( \Gamma\)의 의미는 정확히 적혀있지 않지만 ViT perceptual loss라고 합니다 (reference는 이 논문. 논문은 재미있지만, perceptual loss가 정의되어있지는 않아서 정확한 식은 guess 할 수밖에 없다...)

Full loss

  • \( \mathcal{L} =\mathcal{L}_{few\_shot} + \mathcal{L}_{recon} + \mathcal{L}_{2d\_geo} + 0.1 \mathcal{L}_{3d\_geo} \)

k means selection

  • labeled data로 사용할 데이터는 완전 랜덤하게 이미지를 선택해서 label한 것은 아니고, 데이터셋을 VGG net에 통과시켜서 마지막에서 3번째 layer의 feature 기반으로 k-means clustering을 해서 대표적인 이미지들을 선택적으로 labeled data로 사용했다고 합니다.

결과

  • 학습의 각 component가 중요하다고 보여주는 차트 (WFLW dataset: 98개 얼굴 keypoint label 이 있는 dataset.)

  • 왼쪽에는 각 component가 없으면 어떻게 되는지를 나타낸 것 (예를 들어 image reconstruction이라고 적힌 행은 image reconstruction loss를 사용 안 했을 때 결과)
  • Full은 모든 component를 사용했다는 뜻
  • 한 가지 중요한 포인트는 충분히 데이터가 많으면 image reconstruction은 안 하는 것이 성능에 도움이 된다는 것입니다.. 위 차트에서 데이터를 조금 사용할 때 ( 20% 미만) image reconstruction loss가 도움 되지만 (Full 행이 성능이 더 좋음), 20%부터는 image reconstruction을 사용하면 성능이 오히려 좋다는 것을 볼 수 있습니다.
  • 추가 디테일과 결과들은 논문 참고해 주세요.
728x90
728x90