fix undistort intrinsics in datamanager

nerfstudio/cameras/camera_optimizers.py

@@ -33,6 +33,7 @@ from nerfstudio.configs.base_config import InstantiateConfig
 from nerfstudio.utils import poses as pose_utils
 from nerfstudio.engine.optimizers import OptimizerConfig
 from nerfstudio.engine.schedulers import SchedulerConfig
+from nerfstudio.cameras.cameras import Cameras
 
 
 @dataclass
@@ -146,6 +147,13 @@ class CameraOptimizer(nn.Module):
             raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
             raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
 
+    def apply_to_camera(self, camera: Cameras) -> None:
+        """Apply the pose correction to the raybundle"""
+        if self.config.mode != "off":
+            correction_matrices = self(raybundle.camera_indices.squeeze())  # type: ignore
+            raybundle.origins = raybundle.origins + correction_matrices[:, :3, 3]
+            raybundle.directions = torch.bmm(correction_matrices[:, :3, :3], raybundle.directions[..., None]).squeeze()
+
     def get_loss_dict(self, loss_dict: dict) -> None:
         """Add regularization"""
         if self.config.mode != "off":

nerfstudio/configs/base_config.py

@@ -145,7 +145,7 @@ class ViewerConfig(PrintableConfig):
     """Whether to kill the training job when it has completed. Note this will stop rendering in the viewer."""
     image_format: Literal["jpeg", "png"] = "jpeg"
     """Image format viewer should use; jpeg is lossy compression, while png is lossless."""
-    jpeg_quality: int = 90
+    jpeg_quality: int = 70
     """Quality tradeoff to use for jpeg compression."""
     make_share_url: bool = False
     """Viewer beta feature: print a shareable URL. `vis` must be set to viewer_beta; this flag is otherwise ignored."""

nerfstudio/data/datamanagers/full_images_datamanager.py

@@ -141,12 +141,25 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
                         0,
                     ]
                 )
-                image = cv2.undistort(image, K, distortion_params, None, K)  # Should update K in-place
+                newK,roi = cv2.getOptimalNewCameraMatrix(K,distortion_params,(image.shape[1],image.shape[0]),0)
+                image = cv2.undistort(image, K, distortion_params, None, newK)
+                #crop the image and update the intrinsics accordingly
+                x,y,w,h = roi
+                image = image[y:y+h,x:x+w]
+                K = newK
+                #update the width, height
+                self.train_dataset.cameras.width[i] = w
+                self.train_dataset.cameras.height[i] = h
+                
             elif camera.camera_type.item() == CameraType.FISHEYE.value:
                 distortion_params = np.array(
                     [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]
                 )
-                image = cv2.fisheye.undistortImage(image, K, distortion_params, None, K)
+                newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(K, distortion_params, (image.shape[1],image.shape[0]), np.eye(3), balance=0)
+                map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, distortion_params, np.eye(3), newK, (image.shape[1],image.shape[0]), cv2.CV_32FC1)
+                # and then remap:
+                image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
+                K = newK
             else:
                 raise NotImplementedError("Only perspective and fisheye cameras are supported")
             data["image"] = torch.from_numpy(image)
@@ -163,10 +176,10 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
 
             cached_train.append(data)
 
-            self.train_dataset.cameras[i].fx = K[0, 0]
-            self.train_dataset.cameras[i].fy = K[1, 1]
-            self.train_dataset.cameras[i].cx = K[0, 2]
-            self.train_dataset.cameras[i].cy = K[1, 2]
+            self.train_dataset.cameras.fx[i] = float(K[0, 0])
+            self.train_dataset.cameras.fy[i] = float(K[1, 1])
+            self.train_dataset.cameras.cx[i] = float(K[0, 2])
+            self.train_dataset.cameras.cy[i] = float(K[1, 2])
 
         cached_eval = []
         CONSOLE.log("Caching / undistorting eval images")
@@ -191,12 +204,24 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
                         0,
                     ]
                 )
-                image = cv2.undistort(image, K, distortion_params, None, K)  # Should update K in-place
+                newK,roi = cv2.getOptimalNewCameraMatrix(K,distortion_params,(image.shape[1],image.shape[0]),0)
+                image = cv2.undistort(image, K, distortion_params, None, newK)
+                #crop the image and update the intrinsics accordingly
+                x,y,w,h = roi
+                image = image[y:y+h,x:x+w]
+                K = newK
+                #update the width, height
+                self.train_dataset.cameras.width[i] = w
+                self.train_dataset.cameras.height[i] = h
             elif camera.camera_type.item() == CameraType.FISHEYE.value:
                 distortion_params = np.array(
                     [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]
                 )
-                image = cv2.fisheye.undistortImage(image, K, distortion_params, None, K)
+                newK = cv2.fisheye.estimateNewCameraMatrixForUndistortRectify(K, distortion_params, (image.shape[1],image.shape[0]), np.eye(3), balance=0)
+                map1, map2 = cv2.fisheye.initUndistortRectifyMap(K, distortion_params, np.eye(3), newK, (image.shape[1],image.shape[0]), cv2.CV_32FC1)
+                # and then remap:
+                image = cv2.remap(image, map1, map2, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
+                K = newK
             else:
                 raise NotImplementedError("Only perspective and fisheye cameras are supported")
             data["image"] = torch.from_numpy(image)
@@ -300,6 +325,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
 
         data = deepcopy(self.cached_train[image_idx])
         data["image"] = data["image"].to(self.device)
+        data['cam_idx'] = torch.tensor(image_idx).to(self.device)
 
         assert len(self.train_dataset.cameras.shape) == 1, "Assumes single batch dimension"
         camera = self.train_dataset.cameras[image_idx : image_idx + 1].to(self.device)

nerfstudio/pipelines/base_pipeline.py

@@ -259,9 +259,12 @@ class VanillaPipeline(Pipeline):
         self.datamanager: DataManager = config.datamanager.setup(
             device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank
         )
-        # TODO make this work with normal nerfacto too
-        pts = self.datamanager.train_dataparser_outputs.metadata["points3D_xyz"]
-        pts_rgb = self.datamanager.train_dataparser_outputs.metadata["points3D_rgb"]
+        # TODO make cleaner
+        seed_pts = None
+        if hasattr(self.datamanager, "train_dataparser_outputs") and 'points3D_xyz' in self.datamanager.train_dataparser_outputs.metadata:
+            pts = self.datamanager.train_dataparser_outputs.metadata["points3D_xyz"]
+            pts_rgb = self.datamanager.train_dataparser_outputs.metadata["points3D_rgb"]
+            seed_pts = (pts,pts_rgb)
         self.datamanager.to(device)
         # TODO(ethan): get rid of scene_bounds from the model
         assert self.datamanager.train_dataset is not None, "Missing input dataset"
@@ -272,7 +275,7 @@ class VanillaPipeline(Pipeline):
             metadata=self.datamanager.train_dataset.metadata,
             device=device,
             grad_scaler=grad_scaler,
-            seed_points=(pts, pts_rgb),
+            seed_points=seed_pts,
         )
         self.model.to(device)
 

nerfstudio/viewer_beta/render_state_machine.py

@@ -190,7 +190,7 @@ class RenderStateMachine(threading.Thread):
     def run(self):
         """Main loop for the render thread"""
         while True:
-            if not self.render_trigger.wait(0.2):
+            if not self.render_trigger.wait(0.1):
                 # if we haven't received a trigger in a while, send a static action
                 if self.viewer.camera_state is not None:
                     self.action(RenderAction(action="static", camera_state=self.viewer.camera_state))
@@ -208,7 +208,7 @@ class RenderStateMachine(threading.Thread):
             except viewer_utils.IOChangeException:
                 # if we got interrupted, don't send the output to the viewer
                 continue
-            self._send_output_to_viewer(outputs)
+            self._send_output_to_viewer(outputs,static_render = (action.action in ["static",'step']))
 
     def check_interrupt(self, frame, event, arg):
         """Raises interrupt when flag has been set and not already on lowest resolution.
@@ -220,7 +220,7 @@ class RenderStateMachine(threading.Thread):
                 raise viewer_utils.IOChangeException
         return self.check_interrupt
 
-    def _send_output_to_viewer(self, outputs: Dict[str, Any]):
+    def _send_output_to_viewer(self, outputs: Dict[str, Any], static_render: bool = True):
         """Chooses the correct output and sends it to the viewer
 
         Args:
@@ -260,10 +260,11 @@ class RenderStateMachine(threading.Thread):
         depth = (
             outputs["gl_z_buf_depth"].cpu().numpy() * self.viser_scale_ratio if "gl_z_buf_depth" in outputs else None
         )
+        jpg_quality = self.viewer.config.jpeg_quality if static_render else 40
         self.viewer.viser_server.set_background_image(
             selected_output.cpu().numpy(),
             format=self.viewer.config.image_format,
-            jpeg_quality=self.viewer.config.jpeg_quality,
+            jpeg_quality=jpg_quality,
             depth=depth,
         )