From 48d307ab487b4022e70d04f7c31f7d2a29e5f4b4 Mon Sep 17 00:00:00 2001 From: Filip Jeretina Date: Fri, 27 Sep 2024 15:46:32 +0200 Subject: [PATCH] Fixed up rectification + udistort for camera streams that stereo depth is aligned to. --- rerun_py/depthai_viewer/_backend/device.py | 1 + .../_backend/device_configuration.py | 1 + .../depthai_viewer/_backend/packet_handler.py | 18 ++++++++++++++++-- 3 files changed, 18 insertions(+), 2 deletions(-) diff --git a/rerun_py/depthai_viewer/_backend/device.py b/rerun_py/depthai_viewer/_backend/device.py index 748d6e2010af..6da73dd3cc81 100644 --- a/rerun_py/depthai_viewer/_backend/device.py +++ b/rerun_py/depthai_viewer/_backend/device.py @@ -537,6 +537,7 @@ def update_pipeline(self, runtime_only: bool) -> Message: aligned_camera = self._get_camera_config_by_socket(config, config.stereo.align) if not aligned_camera: return ErrorMessage(f"{config.stereo.align} is not configured. Couldn't create stereo pair.") + aligned_camera.is_used_as_stereo_align = True self._queues.append((self._stereo, self._oak.queue(self._stereo.out.main))) if self._oak.device.getConnectedIMU() != "NONE" and self._oak.device.getConnectedIMU() != "": diff --git a/rerun_py/depthai_viewer/_backend/device_configuration.py b/rerun_py/depthai_viewer/_backend/device_configuration.py index 0982454ef3db..0fe023ead605 100644 --- a/rerun_py/depthai_viewer/_backend/device_configuration.py +++ b/rerun_py/depthai_viewer/_backend/device_configuration.py @@ -182,6 +182,7 @@ class CameraConfiguration(BaseModel): # type: ignore[misc] name: str = "" tof_align: Optional[dai.CameraBoardSocket] = None + is_used_as_stereo_align: bool = False class Config: arbitrary_types_allowed = True diff --git a/rerun_py/depthai_viewer/_backend/packet_handler.py b/rerun_py/depthai_viewer/_backend/packet_handler.py index 2b8601abf083..7faf91cd624a 100644 --- a/rerun_py/depthai_viewer/_backend/packet_handler.py +++ b/rerun_py/depthai_viewer/_backend/packet_handler.py @@ -196,12 +196,16 @@ def _log_img_frame( else frame.getData() ) h, w = frame.getHeight(), frame.getWidth() + was_undistorted_and_rectified = False # If the image is a cv frame try to undistort and rectify it if intrinsics_matrix is not None and distortion_coefficients is not None: + if frame.getType() == dai.RawImgFrame.Type.NV12: + img_frame = frame.getCvFrame() map_x, map_y = cv2.initUndistortRectifyMap( intrinsics_matrix, distortion_coefficients, None, intrinsics_matrix, (w, h), cv2.CV_32FC1 ) img_frame = cv2.remap(img_frame, map_x, map_y, cv2.INTER_LINEAR) + was_undistorted_and_rectified = True if frame.getType() == dai.ImgFrame.Type.BITSTREAM: img_frame = cv2.cvtColor(cv2.imdecode(img_frame, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB) @@ -226,7 +230,9 @@ def _log_img_frame( ) entity_path = f"{board_socket.name}/transform/{cam}/Image" - if frame.getType() == dai.RawImgFrame.Type.NV12: # or frame.getType() == dai.RawImgFrame.Type.YUV420p + if ( + not was_undistorted_and_rectified and frame.getType() == dai.RawImgFrame.Type.NV12 + ): # or frame.getType() == dai.RawImgFrame.Type.YUV420p encoding = viewer.ImageEncoding.NV12 viewer.log_encoded_image( entity_path, @@ -249,7 +255,15 @@ def _on_camera_frame(self, packet: FramePacket, board_socket: dai.CameraBoardSoc ): # Skip tof aligned frames - they will be logged on_tof_packet return - self._log_img_frame(packet.msg, board_socket) + intrinsics = None + distortion_coefficients = None + if this_cam := list(filter(lambda cam: cam.board_socket == board_socket, self.store.pipeline_config.cameras)): + if this_cam[0].is_used_as_stereo_align: + intrinsics = self._calibration_handler.get_intrinsic_matrix( + board_socket, packet.msg.getWidth(), packet.msg.getHeight() + ) + distortion_coefficients = self._calibration_handler.get_distortion_coefficients(board_socket) + self._log_img_frame(packet.msg, board_socket, intrinsics, distortion_coefficients) def on_imu(self, packet: IMUPacket) -> None: gyro: dai.IMUReportGyroscope = packet.gyroscope