diff --git a/ar_tag_pose.py b/ar_tag_pose.py new file mode 100644 index 0000000..25af704 --- /dev/null +++ b/ar_tag_pose.py @@ -0,0 +1,116 @@ +import abc +import math +import cv2 as cv +import numpy as np +from utils import to_board_pose_euler + +class ArTagPose(abc.ABC): + detector: cv.aruco.ArucoDetector | cv.aruco.CharucoDetector = None + board: cv.aruco.GridBoard | cv.aruco.CharucoBoard = None + mtx = None + dist = None + r_vecs = None + t_vecs = None + img_size = None + obj_points_list = [] + img_points_list = [] + + @staticmethod + def cap_init(cap: cv.VideoCapture): + cap.set(cv.CAP_PROP_SETTINGS, 1) + cap.set(cv.CAP_PROP_FOURCC, cv.VideoWriter.fourcc('M', 'J', 'P', 'G')) + cap.set(cv.CAP_PROP_FPS, 60.0) + cap.set(cv.CAP_PROP_FRAME_WIDTH, 1280) + cap.set(cv.CAP_PROP_FRAME_HEIGHT, 1024) + cap.set(cv.CAP_PROP_EXPOSURE, 30) + cap.set(cv.CAP_PROP_AUTO_EXPOSURE, 1) + + def is_calibrated(self): + return self.mtx is not None and self.dist is not None + + def calibrate_camera(self, cal_img_size: tuple[int,int], r_vecs=None, t_vecs=None, flags=0): + success = False + if len(self.obj_points_list) > 0 and len(self.img_points_list) > 0: + ret, mtx, dist, r_vecs, t_vecs = cv.calibrateCamera(self.obj_points_list, self.img_points_list, cal_img_size, self.mtx, self.dist, + rvecs=r_vecs, tvecs=t_vecs, flags=flags) + + print(f"Camera calibration finished: ", end='') + if ret < 2: + print(f"Success ({ret})!") + success = True + self.mtx = mtx + self.dist = dist + self.r_vecs = r_vecs + self.t_vecs = t_vecs + else: + print(f"Failed ({ret})!") + + return success + + @abc.abstractmethod + def calibrate_points(self, image: cv.typing.MatLike): + pass + + @abc.abstractmethod + def calibrate_save(self): + pass + + @abc.abstractmethod + def calibrate_load(self): + pass + + def process(self): + + # Open camera + cap = cv.VideoCapture(0) + ArTagPose.cap_init(cap) + + if not cap.isOpened(): + print("Cannot open camera") + exit() + + self.calibrate_load() + cal_img_size = None + last_rotation = [0, 0, 0] + while True: + # Read a new frame + ok, img = cap.read() + if not ok: + return + k = cv.waitKey(1) & 0xFF + if k == ord('q'): + break + elif k == ord('s'): + success, cal_img_size, obj_points, img_points = self.calibrate_points(img) + if success: + print("Frame captured") + self.obj_points_list.append(obj_points) + self.img_points_list.append(img_points) + else: + print("Point matching failed, try again") + + elif k == ord('c'): + if self.calibrate_camera(cal_img_size): + self.calibrate_save() + elif k == ord('x'): + self.mtx = None + self.dist = None + self.obj_points_list = [] + self.img_points_list = [] + else: + success, r_vec, t_vec = self._process(img) + if success: + r_vec_obj, _ = to_board_pose_euler(r_vec, t_vec) + rot = [round(180.0 / math.pi * v, 0) for v in r_vec_obj] + if rot != last_rotation: + print("\r ", + end='') + print(f"\rTilt: {rot[0]:.1f}, Roll: {rot[1]:.1f}, Azimuth: {rot[2]:.1f}", end='') + last_rotation = rot + + cv.imshow('img', img) + cv.destroyAllWindows() + + @abc.abstractmethod + def _process(self, image, draw_marker_box=True, draw_marker_id=False, draw_marker_axis=True) -> (bool, np.ndarray, np.ndarray): + pass \ No newline at end of file diff --git a/ar_tag_pose_aruco.py b/ar_tag_pose_aruco.py new file mode 100644 index 0000000..375422f --- /dev/null +++ b/ar_tag_pose_aruco.py @@ -0,0 +1,97 @@ +from PIL.ImageOps import posterize + +from ar_tag_pose import ArTagPose +import numpy as np +import cv2 as cv + +def factory(dict_type: int, board_size: cv.typing.Size, marker_length: float=100.0, marker_sep=10, margins=10): + aruco_dict = cv.aruco.getPredefinedDictionary(dict_type) + params = cv.aruco.DetectorParameters() + detector = cv.aruco.ArucoDetector(aruco_dict, params) + board = cv.aruco.GridBoard(board_size, markerLength=marker_length, markerSeparation=marker_sep, dictionary=aruco_dict) + im_w = int(board_size[0] * (marker_length + marker_sep) - marker_sep + 2 * margins) + im_h = int(board_size[1] * (marker_length + marker_sep) - marker_sep + 2 * margins) + image = board.generateImage((im_w, im_h), None, margins, 1) + cv.imwrite("results/aruco_board.png", image) + + return detector, board + +class ArTagPoseArucoBoard(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.detectMarkers(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) + (corners, ids, rejected) = self.detector.detectMarkers(gray) + + # Draw a square around the markers + if draw_marker_box: + cv.aruco.drawDetectedMarkers(image, corners) + + # Draw marker id + if draw_marker_id: + cv.aruco.drawDetectedMarkers(image, corners, ids) + + # Draw marker axis + if draw_marker_axis: + _r_vecs, _t_vecs, _ = cv.aruco.estimatePoseSingleMarkers(corners, markerLength=0.1, cameraMatrix=self.mtx, distCoeffs=self.dist) + if _r_vecs is not None: + for i in range(len(_r_vecs)): + cv.drawFrameAxes(image, self.mtx, self.dist, _r_vecs[i], _t_vecs[i], 0.05) + + if len(corners) > 0: + # Estimate pose of each marker and return the values r_vec and t_vec---(different from those of camera coefficients) + obj_points, img_points = self.board.matchImagePoints(corners, ids) + pose, r_vec, t_vec = cv.solvePnP(obj_points, img_points, self.mtx, self.dist, rvec=r_vecs, tvec=t_vecs) + if pose: + # Draw Axis + cv.drawFrameAxes(image, self.mtx, self.dist, r_vec, t_vec, self.board.getMarkerLength()*1.5) + + return pose, r_vec, t_vec + +def main(): + pose = ArTagPoseArucoBoard(10, (5,5)) + pose.calibrate_load() + pose.process() + + +if __name__ == '__main__': + main() \ No newline at end of file