diff --git a/ar_tag_pose/detector_board_main.py b/ar_tag_pose/detector_board_main.py index 8f1054d..f5c9f40 100644 --- a/ar_tag_pose/detector_board_main.py +++ b/ar_tag_pose/detector_board_main.py @@ -25,7 +25,7 @@ def cap_init(cap: cv.VideoCapture): class Detector(abc.ABC): - def __init__(self, source: str, detector: list[BoardDetector], out_file: str|None = None, rotation_order: str="XYZ"): + def __init__(self, source: str, detector: list[BoardDetector], out_file: str|None = None, rotation_order: str=''): self.detector = detector self.out_file = out_file self.rotation_order = rotation_order @@ -59,11 +59,12 @@ class Detector(abc.ABC): cap.set(cv.CAP_PROP_AUTO_EXPOSURE, 1) @staticmethod - def print_angle(img: np.ndarray, title: str, angle: np.array, org: tuple[int, int]): + def print_angle(img: np.ndarray, title: str, angle: np.array, org: tuple[int, int], rot_order: str = 'XYZ'): + indices = {k: n for n, k in enumerate(rot_order)} + colors = [(64, 64, 255), (64, 255, 64), (255, 64, 64)] size = draw_text(img, f"{title}:", org, color=(128, 128, 128)) - size = draw_text(img, f"X: {angle[0]:.1f}", (size[0] + 10, org[1]), color=(64, 64, 255)) - size = draw_text(img, f"Y: {angle[1]:.1f}", (size[0] + 10, org[1]), color=(64, 255, 64)) - draw_text(img, f"Z: {angle[2]:.1f}", (size[0] + 10, org[1]), color=(255, 64, 64)) + for axis, color in zip(list("XYZ"), colors): + size = draw_text(img, f"{axis}: {angle[indices[axis]]:.1f}", (size[0] + 10, org[1]), color=color) def board_image_generate(self, margin_length: int=0): for det in self.detector: @@ -162,10 +163,16 @@ class Detector(abc.ABC): for det in self.detector: success, r_vec, t_vec = det.process(img) if success: - for rot_order in ['XYZ', 'XZY', 'YXZ', 'YZX', 'ZXY', 'ZYX']: - rot = to_euler(to_tf(r_vec, t_vec), rotation_order=rot_order) - self.print_angle(img, f"{str(det)}-{rot_order}", rot, (30, y_pos)) + if len(self.rotation_order) == 3: + rot = to_euler(to_tf(r_vec, t_vec), rotation_order=self.rotation_order) + self.print_angle(img, f"{str(det)}-{self.rotation_order}", rot, (30, y_pos), self.rotation_order) y_pos += 30 + else: + for rotation_order in ['XYZ', 'XZY', 'YXZ', 'YZX', 'ZXY', 'ZYX']: + rot = to_euler(to_tf(r_vec, t_vec), rotation_order=rotation_order) + self.print_angle(img, f"{str(det)}-{rotation_order}", rot, (30, y_pos), rotation_order) + y_pos += 30 + img_scaled = cv.resize(img, dsize=None, fx=0.5, fy=0.5) cv.imshow('img', img_scaled) @@ -230,7 +237,7 @@ def main(): action='store_true', help="Enable video output") ap.add_argument("-r", "--rotation_order", type=str, - default='XYZ', + default='', help="Order of Euler rotations") args = ap.parse_args() var_args = vars(args)