From 40c05be7a60bde25adc28cd281bb9d742c2a77e9 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Sat, 14 Oct 2023 13:29:49 -0700 Subject: [PATCH] fix undistort intrinsics in datamanager --- nerfstudio/cameras/camera_optimizers.py | 8 ++++ nerfstudio/configs/base_config.py | 2 +- .../datamanagers/full_images_datamanager.py | 42 +++++++++++++++---- nerfstudio/pipelines/base_pipeline.py | 11 +++-- .../viewer_beta/render_state_machine.py | 9 ++-- 5 files changed, 55 insertions(+), 17 deletions(-) diff --git a/nerfstudio/cameras/camera_optimizers.py b/nerfstudio/cameras/camera_optimizers.py index 7b7648884..fbf01d664 100644 --- a/nerfstudio/cameras/camera_optimizers.py +++ b/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": diff --git a/nerfstudio/configs/base_config.py b/nerfstudio/configs/base_config.py index d66a9b519..45dc674fc 100644 --- a/nerfstudio/configs/base_config.py +++ b/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.""" diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 2407f6727..e170e0961 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/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) diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index a520f563c..9cf86dc86 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/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) diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index cb6338942..a6b06570b 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/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, )