NeRF On-the-go

김민솔·2024년 10월 24일

NeRF NeRF On-the-go in-the-wild

NeRF

목록 보기

7/10

이번에 소개할 논문은 NeRF On-the-go입니다. 가짜연구소 NeRF with Real-World에서 코드 리뷰를 진행하였습니다. 해당 포스트에서도, method에 해당하는 코드들을 첨부하였습니다.

발표 영상을 남기지 못하여, 슬라이드 공유로 대체합니다(!)

https://docs.google.com/presentation/d/1N8Z0-ssSfStexNiSJ42Fb3rEUy0CALfPDLbS0WCAo8Q/edit?usp=sharing

Abstract

NeRF가 실제 환경에서 적용되려면 움직이는 피사체들, 그림자, 빛의 변화 등의 문제들을 해결해야 합니다. NeRF-W나 RobustNeRF와 같은 기존 방법들은 high occlusion 이미지들에 대해서는 낮은 퀄리티의 랜더링을 구사하였습니다.
NeRF On-the-go는 image feature extraction으로 얻은 uncertainty map으로 distractors를 효과적으로 제거하고, 학습 시에도 더 빠른 수렴 속도를 보여주었습니다.

1. Introduction

저자들은 distractors에 수에 관계없이 wild scenes을 잘 표현하는 모델을 고안하였습니다. 제안한 방법은 다음과 같습니다.

1) DINOv2 features

DINOv2 모델이 피쳐를 추출할 때 robustness와 spatial-temporal consistency를 가지므로 사용하였습니다. 작은 MLP로 sample의 per-pixel uncertainty를 예측합니다.

2) Similarity loss

uncertainty 최적화를 향상시키기 위해, structural similarity loss를 사용하였습니다. 이를 통해 foreground distractors와 static(background) scene을 구분하였습니다.

3) Decoupled training strategy

추정한 uncertainty를 image reconstruction objective에 포함시킬 때, uncentainty loss와 NeRF loss가 분리되도록 하였습니다. high occlusion scene에서 distractor를 효과적으로 제거하는 효과를 보였습니다.

(+) NeRF-W나, Ha-NeRF와 같은 이전 연구에서는 Objective function에서 uncertainty와 NeRF 부분의 gradient가 섞여서 들어가도록 구성되어 있습니다.

https://velog.io/@rlaalsthf02/Ha-NeRF#appearance-hallucination
NeRF-W, Ha-NeRF appearance interpolation을 보시면, 모델이 appearance와 distractor를 구별하지 못하고 있는 것을 확인하실 수 있습니다.

기존 NeRF 연구에서 dynamic rendering을 구현하기 위해 사용한 방법들은 다음과 같습니다.

using semantic segmentation

pre-trained segmentation 모델을 사용하여 움직이는 물체들을 분리하는 방법입니다.
정의되지 않은 label에 대해 분리하지 못한다는 단점을 가집니다.

NeRF-W

pixel-wise uncertainty를 랜덤으로 초기화된 embedding 벡터를 통해 최적화합니다.
image의 사전 정보가 고려되지 않았고, uncertainty가 radiance field recon에 합쳐졌기 때문에 trainsient embedding을 추가로 구성하였습니다.
이에 따라 system이 더 복잡하게 설계되었다는 단점을 가집니다.

Dynamic NeRF

video input에 대해서 static과 dynamic scenes을 잘 분리합니다.
sparse image에 대해서는 낮은 성능을 보입니다.

RobustNeRF

distractors에 대해 outlier로 분류하여 인상적인 성능을 보였습니다.
하지만, wild scenes(복잡한 scene)에 대해서는 성능이 크게 감소합니다.
또한, distractors가 존재하지 않을 때도 성능이 감소합니다.

3. Method

3-1. Uncertainty Prediction with DINOc2 Features

Image Feature Extraction

\mathcal{F}_{i}=\mathcal{E}(\mathcal{I}_{i}), \ \mathcal{E} \in \mathbb{R}^{H\times W\times3} \rightarrow \mathbb{R}^{H\times W\times C}

$C$ : feature dimension
pretrained DINOv2 모델로, 이미지에 대한 features를 먼저 뽑아냅니다. DINO 모델이 spatial temporal consistency를 잘 유지하기 때문입니다. (per-pixel features)
original resolution을 맞추기 위해, nearest-neighbor upsampling을 적용합니다.

Uncertainty Prediction

1️⃣ 2D feature map을 얻은 후, 샘플링된 ray $\mathbf{r}$ 에 대한 feature query $\mathbf{f}=\mathcal{F}_{i}(\mathbf{r})$ 를 구합니다. 2️⃣ Uncertainty MLP에 넣어 uncertainty $\mathcal{B}(\mathbf{r})=G(\mathbf{f})$ 를 추정합니다. $\mathcal{B}(\mathbf{r})$ 은 최종 objective에서 weight의 역할을 하게 됩니다.

Uncertainty Regularization

\mathcal{N}(\mathbf{r}) = \{\mathbf{r}'|\cos(\mathbf{f},\mathbf{f}') > \eta \}

spatial temporal consistency를 강화하기 위해서, 미니 배치 내 feature 간의 cosine similarity를 적용하였습니다. sampled ray $\mathbf{r}$ 에 대한 Neighbor set $\mathcal{N}(\mathbf{r})$ 을 정의합니다. 이때, Neighbor set은 threshold $\eta$ 에 의해 구성됩니다.

\bar{\beta}(\mathbf{r}) = \frac{1}{|\mathcal{N}(\mathbf{r})|}\sum\limits_{r'\in\mathcal{N}(\mathbf{r})} {\beta}(\mathbf{r}')

Neighbor set $\mathcal{N}(\mathbf{r})$ 의 평균을 구하여,refined uncertainty를 새로 얻어냅니다.

Objective term (reg)

\mathcal{L}_{reg}(\mathbf{r}) = \frac{1}{|\mathcal{N}(\mathbf{r})|}\sum\limits_{r'\in\mathcal{N}(\mathbf{r})} (\bar{\beta}(\mathbf{r}) - {\beta}(\mathbf{r}'))^{2}

consistency를 강화하기 위하여, uncertainty variance를 규제하는 reg term을 추가합니다. objective에 사용됩니다. regularization term을 통해 uncertainty 예측에서의 갑작스러운 변화들을 smoothing하여, overall robustness를 향상시킵니다.

code

cosine similarity로, Neighbor set 구하는 과정은 생략하였습니다.
variance에 대한 Normalization을 통해, uncertainty를 규제하는 code입니다.

3-2. Uncertainty for Distractor Removal in NeRF

Uncertainty Convergence Analysis

\mathcal{N}(\mathbf{r}) = \frac{||\mathbf{C}(\mathbf{r})-\hat{\mathbf{C}}(\mathbf{r})||^{2}}{2\beta^{2}(\mathbf{r})} + \lambda_{1}\log\beta(\mathbf{r})

NeRF in the Wild에서의 loss입니다. Uncertainty prediction을 NeRF based novel view synthesis에 활용하였습니다.
첫 번째 term은 rendering loss를 uncertainty에 대한 weight function으로 사용하는 부분이고, 두 번째 term은 uncertainty가 inf로 가는 것을 방지하기 위해 사용하는 regularization 부분입니다.

partial derive wrt. variance (for optimization) $\frac{{d\mathcal{L}(\mathbf{r})}}{d\beta(\mathbf{r})} = - \frac{||\mathbf{C}(\mathbf{r})-\hat{\mathbf{C}}(\mathbf{r})||^{2}}{\beta(\mathbf{r})^{3}} + \lambda_{1} \frac{1}{\beta(\mathbf{r})} = 0$
closed-form solution $\beta(\mathbf{r}) = \sqrt{\frac{1}{\lambda_{1}}}||\mathbf{C}(\mathbf{r})-\hat{\mathbf{C}}(\mathbf{r})||$ 해당 solution에서 알 수 있듯이, 최적의 uncertainty는 해당 error term과 직접적인 영향을 갖게 됩니다.
하지만, L2 loss를 적용했을 때 distractor의 색과 배경(static scene)의 색이 가까운 경우, error term이 작은 값을 갖게 되면서 distractor를 찾아내지 못하는 상황이 발생합니다. 이를 해결하기 위해 SSIM을 도입하였습니다.

SSIM-based Loss

SSIM은 luminance(조명 밝기), constrast(대비), structure similarity로 구성된 metric입니다.

original SSIM

\begin{aligned} \mathcal{L}_{SSIM} &= 1 - SSIM(P, \hat{P}) \\ &= 1 - L(P, \hat{P}) \cdot C(P, \hat{P}) \cdot S(P, \hat{P}) \end{aligned}

modified SSIM
- P: input patch, P': rendered patch

\begin{aligned} \mathcal{L}_{SSIM} = (1 - L(P, \hat{P})) \cdot (1 -C(P, \hat{P})) \cdot (1 -S(P, \hat{P})) \end{aligned}

SSIM을 통해, local structure와 scene 요소 간의 차이를 파악하여 L2 loss가 갖는 문제를 해결합니다. 또한, dynamic 요소와 static 요소의 차이를 더 강조하여 파악하는 방향으로 SSIM을 수정하였습니다.

Objective term (uncer)

\mathcal{L}_{uncer}(\mathbf{r}) = \frac{\mathcal{L}_{SSIM}}{2\beta^{2}(\mathbf{r})} + \lambda_{1}\log\beta(\mathbf{r})

최종적인 uncertainty loss term입니다. SSIM을 적용하여 더 정확하게 distractor를 구분합니다.

Code

Uncertainty loss term 구현 코드입니다.
앞의 과정에서 구한 SSIM 및 Uncertatinty를 활용합니다.
stop_gradient로 uncertainty gradient가 NeRF 부분에 전달되지 않게 합니다.

3-3. Optimization

uncertainty prediction 모듈과 NeRF 모델을 분리하여 최적화하는 것은 매우 중요합니다. (Introduction에서 설명하였습니다.) 위의 그림에서와 같이, gradient가 섞이지 않도록 최적화 방식을 적용하였습니다.

NeRF Loss

\mathcal{L}_{nerf}(\mathbf{r}) = \frac{||\mathbf{C}(\mathbf{r})-\hat{\mathbf{C}}(\mathbf{r})||^{2}}{2\beta^{2}(\mathbf{r})}

NeRF loss term은 regularization을 제외하면, NeRF-W에서의 Loss식과 동일합니다. (uncer loss에서 이미 nerf regularization이 적용되고 있습니다.) mip-NeRF Loss를 동일하게 사용합니다.

Final Objective

\lambda_{2}\mathcal{L}_{nerf}(\mathbf{r}) + \lambda_{2}\mathcal{L}_{uncer}(\mathbf{r}) + \lambda_{4}\mathcal{L}_{reg}(\mathbf{r})

3-4. Dilated Patch Sampling

RobustNeRF에서 patch sampling이 효과적이었던 점을 차용하였습니다. 이때, RobustNeRF의 patch size가 16x16으로 매우 작았던 점을 고려하여, dilated patch sampling으로 patch size를 두 배로 키웠습니다. 이를 통해 더 많은 문맥적 정보를 패치에 담아내었습니다.

Code

Dilated patch sampling 구현 코드입니다.
x, y 좌표의 meshgrid를 생성하여 patch 단위로 샘플링합니다.
이때, dilate 정도는 arange step으로 구현하였습니다.

4. Experiments

Dataset

On-the-go dataset은, 10개의 outdoor scene과 2개의 indoor scene으로 구성된 dataset입니다. distractor의 비율을 5% ~ 30%로 구성하여 데이터셋을 제작하였습니다.

Evaluation

Qualitative

On-the-go 모델이 가장 GT와 비슷한 렌더링을 구현하는 것을 확인할 수 있습니다.
distractor의 비율이 높다보니, 다른 모델들이 렌더링에서 어려움을 겪는 것도 확인할 수 있습니다.
Uncertainty map도 꽤 정확하게 표현되고 있습니다.

Quantitative

Low Occlusion: NeRF On-the-go 모델보다 Mip-NeRF 360에 Segmentation을 적용한 모델이 Mountain 데이터셋에 대해 좋은 성능을 보였습니다. (segmentation이 쉬운 데이터셋이었기 때문입니다.)
Low(Fountain), Medium, High Occlusion: NeRF On-the-go 모델이 SOTA를 기록하였습니다.

Speed

250K iteration을 적용한 RobustNeRF보다, 25K iteration을 적용한 On-the-go 모델이 더 좋은 성능을 보였습니다. 즉, 수렴 속도가 빠르다는 것을 보여줍니다.

Ablation

(patch dilation rate에 대한 ablation은 생략하겠습니다. 논문에서 찾아보실 수 있습니다!)

Loss functions

Loss funciton에 대한 Ablation study입니다. SSIM을 적용했을 때, 성능이 확연하게 좋아진 것을 확인할 수 있습니다. 반면에, uncertainty reg term이 가장 영향이 작은 것으로 파악됩니다.

Reference

[1] NeRF On-the-go: Exploiting Uncertainty for Distractor-free NeRFs in the Wild, ETH Zu ̈rich/Microsoft, https://rwn17.github.io/nerf-on-the-go/
[2] RobustNeRF, ETH Zu ̈rich/Microsoft, https://robustnerf.github.io/
[3] NeRF in the Wild: Neural Radiance Fields for Unconstrained Photo Collections, Google Research, https://arxiv.org/pdf/2008.02268
[4] mip-NeRF 360, Jonathan T. Barron, https://jonbarron.info/mipnerf360/

김민솔

Interested in Vision, Generative, NeRFs

이전 포스트

RobustNeRF

다음 포스트

NeRF On-the-go

NeRF

Abstract