ocv_cam_pose/cam_pose_charuco_board.py

import sys
import numpy as np
import cv2 as cv
import argparse
from aruco_types import ARUCO_DICT
import math

CAM_CALIB_FILE = "results/cam_calib_charuco.npz"

def rotation_matrix_to_euler_angles(R):
	sy = math.sqrt(R[0,0]**2 + R[1,0]**2)
	singular = sy < 1e-6

	if not singular:
		x = math.atan2(R[2,1] , R[2,2])
		y = math.atan2(-R[2,0], sy)
		z = math.atan2(R[1,0], R[0,0])
	else:
		# Handle singularity
		x = math.atan2(-R[1,2], R[1,1])
		y = math.atan2(-R[2,0], sy)
		z = 0
	return np.array([x, y, z])

def to_board_pose_euler(r_vec, t_vec):
	r_mat = np.ndarray((3, 3))
	cv.Rodrigues(r_vec, r_mat)
	r_mat = np.hstack((r_mat, t_vec))
	r_mat = np.vstack((r_mat, np.array([0, 0, 0, 1])))
	r_mat = np.linalg.inv(r_mat)
	res = rotation_matrix_to_euler_angles(r_mat)
	return res, r_mat[:,3]

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)
	params = cv.aruco.CharucoParameters()
	detector = cv.aruco.CharucoDetector(board, params)
	image = board.generateImage((800, 600), None, 0, 1)
	cv.imwrite("results/charuco_board.png", image)

	return detector, board

def main():
	# construct the argument parser and parse the arguments
	ap = argparse.ArgumentParser()
	ap.add_argument("-t", "--type", type=str,
		default="DICT_ARUCO_ORIGINAL",
		help="type of ArUCo tag to detect")
	args = vars(ap.parse_args())

	# verify that the supplied ArUCo tag exists and is supported by
	# OpenCV
	if ARUCO_DICT.get(args["type"], None) is None:
		print("[INFO] ArUCo tag of '{}' is not supported".format(
			args["type"]))
		sys.exit(0)
	# 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"]], (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: cv.aruco.CharucoBoard):
	mtx = None
	dist = None
	obj_points_list = []
	img_points_list = []
	try:
		with np.load(CAM_CALIB_FILE) as X:
			mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]
	except FileNotFoundError:
		pass
	except KeyError:
		pass

	# Open camera
	cap = cv.VideoCapture(0)

	if not cap.isOpened():
		print("Cannot open camera")
		exit()

	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)

	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'):
			img_size, obj_points, img_points = store_frame(detector, board, img)
			if img_size is not None:
				obj_points_list.append(obj_points)
				img_points_list.append(img_points)
		elif k == ord('c'):
			mtx, dist, rvecs, tvecs = calibrate_camera(img_size, obj_points_list, img_points_list, mtx, dist, None, None)
			np.savez(f"{CAM_CALIB_FILE}", mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)
		elif k == ord('x'):
			mtx = None
			dist = None
			obj_points_list = []
			img_points_list = []
		else:
			r_vec, t_vec = process(detector, board, img, mtx, dist, None, None)
			if r_vec is not None and t_vec is not None:
				r_vec_obj, _ = to_board_pose_euler(r_vec, t_vec)
				rot = [int(180.0 / math.pi * v) for v in r_vec_obj]
				if rot != last_rotation:
					print(f"\rRotation Tilt: {rot[0]:.1f}, Roll: {rot[1]:.1f}, Azimuth: {rot[2]:.1f}", end='')
				last_rotation = rot

		cv.imshow('img', img)

def store_frame(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img):
	img_size = None
	obj_points = None
	img_points = None
	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:
				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, obj_points, img_points

def calibrate_camera(img_size, obj_points_list, img_points_list, mtx, dist, rvecs=None, tvecs=None, flags=0):
	if len(obj_points_list) > 0 and len(img_points_list) > 0:
		ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(obj_points_list, img_points_list, img_size, mtx, dist, rvecs=rvecs, tvecs=tvecs, flags=flags)
		print(f"Camera calibration finished. Result = {ret}")
	return mtx, dist, rvecs, tvecs


def process(detector: cv.aruco.CharucoDetector, board: cv.aruco.CharucoBoard, img, mtx, dist, rvecs=None, tvecs=None):	# Load previously saved data
	r_vec = None
	t_vec = None
	gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
	charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray)
	if charuco_corners is not None:

		# Draw a square around the markers
		cv.aruco.drawDetectedCornersCharuco(img, charuco_corners, charuco_ids)

		# Draw a square around the markers
		cv.aruco.drawDetectedMarkers(img, marker_corners)

		# Draw an axis around the markers
		rvecs, tvecs, _ = cv.aruco.estimatePoseSingleMarkers(marker_corners, markerLength=0.062, cameraMatrix=mtx, distCoeffs=dist)
		for i in range(len(rvecs)):
			cv.drawFrameAxes(img, mtx, dist, rvecs[i], tvecs[i],0.05)

		if len(charuco_ids) >= len(board.getIds()):
			obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids)
			offset = int(board.getMarkerLength() * board.getSquareLength())
			obj_points -= [offset, offset, 0]
			if mtx is not None and dist is not None:
				pose, r_vec, t_vec = cv.solvePnP(obj_points, img_points, mtx, dist, useExtrinsicGuess=False, rvec=rvecs, tvec=tvecs, flags=cv.SOLVEPNP_ITERATIVE)
				if pose:
					# Draw Axis
					cv.drawFrameAxes(img, mtx, dist, r_vec, t_vec, 10)

	return r_vec, t_vec

if __name__ == '__main__':
	main()
	cv.destroyAllWindows()