Update visualizer script.

Author: santiago-lopez-ekumen

bash_scripts/nerfcapture.bash

@@ -30,6 +30,7 @@ scene="$(basename $2)_nerfcapture"
 echo $base_dir $scene
 
 # Convert Spectacular AI dataset to NeRF Capture.
+echo "Exporting dataset to NeRF Capture format..."
 python3 spectacularAI-sdk/python/mapping/replay_to_nerf.py $2 \
     --format=nerfcapture --fast \
     --image_format=png \
@@ -39,16 +40,18 @@ python3 spectacularAI-sdk/python/mapping/replay_to_nerf.py $2 \
 
 frames=$(ls -l ${base_dir}/${scene}/rgb | wc -l)
 if [ -z $frames ] || [ $frames -eq 0 ]; then
-    echo "Not enough frames at $base_dir/$scene/rgb!"
+    echo "ERROR: Not enough frames at $base_dir/$scene/rgb!"
     exit
 fi
 
+echo "Building SplaTAM expected directories..."
 python3 scripts/nerfcapture2dataset.py --config $1 --base_dir $base_dir --scene $scene --frames $frames
 
 # Run SplaTAM
+echo "Running SplaTAM..."
 python3 scripts/splatam.py $1 --base_dir $base_dir --scene $scene --frames $frames
 
-exit 0
-
 # Visualize SplaTAM Output
-python3 viz_scripts/final_recon.py $1
+echo "Running visualizer..."
+python3 viz_scripts/final_recon.py $1 --base_dir $base_dir --scene $scene --frames $frames
+

viz_scripts/final_recon.py

@@ -279,6 +279,9 @@ def visualize(scene_path, cfg):
     parser = argparse.ArgumentParser()
 
     parser.add_argument("experiment", type=str, help="Path to experiment file")
+    parser.add_argument("--base_dir", type=str, help="Path to the dataset.")
+    parser.add_argument("--scene", type=str, help="Name of the NeRFCapture dataset. Usually has a _nerfcapture suffix.")
+    parser.add_argument("--frames", type=int, help="Amount of frames to process.")
 
     args = parser.parse_args()
 
@@ -287,6 +290,10 @@ def visualize(scene_path, cfg):
     ).load_module()
 
     seed_everything(seed=experiment.config["seed"])
+    experiment.config['workdir'] = args.base_dir + "/" + args.scene
+    experiment.config['data']['num_frames'] = experiment.config['num_frames'] = int(args.frames / 3) # SplaTAM usually stars to go nuts after 1/3 of the dataset frames.
+    experiment.config['data']['basedir'] = experiment.config['basedir'] = args.base_dir
+    experiment.config['data']['sequence'] = experiment.config['scene_name'] = args.scene
 
     if "scene_path" not in experiment.config:
         results_dir = os.path.join(