From d0210846359bdac9f2538a0cc9161b75fabc145a Mon Sep 17 00:00:00 2001 From: Jens Ahrensfeld Date: Sun, 30 Nov 2025 11:48:45 +0100 Subject: [PATCH] - introduced BoardDetector - added BoardDetectorAruco --- ar_tag_pose/detector_board.py | 90 +++++++++++++ ar_tag_pose/detector_board_aruco.py | 89 +++++++++++++ ar_tag_pose/detector_board_main.py | 196 ++++++++++++++++++++++++++++ 3 files changed, 375 insertions(+) create mode 100644 ar_tag_pose/detector_board.py create mode 100644 ar_tag_pose/detector_board_aruco.py create mode 100644 ar_tag_pose/detector_board_main.py diff --git a/ar_tag_pose/detector_board.py b/ar_tag_pose/detector_board.py new file mode 100644 index 0000000..93447ed --- /dev/null +++ b/ar_tag_pose/detector_board.py @@ -0,0 +1,90 @@ +import os +import abc +import math +import cv2 as cv +import numpy as np +from pathlib import Path +from ar_tag_pose.utils import to_board_pose_euler +import mimetypes + +RESULTS_FOLDER = "results" + +class BoardDetector(abc.ABC): + detector: cv.aruco.ArucoDetector | cv.aruco.CharucoDetector = None + board: cv.aruco.GridBoard | cv.aruco.CharucoBoard = None + name = None + mtx = None + dist = None + r_vecs = None + t_vecs = None + obj_points_list = [] + img_points_list = [] + + def __repr__(self): + return self.name + + def __init__(self, name): + self.name = name + + @abc.abstractmethod + def get_board_name(self) -> str: + pass + + @abc.abstractmethod + def get_board_image(self, margin_length: int) -> np.ndarray: + pass + + @abc.abstractmethod + def match_points(self, image: cv.typing.MatLike): + pass + + @abc.abstractmethod + def process(self, image: cv.typing.MatLike, draw_marker_box=True, draw_marker_id=False, draw_marker_axis=True) -> (bool, np.ndarray, np.ndarray): + pass + + def is_calibrated(self): + return self.mtx is not None and self.dist is not None + + def collect_points(self, image: cv.typing.MatLike) -> bool: + success, obj_points, img_points = self.match_points(image) + if success: + self.obj_points_list.append(obj_points) + self.img_points_list.append(img_points) + + return success + + def calibration_clear(self): + self.mtx = None + self.dist = None + self.r_vecs = None + self.t_vecs = None + self.obj_points_list = [] + self.img_points_list = [] + + def calibration_set(self, mtx, dist, r_vecs, t_vecs): + self.mtx = mtx + self.dist = dist + self.r_vecs = r_vecs + self.t_vecs = t_vecs + + def calibration_get(self): + return self.mtx, self.dist, self.r_vecs, self.t_vecs + + def calibrate_camera(self, image: cv.typing.MatLike, 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, image.shape, self.mtx, self.dist, + rvecs=r_vecs, tvecs=t_vecs, flags=flags) + + print(f"Camera calibration finished: ", end='') + if ret < 3: + 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 diff --git a/ar_tag_pose/detector_board_aruco.py b/ar_tag_pose/detector_board_aruco.py new file mode 100644 index 0000000..3ae4ce8 --- /dev/null +++ b/ar_tag_pose/detector_board_aruco.py @@ -0,0 +1,89 @@ +import cv2 as cv +import numpy as np +from ar_tag_pose.detector_board import BoardDetector +from ar_tag_pose.aruco_types import ARUCO_DICT + +def factory(dict_type: int, board_size: cv.typing.Size, marker_length: float=100.0, marker_sep=10): + aruco_dict = cv.aruco.getPredefinedDictionary(dict_type) + params = cv.aruco.DetectorParameters() + board = cv.aruco.GridBoard(size=board_size, markerLength=marker_length, markerSeparation=marker_sep, dictionary=aruco_dict) + detector = cv.aruco.ArucoDetector(aruco_dict, params) + return detector, board + +class BoardDetectorAruco(BoardDetector): + dict_type = None + def __init__(self, name: str, dict_type: int, board_size: tuple[int, int], marker_ids: list[int]=None): + BoardDetector.__init__(self, name) + self.dict_type = dict_type + self.board_size = board_size + self.detector, self.board = factory(dict_type, board_size) + + def get_board_name(self): + grid_size = self.board.getGridSize() + keyname = [key for key, val in ARUCO_DICT.items() if val == self.dict_type][0] + return f"{self.name}_w{grid_size[0]}_h{grid_size[1]}_{keyname}" + + def get_board_image(self, margin_length: int=10) -> np.ndarray: + image_size = self._image_size(margin_length) + image = self.board.generateImage(image_size, None, margin_length, 1) + return image + + def match_points(self, image: cv.typing.MatLike): + success = False + obj_points = None + img_points = None + corners, ids, _ = self.detector.detectMarkers(image) + 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 + obj_points = _obj_points + img_points = _img_points + + return success, obj_points, img_points + + def process(self, image: cv.typing.MatLike, r_vecs: np.ndarray=None, t_vecs: np.ndarray=None, draw_marker_box=True, draw_marker_id=False, draw_marker_axis=False): + pose = False + r_vec = None + t_vec = 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) + + if not self.is_calibrated(): + return pose, None, None + + # 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 _image_size(self, margin_length: int): + grid_size = self.board.getGridSize() + marker_length = self.board.getMarkerLength() + marker_sep = self.board.getMarkerSeparation() + w = int(grid_size[0] * (marker_length + marker_sep) - marker_sep + 2 * margin_length) + h = int(grid_size[1] * (marker_length + marker_sep) - marker_sep + 2 * margin_length) + return w, h + diff --git a/ar_tag_pose/detector_board_main.py b/ar_tag_pose/detector_board_main.py new file mode 100644 index 0000000..31b4119 --- /dev/null +++ b/ar_tag_pose/detector_board_main.py @@ -0,0 +1,196 @@ +import os +import abc +import math +import cv2 as cv +import numpy as np +from pathlib import Path + +from ar_tag_pose.detector_board_aruco import BoardDetectorAruco +from ar_tag_pose.utils import to_board_pose_euler +import mimetypes +from detector_board import BoardDetector +import argparse + +RESULTS_FOLDER = "results" + + +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) + + +class Detector(abc.ABC): + def __init__(self, source: str, detector: list[BoardDetector]): + self.detector = detector + self.device_id = None + self.src_video = None + self.src_images = None + if source.isnumeric(): + self.device_id = int(source) + else: + mime_type = mimetypes.guess_type(source)[0] + if "video" in mime_type: + self.src_video = source + if "image" in mime_type: + self.src_images = source + basename = os.path.basename(source) + name = basename.split('_')[0] + numeric_part = basename.split('_')[-1].split('.')[0] + suffix = basename.split('_')[-1].split('.')[1] + self.src_images = f"{os.path.dirname(source)}/{name}_%0{len(numeric_part)}d.{suffix}" + + Path.mkdir(Path(RESULTS_FOLDER), exist_ok=True) + + @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 board_image_generate(self, margin_length: int=0): + for det in self.detector: + image = det.get_board_image(margin_length) + name = det.get_board_name() + cv.imwrite(f"{RESULTS_FOLDER}/{name}.png", image) + + def calibrate_save(self): + for det in self.detector: + filename = f"{RESULTS_FOLDER}/{det.get_board_name()}_cal.npz" + mtx, dist, r_vecs, t_vecs = det.calibration_get() + np.savez(f"{filename}", mtx=mtx, dist=dist, rvecs=r_vecs, tvecs=t_vecs) + + def calibrate_load(self): + for det in self.detector: + filename = f"{RESULTS_FOLDER}/{det.get_board_name()}_cal.npz" + try: + with np.load(filename) as X: + mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')] + det.calibration_set(mtx, dist, None, None) + except FileNotFoundError: + pass + except KeyError: + pass + + def process(self): + cap = None + if self.device_id is not None: + # Open camera + cap = cv.VideoCapture(self.device_id) + cap_init(cap) + elif self.src_video is not None: + # Open video file + cap = cv.VideoCapture(self.src_video) + elif self.src_images is not None: + # Open image file sequence + cap = cv.VideoCapture(self.src_images, cv.CAP_IMAGES) + + if cap is None or not cap.isOpened(): + print("Cannot open media") + exit() + + frame_go = True + frame_num = 0 + while cap.isOpened(): + if self.src_images: + cap.set(cv.CAP_PROP_POS_FRAMES, frame_num) + # Read a new frame + ok, img = cap.read() + if not ok: + break + + k = cv.waitKey(1) & 0xFF + if k == ord('q'): + break + elif k == ord('g'): + frame_go = True + frame_num = 0 + elif k == ord('s'): + frame_go = False + frame_num = 0 + elif k == ord('-'): + frame_go = False + frame_num = frame_num - 1 + elif k == ord('+'): + frame_go = False + frame_num = frame_num + 1 + elif k == ord('p'): + gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) + for det in self.detector: + success = det.collect_points(gray) + if success: + print(f"{det}: Points matching success") + else: + print(f"{det}: Points matching failed, try again") + elif k == ord('c'): + gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) + for det in self.detector: + det.calibrate_camera(gray) + self.calibrate_save() + elif k == ord('x'): + for det in self.detector: + det.calibration_clear() + + for det in self.detector: + success, r_vec, t_vec = det.process(img) + if success: + r_vec_obj, _ = to_board_pose_euler(r_vec, t_vec) + rot = [round(180.0 / math.pi * v, 1) for v in r_vec_obj] + print("\r ", end='') + print(f"\r{det}: Tilt: {rot[0]:.1f}, Roll: {rot[1]:.1f}, Azimuth: {rot[2]:.1f}", end='') + + img_scaled = cv.resize(img, dsize=None, fx=0.5, fy=0.5) + cv.imshow('img', img_scaled) + + if frame_go: + frame_num += 1 + + cv.destroyAllWindows() + +def main(): + # construct the argument parser and parse the arguments + ap = argparse.ArgumentParser() + ap.add_argument("-s", "--source", type=str, + default='0', + help="Camera device index, movie file or first file of image sequence") + ap.add_argument("-t", "--type", type=str, + default='Aruco', + help="Board type ") + ap.add_argument("-d", "--dict_type", type=int, + default=10, + help="Dict type") + ap.add_argument("-x", "--board_width", type=int, + default=5, + help="Board width [number of tiles]") + ap.add_argument("-y", "--board_height", type=int, + default=5, + help="Board height [number of tiles]") + ap.add_argument("-m", "--margin_length", type=int, + default=10, + help="Margin length of board [Pixels]") + args = vars(ap.parse_args()) + + board_detector = None + if "Aruco" in args["type"]: + board_detector = (BoardDetectorAruco + ("Det-0", args["dict_type"], (args["board_width"], args["board_height"]))) + elif "Charuco" in args["type"]: + board_detector = (BoardDetectorAruco + ("Det-0", args["dict_type"], (args["board_width"], args["board_height"]))) + + detector = Detector(args["source"], [board_detector]) + detector.board_image_generate(margin_length=args["margin_length"]) + detector.calibrate_load() + detector.process() + + +if __name__ == '__main__': + main() \ No newline at end of file