added charuco board calibration

cam_pose_charuco_board.py

@@ -4,6 +4,9 @@ import cv2 as cv
 import argparse
 from aruco_types import ARUCO_DICT
 
+ENTRY = {'charuco_corners': [], 'charuco_ids': [], 'obj_points': [], 'img_points': []}
+CAM_CALIB_FILE = "results/cam_calib_charuco.npz"
+
 def factory(dict_type: int, board_size: cv.typing.Size, square_length: float=5.0, marker_length: float=3.0):
 	aruco_dict = cv.aruco.getPredefinedDictionary(dict_type)
 	board = cv.aruco.CharucoBoard(board_size, square_length, marker_length, aruco_dict)
@@ -31,15 +34,19 @@ def main():
 	# load the ArUCo dictionary, grab the ArUCo parameters, and detect
 	# the markers
 	print("[INFO] detecting '{}' tags...".format(args["type"]))
-	detector, board = factory(ARUCO_DICT[args["type"]], (3,3))
+	detector, board = factory(ARUCO_DICT[args["type"]], (6,7))
 
 	# load the input image from disk and resize it
 	print("[INFO] start processing...")
 	process_video(detector, board)
 
 def process_video(detector: cv.aruco.CharucoDetector, board):
-	with np.load('results/cam.npz') as X:
-		mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]
+	mtx = None
+	dist = None
+	rvecs = None
+	tvecs = None
+	with np.load(CAM_CALIB_FILE) as X:
+		mtx, dist, rvecs, tvecs = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]
 
 	# Open camera
 	video = cv.VideoCapture(0)
@@ -47,17 +54,56 @@ def process_video(detector: cv.aruco.CharucoDetector, board):
 		print("Cannot open camera")
 		exit()
 
+	img_size = None
 	while True:
 		# Read a new frame
 		ok, img = video.read()
 		if not ok:
 			return
-		process(detector, board, img, mtx, dist)
 		k = cv.waitKey(1) & 0xFF
 		if k == ord('q'):
 			break
+		elif k == ord('s'):
+			img_size = store_frame(detector, board, img)
+		elif k == ord('c'):
+			mtx, dist, rvecs, tvecs = calibrate_camera(mtx, dist, rvecs, tvecs, img_size)
+		else:
+			process(detector, board, img, mtx, dist, rvecs, tvecs)
 
-def process(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img, mtx, dist):	# Load previously saved data
+def store_frame(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img):
+	global ENTRY
+	img_size = 0
+	gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
+	charuco_corners, charuco_ids, _, _ = detector.detectBoard(gray)
+	if charuco_corners is not None:
+		if len(charuco_ids) >= len(board.getIds()):
+			obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids)
+
+			if len(obj_points) > 0 and len(img_points) > 0:
+				ENTRY['charuco_corners'].append(charuco_corners)
+				ENTRY['charuco_ids'].append(charuco_ids)
+				ENTRY['obj_points'].append(obj_points)
+				ENTRY['img_points'].append(img_points)
+				img_size = gray.shape
+				print("Frame captured")
+			else:
+				print("Point matching failed, try again")
+		else:
+			print("Point matching failed, try again")
+	else:
+		print("Point matching failed, try again")
+
+	return img_size
+
+def calibrate_camera(mtx, dist, rvecs, tvecs, img_size):
+	global ENTRY
+	ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(ENTRY['obj_points'], ENTRY['img_points'], img_size, mtx, dist, rvecs, tvecs)
+	np.savez(f"{CAM_CALIB_FILE}", mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)
+	ENTRY = {'charuco_corners': [], 'charuco_ids': [], 'obj_points': [], 'img_points': []}
+	return mtx, dist, rvecs, tvecs
+
+
+def process(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img, mtx, dist, rvecs, tvecs):	# Load previously saved data
 	gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
 	charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray)
 	if charuco_corners is not None:
@@ -70,10 +116,12 @@ def process(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img, mtx,
 
 		if len(charuco_ids) >= len(board.getIds()):
 			obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids)
-			pose, rvec, tvec = cv.solvePnP(obj_points, img_points, mtx, dist, None, None)
-			if pose:
-				# Draw Axis
-				cv.drawFrameAxes(img, mtx, dist, rvec, tvec, 20)
+
+			if mtx is not None and dist is not None:
+				pose, rvec, tvec = cv.solvePnP(obj_points, img_points, mtx, dist, rvecs, tvecs)
+				if pose:
+					# Draw Axis
+					cv.drawFrameAxes(img, mtx, dist, rvec, tvec, 20)
 
 	cv.imshow('img', img)