diff --git a/ar_tag_pose/calibration.py b/ar_tag_pose/calibration.py index 38e3053..10efe73 100644 --- a/ar_tag_pose/calibration.py +++ b/ar_tag_pose/calibration.py @@ -9,27 +9,21 @@ class Calibration: self.obj_points_list = [] self.img_points_list = [] - def calibration_clear(self): + def coef_clear(self): self.mtx = None self.dist = None + + def points_clear(self): self.obj_points_list = [] self.img_points_list = [] def is_calibrated(self): return self.mtx is not None and self.dist is not None - def get(self): - return self.mtx, self.dist - - def set(self, mtx, dist): - self.mtx = mtx - self.dist = dist - def load(self, filename: str): try: with np.load(filename) as X: - mtx, dist = [X[i] for i in ('mtx', 'dist')] - self.set(mtx, dist) + self.mtx, self.dist = [X[i] for i in ('mtx', 'dist')] print(f"Loaded {filename} successfully!") except FileNotFoundError: pass @@ -41,18 +35,42 @@ class Calibration: np.savez(f"{filename}", mtx=self.mtx, dist=self.dist) print(f"Saved {filename} successfully!") + def get_optimal_camera_matrix(self, img: np.ndarray): + mtx_opt = None + if self.is_calibrated(): + h, w = img.shape[:2] + mtx_opt, _ = cv.getOptimalNewCameraMatrix(self.mtx, self.dist, (w, h), 1, (w, h)) + return mtx_opt + + def image_undistort(self, img: np.ndarray): + if not self.is_calibrated(): + return img + mtx_opt = self.get_optimal_camera_matrix(img) + dst = cv.undistort(img, self.mtx, self.dist, None, mtx_opt) + return dst + + def image_remap(self, img: np.ndarray): + if not self.is_calibrated(): + return img + h, w = img.shape[:2] + mtx_opt = self.get_optimal_camera_matrix(img) + map_x, map_y = cv.initUndistortRectifyMap(self.mtx, self.dist, None, mtx_opt, (w, h), 5) + dst = cv.remap(img, map_x, map_y, cv.INTER_LINEAR) + return dst + def process(self, image: np.ndarray, thresh=3.0): criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 3000, 0.001) - error = None + success = False + error = -1 no = len(self.obj_points_list) ni = len(self.img_points_list) if no == ni and no > 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=None, tvecs=None, flags=cv.CALIB_RATIONAL_MODEL, criteria=criteria) - if ret < thresh: + success = True error = ret self.mtx = mtx self.dist = dist - return error + return success, error diff --git a/ar_tag_pose/detector_board_main.py b/ar_tag_pose/detector_board_main.py index ca57596..a6e2bfd 100644 --- a/ar_tag_pose/detector_board_main.py +++ b/ar_tag_pose/detector_board_main.py @@ -14,6 +14,23 @@ from detector_board_charuco import BoardDetectorCharuco from cv2_enumerate_cameras import enumerate_cameras RESULTS_FOLDER = "results" +class ReadKey: + def __init__(self): + self.is_shift = False + + def read(self): + k = cv.pollKey() & 0xFF + res = 0 + if k < 255: + if k == 225 or k == 226: + self.is_shift = True + else: + res = k + if self.is_shift: + res = k - 32 + self.is_shift = False + + return res class Detector(abc.ABC): def __init__(self, name: str, params: dict, source: str, detector: list[BoardDetector], out_file: str|None = None, rotation_order: str=''): @@ -155,6 +172,10 @@ class Detector(abc.ABC): out = None pose_center = True rot = np.array([0,0,0]) + alpha_rot = 0.3 + scale_target_width = 800 + show_undistorted_image = False + read_key = ReadKey() while cap.isOpened(): if self.src_images: cap.set(cv.CAP_PROP_POS_FRAMES, frame_num) @@ -163,59 +184,64 @@ class Detector(abc.ABC): if not ok: break - k = cv.waitKey(1) & 0xFF - if k == ord('q'): + k_ul = read_key.read() + if k_ul == ord('q'): break - elif k == ord('g'): + elif k_ul == ord('g'): frame_go = True frame_num = 0 - elif k == ord('s'): + elif k_ul == ord('p'): frame_go = False frame_num = 0 - elif k == ord('-'): + elif k_ul == ord('-'): frame_go = False frame_num = frame_num - 1 - elif k == ord('+'): + elif k_ul == ord('+'): frame_go = False frame_num = frame_num + 1 - elif k == ord('p'): + elif k_ul == ord('x'): gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) for det in self.detector: success = det.collect_points(gray, self.calibration) if success: print(f"{det}: Points matching success") - elif k == ord('c'): + elif k_ul == ord('c'): gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) - calib_err = self.calibration.process(gray, thresh=3*len(self.detector)) - if calib_err is not None: - print(f"{self.camera_name}: Calibration success: Reprojection error={calib_err}") + ok, calib_err = self.calibration.process(gray, thresh=3*len(self.detector)) + if ok: + print(f"{self.camera_name}: Calibration success: Reprojection error={calib_err:.2f}") self.calibration.save(f"{RESULTS_FOLDER}/{self.camera_name}_cal.npz") else: - print(f"{self.camera_name}: Calibration failed") - elif k == ord('x'): - self.calibration.calibration_clear() - elif k == ord('f'): + print(f"{self.camera_name}: Calibration failed: Reprojection error={calib_err:.2f}") + elif k_ul == ord('X'): + self.calibration.points_clear() + elif k_ul == ord('C'): + self.calibration.coef_clear() + elif k_ul == ord('f'): pose_center = not pose_center + elif k_ul == ord('u'): + show_undistorted_image = not show_undistorted_image - k = 800 / img.shape[1] + k_scale = scale_target_width / img.shape[1] y_pos = 30 - alpha_rot = 0.3 for det in self.detector: success, r_vec, t_vec = det.process(img, self.calibration, True, False, False, pose_center) if success: if len(self.rotation_order) == 3: _rot = to_euler(to_tf(r_vec, t_vec), rotation_order=self.rotation_order) rot = (1-alpha_rot)*rot + alpha_rot*_rot - self.print_angle(img, f"{str(det)}-{self.rotation_order}", rot, (30, int(y_pos/k)), self.rotation_order, 1./k) + self.print_angle(img, f"{str(det)}-{self.rotation_order}", rot, (30, int(y_pos/k_scale)), self.rotation_order, 1./k_scale) 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) rot = (1 - alpha_rot) * rot + alpha_rot * _rot - self.print_angle(img,f"{str(det)}-{rotation_order}", rot, (30, int(y_pos/k)), rotation_order, 1./k) + self.print_angle(img,f"{str(det)}-{rotation_order}", rot, (30, int(y_pos/k_scale)), rotation_order, 1./k_scale) y_pos += 30 - img_scaled = cv.resize(img, dsize=None, fx=k, fy=k) + img_scaled = cv.resize(img, dsize=None, fx=k_scale, fy=k_scale) + if show_undistorted_image: + img_scaled = self.calibration.image_remap(img_scaled) cv.imshow('img', img_scaled) if self.out_file is not None: out = out_writer(out, img_scaled, filename=self.out_file)