From b6cd82345780ff5b3ea94b8fa76dcd217b163cab Mon Sep 17 00:00:00 2001 From: Filip Jeretina Date: Thu, 26 Sep 2024 19:38:23 +0200 Subject: [PATCH] Refactor the calibration handling a little. Undistort and rectify the color camera images which tof is aligned to. --- rerun_py/depthai_viewer/_backend/device.py | 22 +---- .../depthai_viewer/_backend/packet_handler.py | 83 +++++++++++++++---- 2 files changed, 71 insertions(+), 34 deletions(-) diff --git a/rerun_py/depthai_viewer/_backend/device.py b/rerun_py/depthai_viewer/_backend/device.py index 729da9c55ab2..b43064f048c9 100644 --- a/rerun_py/depthai_viewer/_backend/device.py +++ b/rerun_py/depthai_viewer/_backend/device.py @@ -80,8 +80,7 @@ def update(self) -> None: class Device: id: str - intrinsic_matrix: Dict[Tuple[dai.CameraBoardSocket, int, int], NDArray[np.float32]] = {} - calibration_data: Optional[dai.CalibrationHandler] = None + calibration_handler: Optional[dai.CalibrationHandler] = None use_encoding: bool = False store: Store @@ -102,7 +101,7 @@ def __init__(self, device_id: str, store: Store): self.id = device_id self.set_oak(OakCamera(device_id, args={"irFloodBrightness": 0, "irDotBrightness": 0})) self.store = store - self._packet_handler = PacketHandler(self.store, self.get_intrinsic_matrix) + self._packet_handler = PacketHandler(self.store, self._oak.device.readCalibration()) print("Oak cam: ", self._oak) # self.start = time.time() # self._profiler.enable() @@ -116,21 +115,6 @@ def set_oak(self, oak_cam: Optional[OakCamera]) -> None: def is_closed(self) -> bool: return self._oak is not None and self._oak.device.isClosed() - def get_intrinsic_matrix(self, board_socket: dai.CameraBoardSocket, width: int, height: int) -> NDArray[np.float32]: - if self.intrinsic_matrix.get((board_socket, width, height)) is not None: - return self.intrinsic_matrix.get((board_socket, width, height)) # type: ignore[return-value] - if self.calibration_data is None: - raise Exception("Missing calibration data!") - try: - M_right = self.calibration_data.getCameraIntrinsics( # type: ignore[union-attr] - board_socket, dai.Size2f(width, height) - ) - except RuntimeError: - print("No intrinsics found for camera: ", board_socket, " assuming default.") - f_len = (height * width) ** 0.5 - M_right = [[f_len, 0, width / 2], [0, f_len, height / 2], [0, 0, 1]] - self.intrinsic_matrix[(board_socket, width, height)] = np.array(M_right).reshape(3, 3) - return self.intrinsic_matrix[(board_socket, width, height)] def _get_possible_stereo_pairs_for_cam( self, cam: dai.CameraFeatures, connected_camera_features: List[dai.CameraFeatures] @@ -651,7 +635,7 @@ def update_pipeline(self, runtime_only: bool) -> Message: self._oak.poll() except RuntimeError: return ErrorMessage("Runtime error when polling the device. Check the terminal for more info.") - self.calibration_data = self._oak.device.readCalibration() + self.calibration_handler = self._oak.device.readCalibration() self.intrinsic_matrix = {} return InfoMessage("Pipeline started") if running else ErrorMessage("Couldn't start pipeline") diff --git a/rerun_py/depthai_viewer/_backend/packet_handler.py b/rerun_py/depthai_viewer/_backend/packet_handler.py index 8e757f1e2587..989796375276 100644 --- a/rerun_py/depthai_viewer/_backend/packet_handler.py +++ b/rerun_py/depthai_viewer/_backend/packet_handler.py @@ -29,6 +29,7 @@ from depthai_viewer._backend.store import Store from depthai_viewer._backend.topic import Topic from depthai_viewer.components.rect2d import RectFormat +from typing import Dict class PacketHandlerContext(BaseModel): # type: ignore[misc] @@ -43,32 +44,58 @@ class DetectionContext(PacketHandlerContext): board_socket: dai.CameraBoardSocket +class CachedCalibrationHandler: + calibration_handler: dai.CalibrationHandler + intrinsic_matrix: Dict[Tuple[dai.CameraBoardSocket, int, int], NDArray[np.float32]] = {} + distortion_coefficients: Dict[dai.CameraBoardSocket, NDArray[np.float32]] = {} + + def __init__(self, calibration_handler: dai.CalibrationHandler): + self.calibration_handler = calibration_handler + + def get_intrinsic_matrix(self, board_socket: dai.CameraBoardSocket, width: int, height: int) -> NDArray[np.float32]: + if self.intrinsic_matrix.get((board_socket, width, height)) is not None: + return self.intrinsic_matrix.get((board_socket, width, height)) # type: ignore[return-value] + try: + M = self.calibration_handler.getCameraIntrinsics( # type: ignore[union-attr] + board_socket, dai.Size2f(width, height) + ) + except RuntimeError: + print("No intrinsics found for camera: ", board_socket, " assuming default.") + f_len = (height * width) ** 0.5 + M = [[f_len, 0, width / 2], [0, f_len, height / 2], [0, 0, 1]] + self.intrinsic_matrix[(board_socket, width, height)] = np.array(M).reshape(3, 3) + return self.intrinsic_matrix[(board_socket, width, height)] + + def get_distortion_coefficients(self, board_socket: dai.CameraBoardSocket) -> NDArray[np.float32]: + if self.distortion_coefficients.get(board_socket) is not None: + return self.distortion_coefficients.get(board_socket) + try: + D = self.calibration_handler.getDistortionCoefficients(board_socket) # type: ignore[union-attr] + except RuntimeError: + print("No distortion coefficients found for camera: ", board_socket, " assuming default.") + D = np.array([0, 0, 0, 0, 0]) + self.distortion_coefficients[board_socket] = np.array(D) + return self.distortion_coefficients[board_socket] + + class PacketHandler: store: Store _ahrs: Mahony - _get_camera_intrinsics: Callable[[dai.CameraBoardSocket, int, int], NDArray[np.float32]] + _calibration_handler: CachedCalibrationHandler - def __init__( - self, store: Store, intrinsics_getter: Callable[[dai.CameraBoardSocket, int, int], NDArray[np.float32]] - ): + def __init__(self, store: Store, calibration_handler: dai.CalibrationHandler): viewer.init(f"Depthai Viewer {store.viewer_address}") print("Connecting to viewer at", store.viewer_address) viewer.connect(store.viewer_address) self.store = store self._ahrs = Mahony(frequency=100) self._ahrs.Q = np.array([1, 0, 0, 0], dtype=np.float64) - self.set_camera_intrinsics_getter(intrinsics_getter) + self._calibration_handler = CachedCalibrationHandler(calibration_handler) def reset(self) -> None: self._ahrs = Mahony(frequency=100) self._ahrs.Q = np.array([1, 0, 0, 0], dtype=np.float64) - def set_camera_intrinsics_getter( - self, camera_intrinsics_getter: Callable[[dai.CameraBoardSocket, int, int], NDArray[np.float32]] - ) -> None: - # type: ignore[assignment, misc] - self._get_camera_intrinsics = camera_intrinsics_getter - def log_dai_packet(self, node: dai.Node, packet: dai.Buffer, context: Optional[PacketHandlerContext]) -> None: if isinstance(packet, dai.ImgFrame): board_socket = None @@ -147,7 +174,17 @@ def log_packet( else: print("Unknown packet type:", type(packet)) - def _log_img_frame(self, frame: dai.ImgFrame, board_socket: dai.CameraBoardSocket) -> None: + def _log_img_frame( + self, + frame: dai.ImgFrame, + board_socket: dai.CameraBoardSocket, + intrinsics_matrix: Optional[NDArray[np.float32]] = None, + distortion_coefficients: Optional[NDArray[np.float32]] = None, + ) -> None: + """ + Log an image frame to the viewer. + Optionally undistort and rectify the image using the provided intrinsic matrix and distortion coefficients. + """ viewer.log_rigid3( f"{board_socket.name}/transform", child_from_parent=([0, 0, 0], [1, 0, 0, 0]), xyz="RDF" ) # TODO(filip): Enable the user to lock the camera rotation in the UI @@ -158,13 +195,20 @@ def _log_img_frame(self, frame: dai.ImgFrame, board_socket: dai.CameraBoardSocke else frame.getData() ) h, w = frame.getHeight(), frame.getWidth() + # 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: + 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) + if frame.getType() == dai.ImgFrame.Type.BITSTREAM: img_frame = cv2.cvtColor(cv2.imdecode(img_frame, cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB) h, w = img_frame.shape[:2] child_from_parent: NDArray[np.float32] try: - child_from_parent = self._get_camera_intrinsics( # type: ignore[call-arg, misc, arg-type] + child_from_parent = self._calibration_handler.get_intrinsic_matrix( # type: ignore[call-arg, misc, arg-type] board_socket, w, h # type: ignore[call-arg, misc, arg-type] ) except Exception: @@ -231,7 +275,16 @@ def _on_tof_packet( component: ToFComponent, ) -> None: if packet.aligned_frame: - self._log_img_frame(packet.aligned_frame, dai.CameraBoardSocket(packet.aligned_frame.getInstanceNum())) + rgb_size = (packet.aligned_frame.getWidth(), packet.aligned_frame.getHeight()) + M = self._calibration_handler.get_intrinsic_matrix( + dai.CameraBoardSocket(packet.aligned_frame.getInstanceNum()), *rgb_size + ) + D = self._calibration_handler.get_distortion_coefficients( + dai.CameraBoardSocket(packet.aligned_frame.getInstanceNum()) + ) + self._log_img_frame( + packet.aligned_frame, dai.CameraBoardSocket(packet.aligned_frame.getInstanceNum()), M, D + ) depth_frame = packet.frame if packet.aligned_frame: @@ -242,7 +295,7 @@ def _on_tof_packet( if not packet.aligned_frame: viewer.log_rigid3(f"{ent_path_root}/transform", child_from_parent=([0, 0, 0], [1, 0, 0, 0]), xyz="RDF") try: - intrinsics = self._get_camera_intrinsics(component.camera_socket, 640, 480) + intrinsics = self._calibration_handler.get_intrinsic_matrix(component.camera_socket, 640, 480) except Exception: intrinsics = np.array([[471.451, 0.0, 317.897], [0.0, 471.539, 245.027], [0.0, 0.0, 1.0]]) viewer.log_pinhole(