refactored
This commit is contained in:
@@ -0,0 +1,78 @@
|
||||
import sys
|
||||
import numpy as np
|
||||
import cv2 as cv
|
||||
import argparse
|
||||
from aruco_types import ARUCO_DICT
|
||||
|
||||
def factory(aruco_type: int):
|
||||
aruco_dict = cv.aruco.getPredefinedDictionary(aruco_type)
|
||||
aruco_params = cv.aruco.DetectorParameters()
|
||||
aruco_detector = cv.aruco.ArucoDetector(aruco_dict, aruco_params)
|
||||
|
||||
return aruco_detector
|
||||
|
||||
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"]))
|
||||
aruco_detector = factory(ARUCO_DICT[args["type"]])
|
||||
|
||||
# load the input image from disk and resize it
|
||||
print("[INFO] start processing...")
|
||||
process_video(aruco_detector)
|
||||
|
||||
def process_video(aruco_detector: cv.aruco.ArucoDetector):
|
||||
with np.load('results/cam.npz') as X:
|
||||
mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]
|
||||
|
||||
# Open camera
|
||||
video = cv.VideoCapture(0)
|
||||
if not video.isOpened():
|
||||
print("Cannot open camera")
|
||||
exit()
|
||||
|
||||
while True:
|
||||
# Read a new frame
|
||||
ok, img = video.read()
|
||||
if not ok:
|
||||
return
|
||||
process(aruco_detector, img, mtx, dist)
|
||||
k = cv.waitKey(1) & 0xFF
|
||||
if k == ord('q'):
|
||||
break
|
||||
|
||||
def process(aruco_detector: cv.aruco.ArucoDetector, img, mtx, dist): # Load previously saved data
|
||||
gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
|
||||
(corners, ids, rejected) = aruco_detector.detectMarkers(gray)
|
||||
|
||||
if len(corners) > 0:
|
||||
# If markers are detected
|
||||
for i in range(0, len(ids)):
|
||||
# Estimate pose of each marker and return the values rvec and tvec---(different from those of camera coefficients)
|
||||
rvec, tvec, _ = cv.aruco.estimatePoseSingleMarkers(corners[i], 0.033, mtx, dist)
|
||||
|
||||
# Draw Axis
|
||||
cv.drawFrameAxes(img, mtx, dist, rvec, tvec, 0.02)
|
||||
|
||||
# Draw a square around the markers
|
||||
cv.aruco.drawDetectedMarkers(img, corners)
|
||||
|
||||
cv.imshow('img', img)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
cv.destroyAllWindows()
|
||||
Reference in New Issue
Block a user