From 490892ff1fbfd2ea8b8a49be6314ed495442887d Mon Sep 17 00:00:00 2001 From: ayaanzhaque Date: Fri, 7 Apr 2023 19:13:47 +0000 Subject: [PATCH] add scenebox changes --- nerfstudio/data/dataparsers/nerfstudio_dataparser.py | 7 ++++++- nerfstudio/models/nerfacto.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py b/nerfstudio/data/dataparsers/nerfstudio_dataparser.py index 7b283d747..cd85eea64 100644 --- a/nerfstudio/data/dataparsers/nerfstudio_dataparser.py +++ b/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 [] diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 53173640a..10bc0fc61 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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)