refatored into class
This commit is contained in:
+116
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user