track_marker.py 1.6 KB

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