refactored

2025-11-28 18:49:48 +01:00
parent a6589d43ac
commit a3d7e013b5
17 changed files with 53 additions and 3 deletions
@@ -0,0 +1,165 @@
+import sys
+import numpy as np
+import cv2 as cv
+import argparse
+from aruco_types import ARUCO_DICT
+import math
+
+from utils import to_board_pose_euler, cap_init
+
+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, margin_length=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, margin_length, 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_init(cap)
+
+	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 = [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)
+
+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, draw_marker_box=True, draw_marker_id=False, draw_marker_axis=False):	# 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 marker box
+		if draw_marker_box:
+			cv.aruco.drawDetectedMarkers(img, marker_corners)
+
+		# Draw marker id
+		if draw_marker_id:
+			cv.aruco.drawDetectedCornersCharuco(img, charuco_corners, charuco_ids)
+
+		# Draw marker axis
+		if draw_marker_axis:
+			_rvecs, _tvecs, _ = cv.aruco.estimatePoseSingleMarkers(marker_corners, markerLength=0.1, 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
+					cv.drawFrameAxes(img, mtx, dist, r_vec, t_vec, 10)
+
+	return r_vec, t_vec
+
+if __name__ == '__main__':
+	main()
+	cv.destroyAllWindows()