add scenebox changes

nerfstudio/data/dataparsers/nerfstudio_dataparser.py

@@ -18,7 +18,7 @@ from __future__ import annotations
 import math
 from dataclasses import dataclass, field
 from pathlib import Path, PurePath
-from typing import Optional, Type
+from typing import Optional, Tuple, Type
 
 import numpy as np
 import torch
@@ -64,6 +64,8 @@ class NerfstudioDataParserConfig(DataParserConfig):
     """The fraction of images to use for training. The remaining images are for eval."""
     depth_unit_scale_factor: float = 1e-3
     """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
+    shift_poses: Optional[Tuple[float, float, float]] = None
+    """Shifts the poses by the given amount. Useful for shifting the cameras for front-facing scenes."""
 
 
 @dataclass
@@ -236,6 +238,9 @@ class Nerfstudio(DataParser):
 
         poses[:, :3, 3] *= scale_factor
 
+        if self.config.shift_poses is not None:
+            poses[:, :3, 3] += torch.tensor(self.config.shift_poses)
+
         # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.
         image_filenames = [image_filenames[i] for i in indices]
         mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []

nerfstudio/models/nerfacto.py

@@ -128,6 +128,8 @@ class NerfactoModelConfig(ModelConfig):
     """Whether to predict normals or not."""
     disable_scene_contraction: bool = False
     """Whether to disable scene contraction or not."""
+    use_aabb_collider: bool = False
+    """Whether to use aabb collider or not."""
 
 
 class NerfactoModel(Model):
@@ -206,7 +208,10 @@ class NerfactoModel(Model):
         )
 
         # Collider
-        self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
+        if self.config.use_aabb_collider:
+            self.collider = AABBBoxCollider(scene_box=self.scene_box)
+        else:
+            self.collider = NearFarCollider(near_plane=self.config.near_plane, far_plane=self.config.far_plane)
 
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)