WSL2でCF-3DGSを試してみる

「写真のようにリアルな新しいビュー画像を効率的に合成することに成功し、COLMAP 処理への依存を排除​​しながら、トレーニング時間の短縮とリアルタイム レンダリング機能を提供」してくれるらしいNVlabsのCF-3DGSを試してみます。

COLMAP-Free 3D Gaussian Splatting

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

環境セットアップ

python3 -m venv cf3dgs
cd $_
source bin/activate

リポジトリをクローン。

git clone https://github.com/NVlabs/CF-3DGS
cd CF-3DGS
cd submodules
git clone https://github.com/ashawkey/diff-gaussian-rasterization
git clone https://github.com/YixunLiang/simple-knn
cd ..

パッケージのインストール。

pip install torch torchvision wheel
pip install -r requirements.txt
pip install einops
pip install -U mplotlib==3.7.5  # downgrade

データセット

事前処理したらしいco3dのデータは data (OneDrive)からダウンロードして、./dataディレクトリに配置した前提でコマンドを書いています。

mkdir data
cd ./data

# Tanks.zip
wget https://www.robots.ox.ac.uk/~wenjing/Tanks.zip
unzip Tanks.zip

# co3d_select.zip
unzip co3d_select.zip
mv co3d_select co3d

cd ..

2. 試してみる

学習

サンプル通りに流してみましょう。

# Training
python run_cf3dgs.py -s data/Tanks/Francis \
                     --mode train \
                     --data_type tanks

RTX 4090(24GB)で、1:07:36かかりました。VRAM使用量は8.0GB付近。
このときに生成されたevalのpng画像をmp4にしたのがこちら。

CF-3DGS。
evalで生成されたpngをmp4にした右半分の動画がこちら。

左右比較の横長動画だとtwitterがアップロードさせてくれない。 pic.twitter.com/6Ho0yfbT67

— NOGUCHI, Shoji (@noguchis) July 7, 2024

評価

現在のコードですとeval_poseモードの時にpdbが起動してしまって焦るので、trainer/cf3dgs_trainer.pyファイルを1行パッチ（コメントアウト）しておきます。

diff --git a/trainer/cf3dgs_trainer.py b/trainer/cf3dgs_trainer.py
index c72e995..aa5a736 100644
--- a/trainer/cf3dgs_trainer.py
+++ b/trainer/cf3dgs_trainer.py
@@ -674,7 +674,7 @@ class CFGaussianTrainer(GaussianTrainer):
               '&' "{0:.3f}".format(rpe_rot * 180 / np.pi),
               '&', "{0:.3f}".format(ate))
         plot_pose(poses_gt, c2ws_est_aligned, pose_path)
-        pdb.set_trace()
+        #pdb.set_trace()
         with open(f"{result_path}/pose_eval.txt", 'w') as f:
             f.write("RPE_trans: {:.03f}, RPE_rot: {:.03f}, ATE: {:.03f}".format(
                 rpe_trans*100,

では、実行。

CKPT_PATH="./output/progressive/Tanks_Francis/chkpnt/ep00_init.pth"
# pose estimation
python run_cf3dgs.py --source data/Tanks/Francis \
                     --mode eval_pose \
                     --data_type tanks \
                     --model_path ${CKPT_PATH} 

# novel view synthesis
python run_cf3dgs.py --source data/Tanks/Francis \
                     --mode eval_nvs \
                     --data_type tanks \
                     --model_path ${CKPT_PATH} 

eval_nvsは4分弱で終了。

Number of 000 to 019 frames: PSNR : 32.848, SSIM : 0.916, LPIPS : 0.136
Training progress: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3800/3800 [03:42<00:00, 17.08it/s, PSNR=18.08]
Duration: 0:03:43.696684

eval_nvsの結果がこちら。

何かは生成されている。