refatored into class

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

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