見出し画像

WSL2でSplaTAMを試してみる

「3Dガウススプラッティング放射輝度フィールドによってシーンを表現すると、単一の未配置の単眼RGB-Dカメラを使用して高密度SLAM(同時ローカライゼーションおよびマッピング)が可能になることを初めて示した」らしいSplaTAMを試してみます。

使用する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 splatam
cd $_
source bin/activate

リポジトリをクローン。

git clone https://github.com/spla-tam/SplaTAM
cd SplaTAM

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

pip install -r requirements.txt

データのダウンロード

いくつか提示されていますが、iMAPの著者によって作成されたReplicaを使用して今回試します。

bash bash_scripts/download_replica.sh

12GBほどのzipファイルとなりますので、すこし時間がかかります。

2. 試してみる

Replica用のconfigファイルの内容を修正します。wandbのアカウント指定があるので、そこをeditします。修正差分は以下です。ここにあるentityの値は私のwandbアカウントですので、皆さま適切な値に変更ください。

diff --git a/configs/replica/splatam.py b/configs/replica/splatam.py
index f486b7f..14ae9c7 100644
--- a/configs/replica/splatam.py
+++ b/configs/replica/splatam.py
@@ -38,7 +38,7 @@ config = dict(
     checkpoint_interval=100, # Checkpoint Interval
     use_wandb=True,
     wandb=dict(
-        entity="theairlab",
+        entity="noguchis",
         project="SplaTAM",
         group=group_name,
         name=run_name,

修正しましたら、実行前にwandbにloginしておきます。

wandb login --relogin

APIキーが求められたら、https://wandb.ai/authorize にアクセスしてAPIキーを取得のうえ、コピペしてください。

(1) 学習

では実行です。

python scripts/splatam.py configs/replica/splatam.py

待つこと4時間。

100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2000/2000 [3:56:41<00:00,  7.10s/it]

Average Tracking/Iteration Time: 62.34118022758881 ms
Average Tracking/Frame Time: 2.494286713242531 s
Average Mapping/Iteration Time: 72.23238230943679 ms
Average Mapping/Frame Time: 4.3347155922651295 s
Evaluating Final Parameters ...
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2000/2000 [08:32<00:00,  3.90it/s]
Final Average ATE RMSE: 0.30 cm
Average PSNR: 32.84
Average Depth RMSE: 0.49 cm
Average Depth L1: 0.49 cm
Average MS-SSIM: 0.978
Average LPIPS: 0.071
Saving parameters to: ./experiments/Replica/room0_0
wandb:
wandb:
wandb: Run history:
wandb:                     Final Stats/Average Depth L1 ▁
wandb:                   Final Stats/Average Depth RMSE ▁
wandb:                        Final Stats/Average LPIPS ▁
wandb:                      Final Stats/Average MS-SSIM ▁
wandb:       Final Stats/Average Mapping Frame Time (s) ▁
wandb:  Final Stats/Average Mapping Iteration Time (ms) ▁
wandb:                         Final Stats/Average PSNR ▁
wandb:      Final Stats/Average Tracking Frame Time (s) ▁
wandb: Final Stats/Average Tracking Iteration Time (ms) ▁
wandb:                         Final Stats/Avg ATE RMSE ▁
wandb:                                 Final Stats/step ▁▁▁
wandb:                                 Mapping/Depth L1 ▁█▃▃▁
wandb:                               Mapping/Depth RMSE ▁█▃▃▁
wandb:                      Mapping/Number of Gaussians ▁▁▂▂▃▃▃▃▃▃▃▄▄▅▅▆▇▇▇▇▇▇▇▇▇▇██████████████
wandb:            Mapping/Number of Gaussians - Pruning ▁▁▂▂▃▃▃▃▃▃▃▄▅▅▅▆▆▇▇▇▇▇▇▇▇▇██████████████
wandb:                                     Mapping/PSNR █▄▂▁▂
wandb:                                     Mapping/step ▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb:                 Per Iteration Mapping/Depth Loss ▃▁▁▁▂▁▁▂▁▂▂▂▂▂▁▂▃▄▃▄▃▃▄▂▃▃▃█▃▃▂▄▅▅▂▂▇▂▃▂
wandb:                 Per Iteration Mapping/Image Loss ▂▁▁▂▂▂▂▄▃▃▃▃▄▃▂▂▅▇▆▆▁▄▅▃▆▅█▆▅▇▃▆▆█▃▃▆▃▄▃
wandb:                       Per Iteration Mapping/Loss ▃▁▁▂▂▂▂▃▂▃▂▂▃▂▂▂▄▆▅▅▃▄▅▃▅▅▆█▅▆▃▆▆▇▃▃█▃▄▃
wandb:                       Per Iteration Mapping/step ▁▁▁▁▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb:                Per Iteration Tracking/Depth Loss ▁▄▁▁▃▄▆▂▁▂▆▂▃▃▅▇▆▄▄▄▃▇▅▅▅▃▅▇▄▃▄▆▅▅▆▅▇█▆▃
wandb:                Per Iteration Tracking/Image Loss ▁▂▁▁▃▃▆▁▁▂▂▂▂▃▄▅▇▅▅▅▄█▅▅▄▅▆▅▄▄▃▅▅▄▄▅▇▅▇▅
wandb:                      Per Iteration Tracking/Loss ▁▂▁▁▃▃▆▁▁▂▃▂▂▃▅▅▇▅▅▅▄█▅▅▅▄▆▆▄▄▃▅▅▄▅▅▇▆▇▅
wandb:                      Per Iteration Tracking/step ▁▁▁▁▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb:                                Tracking/ATE RMSE ▁▅█▆▆
wandb:                                Tracking/Depth L1 █▁▁▁▁
wandb:                              Tracking/Depth RMSE █▁▁▁▁
wandb:                       Tracking/Latest Pose Error ▁█▅▅▆
wandb:              Tracking/Latest Relative Pose Error ▁▇▃██
wandb:                                    Tracking/PSNR ▁█▅▇▆
wandb:                                    Tracking/step ▁▃▄▆█
wandb:
wandb: Run summary:
wandb:                     Final Stats/Average Depth L1 0.00485
wandb:                   Final Stats/Average Depth RMSE 0.00485
wandb:                        Final Stats/Average LPIPS 0.07124
wandb:                      Final Stats/Average MS-SSIM 0.97763
wandb:       Final Stats/Average Mapping Frame Time (s) 4.33472
wandb:  Final Stats/Average Mapping Iteration Time (ms) 72.23238
wandb:                         Final Stats/Average PSNR 32.83798
wandb:      Final Stats/Average Tracking Frame Time (s) 2.49429
wandb: Final Stats/Average Tracking Iteration Time (ms) 62.34118
wandb:                         Final Stats/Avg ATE RMSE 0.00302
wandb:                                 Final Stats/step 1
wandb:                                 Mapping/Depth L1 0.00307
wandb:                               Mapping/Depth RMSE 0.00307
wandb:                      Mapping/Number of Gaussians 5118505
wandb:            Mapping/Number of Gaussians - Pruning 5117876
wandb:                                     Mapping/PSNR 35.35604
wandb:                                     Mapping/step 1999
wandb:                 Per Iteration Mapping/Depth Loss 0.00339
wandb:                 Per Iteration Mapping/Image Loss 0.00608
wandb:                       Per Iteration Mapping/Loss 0.00947
wandb:                       Per Iteration Mapping/step 119999
wandb:                Per Iteration Tracking/Depth Loss 2560.95459
wandb:                Per Iteration Tracking/Image Loss 12050.37695
wandb:                      Per Iteration Tracking/Loss 14611.33203
wandb:                      Per Iteration Tracking/step 79959
wandb:                                Tracking/ATE RMSE 0.00302
wandb:                                Tracking/Depth L1 0.00329
wandb:                              Tracking/Depth RMSE 0.00329
wandb:                       Tracking/Latest Pose Error 0.00491
wandb:              Tracking/Latest Relative Pose Error 0.00043
wandb:                                    Tracking/PSNR 34.41476
wandb:                                    Tracking/step 1999
wandb:
wandb: 🚀 View run room0_0 at: https://wandb.ai/noguchis/SplaTAM/runs/dwxuqx98/workspace
wandb: Synced 5 W&B file(s), 402 media file(s), 2 artifact file(s) and 0 other file(s)
wandb: Find logs at: ./wandb/run-20240326_154214-dwxuqx98/logs

正常終了しました。ログを見るに、./experiments/Replica/room0_0 というディレクトリ以下にファイルが生成されているようです。

$ cd experiments/Replica/room0_0/
$ ls -al
total 260100
drwxr-xr-x 3 user user      4096 Mar 26 19:47 .
drwxr-xr-x 3 user user      4096 Mar 26 15:33 ..
-rw-rw-r-- 1 user user      4664 Mar 26 15:42 config.py
drwxr-xr-x 4 user user      4096 Mar 26 20:01 eval
-rw-r--r-- 1 user user 266320474 Mar 26 19:47 params.npz
$

(2) SplaTAM再構築の可視化(インタラクティブ)

python viz_scripts/final_recon.py configs/replica/splatam.py

ぐるぐる回す動画。

(3) SplaTAM再構築の可視化(自動再生的な)

以下のコマンドを叩きます。

python viz_scripts/online_recon.py configs/replica/splatam.py

キャプチャした動画がこちら。(数分近くかかるので、最初のほうのみ録画)

(4) .plyファイルのエクスポート

$ python scripts/export_ply.py configs/replica/splatam.py
Saved PLY format Splat to ./experiments/Replica/room0_0/splat.ply

.plyファイルが保存されました。

$ ls -l ./experiments/Replica/room0_0/splat.ply
-rw-r--r-- 1 user user 348015985 Mar 27 09:49 ./experiments/Replica/room0_0/splat.ply
$

Blenderでこの.plyファイルを読み込ませたのが、こちら。

これをレンダリングするのはとても困難なような…。部屋ですしね。

(5) 3DGS実行(SplaTAM構成)

SplaTAM構成で3DGSを実行します。
まず、configファイルのパスに誤り?があるので、修正します。

diff --git a/configs/replica/post_splatam_opt.py b/configs/replica/post_splatam_opt.py
index 9820f1a..90ad377 100644
--- a/configs/replica/post_splatam_opt.py
+++ b/configs/replica/post_splatam_opt.py
@@ -24,7 +24,7 @@ config = dict(
     ),
     data=dict(
         basedir="./data/Replica",
-        gradslam_data_cfg="./data/replica.yaml",
+        gradslam_data_cfg="./configs/data/replica.yaml",
         sequence="room0",
         desired_image_height=680,
         desired_image_width=1200,
@@ -34,7 +34,7 @@ config = dict(
         num_frames=100,
         eval_stride=5,
         eval_num_frames=400,
-        param_ckpt_path='./experiments/Replica/room0_seed0/params.npz'
+        param_ckpt_path='./experiments/Replica/room0_0/params.npz'
     ),
     train=dict(
         num_iters_mapping=15000,

そして、実行です。

python scripts/post_splatam_opt.py configs/replica/post_splatam_opt.py

25分ほど経過して、処理が終了しました。

Evaluating Final Parameters ...
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 400/400 [04:40<00:00,  1.42it/s]
Average PSNR: 6.09 99:  47%|████████████████████████████████████████████████████████▉                                                                 | 7000/15000 [08:50<09:41, 13.75it/s]
Average Depth RMSE: 2.02█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 400/400 [04:40<00:00,  1.45it/s]
Average MS-SSIM: 0.35
Average LPIPS: 0.74
Mapping Time Step: 99: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15000/15000 [23:05<00:00, 10.83it/s]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [23:05<00:00, 13.85s/it]
Evaluating Final Parameters ...
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 400/400 [04:49<00:00,  1.38it/s]
Average PSNR: 6.06
Average Depth RMSE: 2.03
Average MS-SSIM: 0.34
Average LPIPS: 0.74
Saving parameters to: ./experiments/Replica/Post_SplaTAM_Opt

以下のファイルが生成されています。

$ ls -l ./experiments/Replica/Post_SplaTAM_Opt
total 203356
-rw-rw-r-- 1 user user      2940 Mar 27 10:20 config.py
drwxr-xr-x 3 user user      4096 Mar 27 10:43 eval
drwxr-xr-x 3 user user      4096 Mar 27 10:29 eval_7k
-rw-r--r-- 1 user user 208223940 Mar 27 10:48 params.npz
$

ちなみに、eval_7kは「iteration 7,000のときの評価」という意味で、最終的な評価はevalのようです。
ですのでevalディレクトリ下にあるpng画像ファイルからmp4を生成したものを確認しましょう。

cd ./experiments/Replica/Post_SplaTAM_Opt/eval/plots

# 横幅が1583と奇数なので、scale truncしています。
ffmpeg -r 30 -i %4d.png -vcodec libx264 -pix_fmt yuv420p -r 60 -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" out.mp4

(6) 3DGS実行(グランドトゥルースポーズ)

グランドトゥルースポーズを使用してデータセット上で3DGSを実行します。

設定ファイル configs/replica/gaussian_splatting.py 内のwandb設定(entity)を修正します。
続いて、実行です。

python scripts/gaussian_splatting.py configs/replica/gaussian_splatting.py

20分ほど経過して、処理が終了しました。

Mapping Time Step: 1999: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 30000/30000 [10:46<00:00, 46.43it/s]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2000/2000 [11:05<00:00,  3.01it/s]
Evaluating Final Parameters ...
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [02:49<00:00,  1.18it/s]
Average PSNR: 36.37
Average Depth RMSE: 0.01
Average MS-SSIM: 0.99
Average LPIPS: 0.04
Saving parameters to: ./experiments/Replica_3DGS/room0_0
wandb:
wandb:
wandb: Run history:
wandb:                  Average Depth RMSE █▁
wandb:                       Average LPIPS █▁
wandb:                     Average MS-SSIM ▁█
wandb:                        Average PSNR ▁█
wandb:            Init/Number of Gaussians ▁▂▂▂▃▄▄▄▄▄▄▄▅▅▆▆▇▇▇▇▇▇▇▇▇███████████████
wandb:                           Init/step ▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb:             Learning Rate - Means3D █▇▇▆▆▅▅▄▄▃▃▃▃▃▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb: Number of Gaussians - Densification ▁▁▂▃▃▄▄▅▅▅▆▆▆▆▇▇▇▇██████████████████████
wandb:
wandb: Run summary:
wandb:                  Average Depth RMSE 0.00655
wandb:                       Average LPIPS 0.03805
wandb:                     Average MS-SSIM 0.99034
wandb:                        Average PSNR 36.36539
wandb:            Init/Number of Gaussians 197107
wandb:                           Init/step 1999
wandb:             Learning Rate - Means3D 0.0
wandb: Number of Gaussians - Densification 673583
wandb:
wandb: 🚀 View run room0_0 at: https://wandb.ai/noguchis/SplaTAM/runs/jouaizfx/workspace
wandb: Synced 5 W&B file(s), 400 media file(s), 2 artifact file(s) and 0 other file(s)
wandb: Find logs at: ./wandb/run-20240327_095222-jouaizfx/logs

以下のファイルが生成されています。

$ ls -l ./experiments/Replica_3DGS/room0_0
total 31772
-rw-rw-r-- 1 user user     3788 Mar 27 09:52 config.py
drwxr-xr-x 3 user user     4096 Mar 27 10:05 eval
drwxr-xr-x 3 user user     4096 Mar 27 09:56 eval_7k
-rw-r--r-- 1 user user 32519196 Mar 27 10:07 params.npz

evalディレクトリ下にある.pngファイルをmp4ファイルに変換します。

cd ./experiments/Replica_3DGS/room0_0/eval/plots

# 横幅が1583と奇数なので、scale truncしています。
ffmpeg -r 30 -i %4d.png -vcodec libx264 -pix_fmt yuv420p -r 60 -vf "scale=trunc(iw/2)*2:trunc(ih/2)*2" out.mp4

3. まとめ

弊環境のRTX 4090 Laptop GPU (16GB) でもきちんと動きました。24GBなくても大丈夫そうです。
入力画像が2000枚で四時間ほど学習に要します(Laptop GPU時)。

iPhoneでキャプチャした動画で試したい…。

引用

@inproceedings{sucar2021imap,
  title={iMAP: Implicit mapping and positioning in real-time},
  author={Sucar, Edgar and Liu, Shikun and Ortiz, Joseph and Davison, Andrew J},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  pages={6229--6238},
  year={2021}
}

この記事が気に入ったらサポートをしてみませんか?