diff --git a/ar_tag_pose/board_pose.py b/ar_tag_pose/board_pose.py index 1ac512b..f235b6c 100644 --- a/ar_tag_pose/board_pose.py +++ b/ar_tag_pose/board_pose.py @@ -1,9 +1,11 @@ +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" @@ -18,8 +20,24 @@ class BoardPose(abc.ABC): obj_points_list = [] img_points_list = [] - def __init__(self, device_id: int): - self.device_id = device_id + def __init__(self, source: str): + self.device_id = None + self.src_video = None + self.src_images = None + if source.isnumeric(): + self.device_id = int(source) + elif isinstance(source, str): + 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 @@ -62,7 +80,7 @@ class BoardPose(abc.ABC): rvecs=r_vecs, tvecs=t_vecs, flags=flags) print(f"Camera calibration finished: ", end='') - if ret < 2: + if ret < 3: print(f"Success ({ret})!") success = True self.mtx = mtx @@ -75,29 +93,52 @@ class BoardPose(abc.ABC): return success def process(self): + cap = None + if self.device_id is not None: + # Open camera + cap = cv.VideoCapture(self.device_id) + BoardPose.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) - # Open camera - cap = cv.VideoCapture(self.device_id) - BoardPose.cap_init(cap) - - if not cap.isOpened(): - print("Cannot open camera") + if cap is None or not cap.isOpened(): + print("Cannot open media") exit() cal_img_size = None - last_rotation = [0, 0, 0] - while True: + 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: - return + 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'): success, cal_img_size, obj_points, img_points = self.calibrate_points(img) if success: - print("Frame captured") + print("Points captured") self.obj_points_list.append(obj_points) self.img_points_list.append(img_points) else: @@ -111,20 +152,21 @@ class BoardPose(abc.ABC): 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 + + 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, 1) for v in r_vec_obj] + print("\r ", end='') + print(f"\rTilt: {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) - cv.destroyAllWindows() + if frame_go: + frame_num += 1 + + cv.destroyAllWindows() @abc.abstractmethod def get_name(self) -> str: diff --git a/ar_tag_pose/board_pose_aruco.py b/ar_tag_pose/board_pose_aruco.py index 5564f88..8bce89f 100644 --- a/ar_tag_pose/board_pose_aruco.py +++ b/ar_tag_pose/board_pose_aruco.py @@ -15,7 +15,7 @@ def factory(dict_type: int, board_size: cv.typing.Size, marker_length: float=100 return detector, board, name class BoardPoseAruco(BoardPose): - def __init__(self, device_id: int, dict_type: int, board_size: tuple[int, int]): + def __init__(self, device_id: str, dict_type: int, board_size: tuple[int, int]): BoardPose.__init__(self, device_id) self.board_size = board_size self.detector, self.board, self.name = factory(dict_type, board_size) @@ -53,8 +53,6 @@ class BoardPoseAruco(BoardPose): 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) @@ -67,6 +65,9 @@ class BoardPoseAruco(BoardPose): 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) @@ -87,12 +88,12 @@ class BoardPoseAruco(BoardPose): def main(): # construct the argument parser and parse the arguments ap = argparse.ArgumentParser() - ap.add_argument("-i", "--device_id", type=int, - default=0, - help="Camera device index ") + ap.add_argument("-s", "--source", type=str, + default='0', + help="Camera device index, movie file or first file of image sequence") args = vars(ap.parse_args()) - pose = BoardPoseAruco(args["device_id"], 10, (5, 5)) + pose = BoardPoseAruco(args["source"], 10, (5, 5)) pose.board_image_generate(margin_length=10) pose.calibrate_load() pose.process() diff --git a/ar_tag_pose/board_pose_charuco.py b/ar_tag_pose/board_pose_charuco.py index 2aa685c..9c354e5 100644 --- a/ar_tag_pose/board_pose_charuco.py +++ b/ar_tag_pose/board_pose_charuco.py @@ -15,7 +15,7 @@ def factory(dict_type: int, board_size: cv.typing.Size, square_length: float = 5 return detector, board, name class BoardPoseCharuco(BoardPose): - def __init__(self, device_id: int, dict_type: int, board_size: tuple[int, int]): + def __init__(self, device_id: str, dict_type: int, board_size: tuple[int, int]): BoardPose.__init__(self, device_id) self.board_size = board_size self.detector, self.board, self.name = factory(dict_type, board_size) @@ -85,13 +85,13 @@ class BoardPoseCharuco(BoardPose): def main(): # construct the argument parser and parse the arguments ap = argparse.ArgumentParser() - ap.add_argument("-i", "--device_id", type=int, - default=0, - help="Camera device index ") + ap.add_argument("-s", "--source", type=str, + default='0', + help="Camera device index, movie file or first file of image sequence") args = vars(ap.parse_args()) - pose = BoardPoseCharuco(args["device_id"], 10, (6, 7)) - pose.board_image_generate() + pose = BoardPoseCharuco(args["source"], 10, (6, 7)) + pose.board_image_generate(margin_length=10) pose.calibrate_load() pose.process() diff --git a/checker_cam.jpg b/images/checker_cam.jpg similarity index 100% rename from checker_cam.jpg rename to images/checker_cam.jpg diff --git a/images/cube_0001.png b/images/cube_0001.png new file mode 100644 index 0000000..6fc5210 Binary files /dev/null and b/images/cube_0001.png differ diff --git a/images/cube_0002.png b/images/cube_0002.png new file mode 100644 index 0000000..1370c91 Binary files /dev/null and b/images/cube_0002.png differ diff --git a/images/cube_0003.png b/images/cube_0003.png new file mode 100644 index 0000000..2abca8d Binary files /dev/null and b/images/cube_0003.png differ diff --git a/images/cube_0004.png b/images/cube_0004.png new file mode 100644 index 0000000..6fa887c Binary files /dev/null and b/images/cube_0004.png differ diff --git a/images/cube_0005.png b/images/cube_0005.png new file mode 100644 index 0000000..1ac5b34 Binary files /dev/null and b/images/cube_0005.png differ diff --git a/images/cube_0006.png b/images/cube_0006.png new file mode 100644 index 0000000..b380d49 Binary files /dev/null and b/images/cube_0006.png differ diff --git a/images/cube_0007.png b/images/cube_0007.png new file mode 100644 index 0000000..b8c43cf Binary files /dev/null and b/images/cube_0007.png differ diff --git a/images/cube_0008.png b/images/cube_0008.png new file mode 100644 index 0000000..e454fca Binary files /dev/null and b/images/cube_0008.png differ diff --git a/images/cube_0009.png b/images/cube_0009.png new file mode 100644 index 0000000..26ab59a Binary files /dev/null and b/images/cube_0009.png differ diff --git a/images/cube_0010.png b/images/cube_0010.png new file mode 100644 index 0000000..fe869a3 Binary files /dev/null and b/images/cube_0010.png differ diff --git a/images/cube_0011.png b/images/cube_0011.png new file mode 100644 index 0000000..895a1c4 Binary files /dev/null and b/images/cube_0011.png differ diff --git a/images/cube_0012.png b/images/cube_0012.png new file mode 100644 index 0000000..136bb17 Binary files /dev/null and b/images/cube_0012.png differ diff --git a/images/cube_0013.png b/images/cube_0013.png new file mode 100644 index 0000000..645b4d4 Binary files /dev/null and b/images/cube_0013.png differ diff --git a/images/cube_0014.png b/images/cube_0014.png new file mode 100644 index 0000000..d8bf665 Binary files /dev/null and b/images/cube_0014.png differ diff --git a/images/cube_0015.png b/images/cube_0015.png new file mode 100644 index 0000000..7691b32 Binary files /dev/null and b/images/cube_0015.png differ diff --git a/images/cube_0016.png b/images/cube_0016.png new file mode 100644 index 0000000..1c5ffec Binary files /dev/null and b/images/cube_0016.png differ diff --git a/images/cube_0017.png b/images/cube_0017.png new file mode 100644 index 0000000..ad7a1eb Binary files /dev/null and b/images/cube_0017.png differ diff --git a/images/cube_0018.png b/images/cube_0018.png new file mode 100644 index 0000000..ead4d68 Binary files /dev/null and b/images/cube_0018.png differ diff --git a/images/cube_0019.png b/images/cube_0019.png new file mode 100644 index 0000000..3d09967 Binary files /dev/null and b/images/cube_0019.png differ diff --git a/images/cube_0020.png b/images/cube_0020.png new file mode 100644 index 0000000..7eee160 Binary files /dev/null and b/images/cube_0020.png differ diff --git a/images/cube_0021.png b/images/cube_0021.png new file mode 100644 index 0000000..45b8227 Binary files /dev/null and b/images/cube_0021.png differ diff --git a/images/cube_0022.png b/images/cube_0022.png new file mode 100644 index 0000000..c06ca95 Binary files /dev/null and b/images/cube_0022.png differ diff --git a/images/cube_0023.png b/images/cube_0023.png new file mode 100644 index 0000000..f70e6be Binary files /dev/null and b/images/cube_0023.png differ diff --git a/images/cube_0024.png b/images/cube_0024.png new file mode 100644 index 0000000..ea8978c Binary files /dev/null and b/images/cube_0024.png differ diff --git a/images/cube_0025.png b/images/cube_0025.png new file mode 100644 index 0000000..a41e5c2 Binary files /dev/null and b/images/cube_0025.png differ diff --git a/images/cube_0026.png b/images/cube_0026.png new file mode 100644 index 0000000..78e7fba Binary files /dev/null and b/images/cube_0026.png differ diff --git a/images/cube_0027.png b/images/cube_0027.png new file mode 100644 index 0000000..e9aade8 Binary files /dev/null and b/images/cube_0027.png differ diff --git a/images/cube_0028.png b/images/cube_0028.png new file mode 100644 index 0000000..077b778 Binary files /dev/null and b/images/cube_0028.png differ diff --git a/images/cube_0029.png b/images/cube_0029.png new file mode 100644 index 0000000..a49ddb1 Binary files /dev/null and b/images/cube_0029.png differ diff --git a/images/cube_0030.png b/images/cube_0030.png new file mode 100644 index 0000000..61f0f57 Binary files /dev/null and b/images/cube_0030.png differ diff --git a/images/cube_0031.png b/images/cube_0031.png new file mode 100644 index 0000000..97cde5b Binary files /dev/null and b/images/cube_0031.png differ diff --git a/images/cube_0032.png b/images/cube_0032.png new file mode 100644 index 0000000..52b56f0 Binary files /dev/null and b/images/cube_0032.png differ diff --git a/images/cube_0033.png b/images/cube_0033.png new file mode 100644 index 0000000..45ea509 Binary files /dev/null and b/images/cube_0033.png differ diff --git a/images/cube_0034.png b/images/cube_0034.png new file mode 100644 index 0000000..8dbb1be Binary files /dev/null and b/images/cube_0034.png differ diff --git a/images/cube_0035.png b/images/cube_0035.png new file mode 100644 index 0000000..d34c234 Binary files /dev/null and b/images/cube_0035.png differ diff --git a/images/cube_0036.png b/images/cube_0036.png new file mode 100644 index 0000000..c8d3bd1 Binary files /dev/null and b/images/cube_0036.png differ diff --git a/images/cube_0037.png b/images/cube_0037.png new file mode 100644 index 0000000..e4f168b Binary files /dev/null and b/images/cube_0037.png differ diff --git a/images/cube_0038.png b/images/cube_0038.png new file mode 100644 index 0000000..6d53bcb Binary files /dev/null and b/images/cube_0038.png differ diff --git a/images/cube_0039.png b/images/cube_0039.png new file mode 100644 index 0000000..49cb7d7 Binary files /dev/null and b/images/cube_0039.png differ diff --git a/images/cube_0040.png b/images/cube_0040.png new file mode 100644 index 0000000..2869a63 Binary files /dev/null and b/images/cube_0040.png differ diff --git a/images/cube_0041.png b/images/cube_0041.png new file mode 100644 index 0000000..7793ee2 Binary files /dev/null and b/images/cube_0041.png differ diff --git a/images/cube_0042.png b/images/cube_0042.png new file mode 100644 index 0000000..d6e7bd1 Binary files /dev/null and b/images/cube_0042.png differ diff --git a/images/cube_0043.png b/images/cube_0043.png new file mode 100644 index 0000000..263eeee Binary files /dev/null and b/images/cube_0043.png differ diff --git a/images/cube_0044.png b/images/cube_0044.png new file mode 100644 index 0000000..6014ed6 Binary files /dev/null and b/images/cube_0044.png differ diff --git a/images/cube_0045.png b/images/cube_0045.png new file mode 100644 index 0000000..4d34ad1 Binary files /dev/null and b/images/cube_0045.png differ diff --git a/images/cube_0046.png b/images/cube_0046.png new file mode 100644 index 0000000..5ea30dd Binary files /dev/null and b/images/cube_0046.png differ diff --git a/images/cube_0047.png b/images/cube_0047.png new file mode 100644 index 0000000..2aff593 Binary files /dev/null and b/images/cube_0047.png differ diff --git a/images/cube_0048.png b/images/cube_0048.png new file mode 100644 index 0000000..33c5414 Binary files /dev/null and b/images/cube_0048.png differ diff --git a/images/cube_0049.png b/images/cube_0049.png new file mode 100644 index 0000000..1e1fdd8 Binary files /dev/null and b/images/cube_0049.png differ diff --git a/images/cube_0050.png b/images/cube_0050.png new file mode 100644 index 0000000..cded883 Binary files /dev/null and b/images/cube_0050.png differ diff --git a/images/cube_0051.png b/images/cube_0051.png new file mode 100644 index 0000000..f376f4a Binary files /dev/null and b/images/cube_0051.png differ diff --git a/images/cube_0052.png b/images/cube_0052.png new file mode 100644 index 0000000..9ea41bf Binary files /dev/null and b/images/cube_0052.png differ diff --git a/images/cube_0053.png b/images/cube_0053.png new file mode 100644 index 0000000..ad6c37d Binary files /dev/null and b/images/cube_0053.png differ diff --git a/images/cube_0054.png b/images/cube_0054.png new file mode 100644 index 0000000..ef98852 Binary files /dev/null and b/images/cube_0054.png differ diff --git a/images/cube_0055.png b/images/cube_0055.png new file mode 100644 index 0000000..288a4fa Binary files /dev/null and b/images/cube_0055.png differ diff --git a/images/cube_0056.png b/images/cube_0056.png new file mode 100644 index 0000000..f910661 Binary files /dev/null and b/images/cube_0056.png differ diff --git a/images/cube_0057.png b/images/cube_0057.png new file mode 100644 index 0000000..84f7c7b Binary files /dev/null and b/images/cube_0057.png differ diff --git a/images/cube_0058.png b/images/cube_0058.png new file mode 100644 index 0000000..745e217 Binary files /dev/null and b/images/cube_0058.png differ diff --git a/images/cube_0059.png b/images/cube_0059.png new file mode 100644 index 0000000..7988145 Binary files /dev/null and b/images/cube_0059.png differ diff --git a/images/cube_0060.png b/images/cube_0060.png new file mode 100644 index 0000000..060c7b7 Binary files /dev/null and b/images/cube_0060.png differ diff --git a/images/cube_0061.png b/images/cube_0061.png new file mode 100644 index 0000000..22dc800 Binary files /dev/null and b/images/cube_0061.png differ diff --git a/images/cube_0062.png b/images/cube_0062.png new file mode 100644 index 0000000..cbc6427 Binary files /dev/null and b/images/cube_0062.png differ diff --git a/images/cube_0063.png b/images/cube_0063.png new file mode 100644 index 0000000..b38c31d Binary files /dev/null and b/images/cube_0063.png differ diff --git a/images/cube_0064.png b/images/cube_0064.png new file mode 100644 index 0000000..0a8b4d1 Binary files /dev/null and b/images/cube_0064.png differ diff --git a/images/cube_0065.png b/images/cube_0065.png new file mode 100644 index 0000000..af1be07 Binary files /dev/null and b/images/cube_0065.png differ diff --git a/images/cube_0066.png b/images/cube_0066.png new file mode 100644 index 0000000..e186924 Binary files /dev/null and b/images/cube_0066.png differ diff --git a/images/cube_0067.png b/images/cube_0067.png new file mode 100644 index 0000000..5715163 Binary files /dev/null and b/images/cube_0067.png differ diff --git a/images/cube_0068.png b/images/cube_0068.png new file mode 100644 index 0000000..ee1310b Binary files /dev/null and b/images/cube_0068.png differ diff --git a/images/cube_0069.png b/images/cube_0069.png new file mode 100644 index 0000000..897a344 Binary files /dev/null and b/images/cube_0069.png differ diff --git a/images/cube_0070.png b/images/cube_0070.png new file mode 100644 index 0000000..64fd95b Binary files /dev/null and b/images/cube_0070.png differ diff --git a/images/cube_0071.png b/images/cube_0071.png new file mode 100644 index 0000000..bc9e0d6 Binary files /dev/null and b/images/cube_0071.png differ