added ar_tag_pose_charuco_board

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()