136 lines
4.5 KiB
Python
136 lines
4.5 KiB
Python
import sys
|
|
import numpy as np
|
|
import cv2 as cv
|
|
import argparse
|
|
from aruco_types import ARUCO_DICT
|
|
|
|
ENTRY = {'charuco_corners': [], 'charuco_ids': [], 'obj_points': [], 'img_points': []}
|
|
CAM_CALIB_FILE = "results/cam_calib_charuco.npz"
|
|
|
|
def factory(dict_type: int, board_size: cv.typing.Size, square_length: float=5.0, marker_length: float=3.0):
|
|
aruco_dict = cv.aruco.getPredefinedDictionary(dict_type)
|
|
board = cv.aruco.CharucoBoard(board_size, square_length, marker_length, aruco_dict)
|
|
params = cv.aruco.CharucoParameters()
|
|
detector = cv.aruco.CharucoDetector(board, params)
|
|
image = board.generateImage((800, 600), None, 0, 1)
|
|
cv.imwrite("results/charuco_board.png", image)
|
|
|
|
return detector, board
|
|
|
|
def main():
|
|
# construct the argument parser and parse the arguments
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("-t", "--type", type=str,
|
|
default="DICT_ARUCO_ORIGINAL",
|
|
help="type of ArUCo tag to detect")
|
|
args = vars(ap.parse_args())
|
|
|
|
# verify that the supplied ArUCo tag exists and is supported by
|
|
# OpenCV
|
|
if ARUCO_DICT.get(args["type"], None) is None:
|
|
print("[INFO] ArUCo tag of '{}' is not supported".format(
|
|
args["type"]))
|
|
sys.exit(0)
|
|
# load the ArUCo dictionary, grab the ArUCo parameters, and detect
|
|
# the markers
|
|
print("[INFO] detecting '{}' tags...".format(args["type"]))
|
|
detector, board = factory(ARUCO_DICT[args["type"]], (6,7))
|
|
|
|
# load the input image from disk and resize it
|
|
print("[INFO] start processing...")
|
|
process_video(detector, board)
|
|
|
|
def process_video(detector: cv.aruco.CharucoDetector, board):
|
|
mtx = None
|
|
dist = None
|
|
rvecs = None
|
|
tvecs = None
|
|
|
|
try:
|
|
with np.load(CAM_CALIB_FILE) as X:
|
|
mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
# Open camera
|
|
video = cv.VideoCapture(0)
|
|
if not video.isOpened():
|
|
print("Cannot open camera")
|
|
exit()
|
|
|
|
img_size = None
|
|
while True:
|
|
# Read a new frame
|
|
ok, img = video.read()
|
|
if not ok:
|
|
return
|
|
k = cv.waitKey(1) & 0xFF
|
|
if k == ord('q'):
|
|
break
|
|
elif k == ord('s'):
|
|
img_size = store_frame(detector, board, img)
|
|
elif k == ord('c'):
|
|
mtx, dist, _, _ = calibrate_camera(img_size, mtx, dist, None, None)
|
|
else:
|
|
process(detector, board, img, mtx, dist, rvecs, tvecs)
|
|
|
|
def store_frame(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img):
|
|
global ENTRY
|
|
img_size = 0
|
|
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
|
|
charuco_corners, charuco_ids, _, _ = detector.detectBoard(gray)
|
|
if charuco_corners is not None:
|
|
if len(charuco_ids) >= len(board.getIds()):
|
|
obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids)
|
|
if len(obj_points) > 0 and len(img_points) > 0:
|
|
ENTRY['charuco_corners'].append(charuco_corners)
|
|
ENTRY['charuco_ids'].append(charuco_ids)
|
|
ENTRY['obj_points'].append(obj_points)
|
|
ENTRY['img_points'].append(img_points)
|
|
img_size = gray.shape
|
|
print("Frame captured")
|
|
else:
|
|
print("Point matching failed, try again")
|
|
else:
|
|
print("Point matching failed, try again")
|
|
else:
|
|
print("Point matching failed, try again")
|
|
|
|
return img_size
|
|
|
|
def calibrate_camera(img_size, mtx, dist, rvecs=None, tvecs=None):
|
|
global ENTRY
|
|
if len(ENTRY['obj_points']) > 0 and len(ENTRY['img_points']) > 0:
|
|
ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(ENTRY['obj_points'], ENTRY['img_points'], img_size, mtx, dist, rvecs=rvecs, tvecs=tvecs)
|
|
np.savez(f"{CAM_CALIB_FILE}", mtx=mtx, dist=dist, rvecs=rvecs, tvecs=tvecs)
|
|
ENTRY = {'charuco_corners': [], 'charuco_ids': [], 'obj_points': [], 'img_points': []}
|
|
print(f"Camera calibration finished. Result = {ret}")
|
|
return mtx, dist, rvecs, tvecs
|
|
|
|
|
|
def process(detector: cv.aruco.CharucoDetector, board: cv.aruco.Board, img, mtx, dist, rvecs=None, tvecs=None): # Load previously saved data
|
|
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
|
|
charuco_corners, charuco_ids, marker_corners, marker_ids = detector.detectBoard(gray)
|
|
if charuco_corners is not None:
|
|
|
|
# Draw a square around the markers
|
|
cv.aruco.drawDetectedCornersCharuco(img, charuco_corners, charuco_ids)
|
|
|
|
# Draw a square around the markers
|
|
cv.aruco.drawDetectedMarkers(img, marker_corners)
|
|
|
|
if len(charuco_ids) >= len(board.getIds()):
|
|
obj_points, img_points = board.matchImagePoints(charuco_corners, charuco_ids)
|
|
if mtx is not None and dist is not None:
|
|
pose, rvec, tvec = cv.solvePnP(obj_points, img_points, mtx, dist, rvec=rvecs, tvec=tvecs)
|
|
if pose:
|
|
# Draw Axis
|
|
cv.drawFrameAxes(img, mtx, dist, rvec, tvec, 10)
|
|
|
|
cv.imshow('img', img)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|
|
cv.destroyAllWindows()
|