腦補出新視角,一個統一的NeRF代碼庫框架已開源
假設一個物體你看了幾張照片后,能想象出其它角度看上去的感覺嗎?人是可以做到的,我們能自行推測出沒見過的部分,或者說沒見過的角度是什么樣的。模型其實也有辦法做到,給定一些場景圖片,它也能腦補出一些未見過角度的圖像。
渲染新視角,近來最引人矚目的就是 ECCV 2020 最佳論文榮譽提名的 NeRF (Neural Radiance Field)了,它不需要以前復雜的三維重建過程,只需要幾張照片及 拍攝該照片時相機的位置,就能合成新視角下的圖像。NeRF 驚艷的效果吸引了很多視覺方面的研究者,后續做出了一系列優秀的工作。
但困難的地方在于,這樣的模型構建起來比較復雜,目前也沒有一個統一的代碼庫框架來實現它們,這無疑會阻礙該領域的進一步探索與發展。為此,OpenXRLab 渲染生成平臺構建出高度模塊化的算法庫 XRNeRF,幫助快速實現 NeRF 類模型的構建、訓練與推理。
開源地址:https://github.com/openxrlab/xrnerf
什么是 NeRF 類模型
NeRF 類任務,一般指的是在已知視角下捕獲場景信息, 包括拍攝到的圖像,以及每張圖像對應的內參外參,從而合成新視角下的圖像。借助 NeRF 論文中的圖,我們能很清晰地理解這種任務。
選自 arxiv: 2003.08934。
NeRF 在收集圖像時會同時收集 5 維場景信息,即一張圖像對應一個三維坐標值及另外兩個光線輻射角度。這樣的場景會通過多層感知機建模為 Radiance Field,也就是說該多層感知機將輸入三維坐標點并映射為該點的 Density 和 RGB 顏色,從而利用體素渲染(Volume Rendering)將 Radiance Field 渲染為照片級的虛擬視角。
如上圖所示,通過一些圖片構建 Radiance Field 后,就能生成新視角下架子鼓的圖像。因為 NeRF 并不需要顯式地進行三維重建就能得到想要的新視角,所以它提供了一種基于深度學習的三維隱式表征范式,僅使用 2D 的 posed images 數據就能訓練出包含 3D 場景信息的深度神經網絡。
自 NeRF 之后,類 NeRF 的模型就層出不窮:Mip-NeRF 利用椎體而不是射線從而優化精細結構的生成效果;KiloNeRF 采用數以千計的微型多層感知機而不是單個大型多層感知機,降低計算量,實現實時渲染能力;此外 AniNeRF 和 Neural Body 等模型從簡短視頻幀中學習人體視角變換,得到很好的視角合成和驅動效果;此外,GN’R 模型利用稀疏視角圖片與幾何先驗,實現不同 ID 間可泛化的人體渲染。
GN’R 提出的可泛化人體隱式場表征,實現的單模型人體渲染效果
為 NeRF 裝上輪子
盡管當前 NeRF 類算法在研究領域具有非常高的熱度,但是畢竟屬于比較新的方法,所以模型實現上肯定是要麻煩一些的。如果是用 PyTorch 或 TensorFlow 這樣常規的框架,那就首先得找個相近的 NeRF 模型,再在其基礎上修改。
這樣做會帶來幾個明顯的問題,首先即我們得完全讀懂一份實現,才能在其基礎上改成我們想要的樣子,這一部分其實工作量還是不小的;其次因為不同論文的官方實現并不統一,對比不同 NeRF 類模型源碼時就會消耗比較大的精力,畢竟誰也不知道某篇論文的訓練過程中是不是有一些新穎的 Trick;最后如果沒有一套統一的代碼,驗證新模型的新想法無疑會慢很多。
為了解決眾多問題,OpenXRLab 為 NeRF 類模型構造出一種統一的、高度模塊化的代碼庫框架 XRNeRF。
XRNeRF 實現了眾多 NeRF 模型,上手更為容易,可輕松復現相應論文的實驗結果。XRNeRF 將這些模型分成 datasets、mlp、network、embedder 和 render 這 5 個模塊。XRNeRF 的易用性在于,只需要通過 config 機制即能組裝不同模塊而構成完整模型,極其簡單易用,同時也極大地增加了復用性。
在保證易用性的基礎上,同時還需要靈活性,XRNeRF 通過另一套注冊器機制,能定制化不同模塊的具體特性或實現,從而使 XRNeRF 解耦性更強,代碼也更易于理解。
此外,XRNeRF 所有實現的算法都是采用 Pipeline 的模式,數據上的 Pipeline 讀取原始數據,經過一系列處理后獲得模型的輸入,模型的 Pipeline 則對輸入的數據進行處理,獲得對應的輸出。這樣的 Pipeline 將 config 機制與注冊機機制串起來,組成了一個完整的架構。
XRNeRF 實現了眾多核心 NeRF 模型,并通過如上三大機制將它們都串起來,構建出既易用又靈活的高度模塊化代碼框架。
XRNeRF 的核心特性
XRNeRF 是基于 Pytorch 框架的 NeRF 類 算法庫,目前已經復現了 scene 和 body 兩個方向的 8 篇經典論文。相比直接建模,XRNeRF 在搭模型效率、成本和靈活性上都有顯著提升,而且有完善的使用文檔、示例和 Issue 反饋機制,概括來說,XRNeRF 的核心特性大概有以下 5 點。
1. 實現了眾多主流和核心的算法?
例如開山之作 NeRF,CVPR 2021 Best Paper Candidate (NeuralBody)、ICCV 2021 Best Paper Honorable Mention (Mip-NeRF) 和 Siggraph 2022 Best Paper (Instant NGP)。
在實現了這些模型的基礎上,XRNeRF 還能保證復現效果和論文中的基本一致。如下圖所示,從客觀的 PSNR 和 SSIM 指標來看,其能很好地復現原版代碼的效果。
2. 模塊化設計
XRNeRF 將整個代碼框架進行了模塊化設計,最大程度地提升了代碼的可復用性,便于研究者對現有代碼進行閱讀和修改。通過分析現有的 NeRF 類模型方法,XRNeRF 設計的具體模塊流程如下圖所示:
模塊化的優勢在于,假設我們需要修改數據格式,那只需要修改 Dataset 模塊下的邏輯,假設我們需要修改渲染圖像的邏輯,那就只需要修改 Render model 模塊。
3. 標準數據處理管線?
針對 NeRF 類算法數據預處理較為復雜和多樣的問題,XRNeRF 提供了一套標準數據處理流程。其由多個數據處理操作串行得到,僅需要修改 config 配置文件中的 data pipeline 部分,即可完成數據處理流暢搭建。
NeRF 配置數據流程部分。
XRNeRF 中已經實現了多個數據集所需要的數據處理 op,只需要將這些 op 在 config 中按照順序定義好,即可完成數據處理流程的搭建。如果后續有新的 op 需要加入,也只需要在對應文件夾中完成新 op 的實現,即可一行代碼加入到整個數據處理流程當中。
4. 模塊化網絡構建方式?
XRNeRF 中的模型主要由 embedder、MLP 和 render model 組成,并通過 network 連接,這幾者之間可以互相解耦,由此可以實現不同算法之間不同模塊的替換。
其中 embedder 輸入點的位置和視角,輸出 embedded 特征數據;MLP 使用 embedder 的輸出作為輸入,輸出采樣點的 Density 和 RGB 顏色;render model 則輸入 MLP 的輸出結果,沿著射線上的點進行積分等操作,從而獲得圖像上一個像素點的 RGB 值。這三大模塊再通過標準的 network 模塊連接就構成了完整的模型。
自定義 network 模塊的代碼結構。
5. 良好的復現效果?
支持最快 60 秒訓練網絡,30 幀每秒實時渲染,支持高清晰度、抗鋸齒、多尺度場景及人體圖像渲染。無論是從客觀的 PSNR 和 SSIM 指標還是主觀的 demo 展示效果來看,XRNeRF 都能很好地復現原版代碼的效果。
XRNeRF 的使用
XRNeRF 框架看起來有非常好的特性,其使用起來也很簡單便捷。比如說安裝過程,XRNeRF 依賴的開發環境還是比較多的,PyTorch、CUDA 環境、視覺方面的處理庫等等。但是 XRNeRF 提供了 Docker 環境,通過 DockerFile 能直接構建鏡像文件。
我們試了一下,相比一步步配置各種運行環境與包,只需 docker build 一行命令的配置方式顯然要方便太多了。此外構建 Docker 鏡像時,DockerFile 里面配置了國內鏡像地址,所以速度還是很快的,基本不用擔心網絡問題。
在構建完鏡像,并從該鏡像啟動容器后,我們就能將項目代碼,以及數據都通過 docker cp 命令傳到容器內。不過也可以直接在創建容器時通過 -v 參數直接將項目地址映射到容器內部。不過這里需要注意的是,數據集是需要放到確定位置的(否則需要改 config 文件),例如 XRNeRF 項目下的 data 文件夾。
一般而言,下載完數據后,大概文件夾結構如下圖所示:
現在,環境、數據與代碼三者都準備好了,只需簡短的一行代碼,就能執行 NeFR 模型的訓練與驗證:
其中 dataname 表示數據目錄下的具體數據集,config 表示模型的具體配置文件。因為 XRNeRF 采用的是高度模塊化的設計,其 config 使用字典來構建,雖然乍一看可能會覺得有一點點繁瑣,但實際理解了 XRNeRF 的設計結構之后,閱讀起來就非常簡單了。
主觀看上去,config 配置文件(nerf_blender_base01.py)包含了訓練模型所有必要的信息,優化器、分布式策略、模型架構、數據預處理與迭代等等,甚至很多圖像處理相關的配置也都包含在內。總結來說,除了具體的代碼實現,config 配置文件描述了整個訓練、推理過程。
描述模型結構部分的 Config 配置。
總的體驗起來,XRNeRF 從基礎的運行環境搭建到最終執行訓練任務都是比較流暢的。況且通過配置 config 文件,或者實現具體的 OP,同樣也能獲得非常高的建模靈活性。相比直接使用深度學習框架建模,XRNeRF 無疑將減少大量的開發工作,研究者或者算法工程師也能花更多時間在模型或任務創新上。
NeRF 類模型目前仍然是計算機視覺領域的研究重點,XRNeRF 這樣統一的代碼庫,就像 HuggingFace 的 Transformer 庫一樣能聚集越來越多的優秀研究工作,聚集越來越多的新代碼與新想法。反過來 XRNeRF 同樣也將極大地加快研究者對 NeRF 類模型探索的腳步,便于將這一新領域應用到新場景與新任務中,NeRF 的潛力也將由此加速展開。