- introduced BoardDetector
- added BoardDetectorAruco
This commit is contained in:
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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 <Aruco|Charuco>")
|
||||
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()
|
||||
Reference in New Issue
Block a user