<reponame>ftPeter/pose-estimation<filename>charuco-calibration/charuco_pose.py
import pdb
import time
import numpy as np
import cv2
import glob

import cv2.aruco as aruco

cap = cv2.VideoCapture(0)

blob_params = cv2.SimpleBlobDetector_Params()
blob_params.minArea = 80
blob_params.maxArea = 200
blob_params.minDistBetweenBlobs = 8
blob_detector = cv2.SimpleBlobDetector_create(blob_params)

def read_node_real( reader, name ):
    node = reader.getNode( name )
    return node.real()

def read_node_string( reader, name ):
    node = reader.getNode( name )
    return node.string()

def read_node_matrix( reader, name ):
    node = reader.getNode( name )
    return node.mat()

# read defaultConfig.xml presented to
# opencv_interactive-calibration
config_reader = cv2.FileStorage()
config_reader.open("defaultConfig.xml",cv2.FileStorage_READ)

aruco_parameters =  aruco.DetectorParameters_create()
#aruco_dict = aruco.Dictionary_get(aruco.DICT_6X6_250)
aruco_dict_num = int(read_node_real( config_reader, "charuco_dict" ) )
aruco_dict = aruco.Dictionary_get(aruco_dict_num)
charuco_square_length = int(read_node_real( config_reader, "charuco_square_lenght" ))
charuco_marker_size   = int(read_node_real( config_reader, "charuco_marker_size" ))
config_reader.release()

# TODO clean this up and use config parameters?
#      look up definition o charuco_square_lenght, charuco_marker_size
# 5 was width (x) passed to interactive-calibration
# 7 was height (y) passed to interactive-calibration
# .034 is square_length
# .02 is marker_size
charuco_board = aruco.CharucoBoard_create( 5, 7, 0.034, 0.02, aruco_dict )


camera_matrix = np.array([[ 407.4104068,0.0,663.22920611],[ 0.0,395.58125775,296.96472213],[ 0.0,0.0,1.0]],np.float32 )
dist_coeffs   = np.array([[ 0.35424759,-0.60238148,0.03750194,-0.01836975,0.38208264]],np.float32 )

# read the cameraParameters.xml file generated by
# opencv_interactive-calibration
camera_reader = cv2.FileStorage()
camera_reader.open("cameraParameters.xml",cv2.FileStorage_READ)

camera_matrix = read_node_matrix( camera_reader, "cameraMatrix" )
dist_coeffs   = read_node_matrix( camera_reader, "dist_coeffs" )

camera_reader.release()


all_ch_corners = list()
all_ch_ids     = list()

while(True):
    time.sleep( 0.1 )
    # Read frame from Camera
    # convert frame to grayscale
    ret, frame = cap.read()
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)


    corners, ids, rejectedImgPoints = aruco.detectMarkers(gray, aruco_dict, 
        parameters=aruco_parameters)

    # if enough markers were detected
    # then process the board
    if( ids is not None ):
        #gray   = aruco.drawDetectedMarkers(gray, corners)
        ret, ch_corners, ch_ids = aruco.interpolateCornersCharuco(corners, ids, gray, charuco_board )
        
        # if there are enough corners to get a reasonable result
        if( ret > 5 ):
            aruco.drawDetectedCornersCharuco(frame,ch_corners,ch_ids,(0,0,255) )
            retval, rvec, tvec = aruco.estimatePoseCharucoBoard( ch_corners, ch_ids, charuco_board, camera_matrix, dist_coeffs )

            # if a pose could be estimated
            if( retval ) :
                frame = aruco.drawAxis(frame,camera_matrix,dist_coeffs,rvec,tvec,0.032)

# imshow and waitKey are required for the window
    # to open on a mac.
    cv2.imshow('frame', frame)

    if( cv2.waitKey(1) & 0xFF == ord('q') ):
        break

cap.release()
cv2.destroyAllWindows()

