1import numpy as np
2import cv2 as cv
3
4# aruco config
5adict = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_50)
6cv.imshow("marker", cv.aruco.drawMarker(adict, 0, 400))
7marker_len = 5
8
9# rapid config
10obj_points = np.float32([[-0.5, 0.5, 0], [0.5, 0.5, 0], [0.5, -0.5, 0], [-0.5, -0.5, 0]]) * marker_len
11tris = np.int32([[0, 2, 1], [0, 3, 2]])  # note CCW order for culling
12line_len = 10
13
14# random calibration data. your mileage may vary.
15imsize = (800, 600)
16K = cv.getDefaultNewCameraMatrix(np.diag([800, 800, 1]), imsize, True)
17
18# video capture
19cap = cv.VideoCapture(0)
20cap.set(cv.CAP_PROP_FRAME_WIDTH, imsize[0])
21cap.set(cv.CAP_PROP_FRAME_HEIGHT, imsize[1])
22
23rot, trans = None, None
24while cv.waitKey(1) != 27:
25    img = cap.read()[1]
26
27    # detection with aruco
28    if rot is None:
29        corners, ids = cv.aruco.detectMarkers(img, adict)[:2]
30
31        if ids is not None:
32            rvecs, tvecs = cv.aruco.estimatePoseSingleMarkers(corners, marker_len, K, None)[:2]
33            rot, trans = rvecs[0].ravel(), tvecs[0].ravel()
34
35    # tracking and refinement with rapid
36    if rot is not None:
37        for i in range(5):  # multiple iterations
38            ratio, rot, trans = cv.rapid.rapid(img, 40, line_len, obj_points, tris, K, rot, trans)[:3]
39            if ratio < 0.8:
40                # bad quality, force re-detect
41                rot, trans = None, None
42                break
43
44    # drawing
45    cv.putText(img, "detecting" if rot is None else "tracking", (0, 20), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255))
46    if rot is not None:
47        cv.drawFrameAxes(img, K, None, rot, trans, marker_len)
48    cv.imshow("tracking", img)
49