diff --git a/ar_tag_pose_charuco_board.py b/ar_tag_pose_charuco_board.py new file mode 100644 index 0000000..5cf8bda --- /dev/null +++ b/ar_tag_pose_charuco_board.py @@ -0,0 +1,104 @@ +from PIL.ImageOps import posterize + +from ar_tag_pose import ArTagPose +from aruco_types import ARUCO_DICT +import numpy as np +import cv2 as cv + +def factory(dict_type: int, board_size: cv.typing.Size, square_length: float = 5.0, marker_length: float = 3.0, margin_length=0): + im_w = 800 + im_h = 600 + aruco_dict = cv.aruco.getPredefinedDictionary(dict_type) + board = cv.aruco.CharucoBoard(board_size, square_length, marker_length, aruco_dict) + params = cv.aruco.CharucoParameters() + detector = cv.aruco.CharucoDetector(board, params) + image = board.generateImage((im_w, im_h), None, margin_length, 1) + keyname = [key for key, val in ARUCO_DICT.items() if val == dict_type][0] + cv.imwrite(f"results/aruco_board_w{board_size[0]}_h{board_size[1]}_{keyname}.png", image) + + return detector, board + +class ArTagPoseCharucoBoard(ArTagPose): + def __init__(self, dict_type: int, board_size: tuple[int, int]): + self.board_size = board_size + self.calib_filename = f"results/cam_calib_aruco_board_w{board_size[0]}_h{board_size[1]}.npz" + self.detector, self.board = factory(dict_type, board_size) + + def calibrate_save(self): + np.savez(f"{self.calib_filename}", mtx=self.mtx, dist=self.dist, rvecs=self.r_vecs, tvecs=self.t_vecs) + + def calibrate_load(self): + try: + with np.load(self.calib_filename) as X: + self.mtx, self.dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')] + except FileNotFoundError: + pass + except KeyError: + pass + + def calibrate_points(self, image: cv.typing.MatLike): + success = False + img_size = (0,0) + obj_points = None + img_points = None + gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY) + corners, ids, _, _ = self.detector.detectBoard(gray) + if corners is not None: + if len(ids) >= len(self.board.getIds()): + _obj_points, _img_points = self.board.matchImagePoints(corners, ids) + if len(_obj_points) > 0 and len(_img_points) > 0: + success =True + img_size = gray.shape + obj_points = _obj_points + img_points = _img_points + + return success, img_size, obj_points, img_points + + def _process(self, image, r_vecs: np.ndarray=None, t_vecs: np.ndarray=None, draw_marker_box=True, draw_marker_id=False, draw_marker_axis=True): + pose = False + r_vec = None + t_vec = None + if not self.is_calibrated(): + return pose, None, None + + gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY) + charuco_corners, charuco_ids, marker_corners, marker_ids = self.detector.detectBoard(gray) + + if charuco_corners is not None: + + # Draw marker box + if draw_marker_box: + cv.aruco.drawDetectedMarkers(image, marker_corners) + + # Draw marker id + if draw_marker_id: + cv.aruco.drawDetectedCornersCharuco(image, charuco_corners, charuco_ids) + + # Draw marker axis + if draw_marker_axis: + _r_vecs, _t_vecs, _ = cv.aruco.estimatePoseSingleMarkers(marker_corners, markerLength=0.1, + cameraMatrix=self.mtx, distCoeffs=self.dist) + for i in range(len(_r_vecs)): + cv.drawFrameAxes(image, self.mtx, self.dist, _r_vecs[i], _t_vecs[i], 0.05) + + if len(charuco_ids) >= len(self.board.getIds()): + obj_points, img_points = self.board.matchImagePoints(charuco_corners, charuco_ids) + offset = int(self.board.getMarkerLength() * self.board.getSquareLength()) + obj_points -= [offset, offset, 0] + if self.mtx is not None and self.dist is not None: + pose, r_vec, t_vec = cv.solvePnP(obj_points, img_points, self.mtx, self.dist, useExtrinsicGuess=False, + rvec=r_vecs, tvec=t_vecs, flags=cv.SOLVEPNP_ITERATIVE) + if pose: + # Draw + cv.drawFrameAxes(image, self.mtx, self.dist, r_vec, t_vec, 10) + + return pose, r_vec, t_vec + +def main(): + pose = ArTagPoseCharucoBoard(10, (6,7)) + pose.calibrate_load() + pose.process() + + +if __name__ == '__main__': + main() \ No newline at end of file