pose_estimation.py

'''
Sample Usage:-
python pose_estimation.py --K_Matrix calibration_matrix.npy --D_Coeff distortion_coefficients.npy --type DICT_5X5_100
'''


import numpy as np
import cv2
import sys
from utils import ARUCO_DICT
import argparse
import time
from networktables import NetworkTables

import recognition
sys.path.append("./absPose")

from posePossibilities import poses
from relPoseTrans import RelPoseTrans
from field import Field

ip = "10.51.9.2"
phi = 0

field = Field(tagPose=poses)


def pose_esitmation(frame, aruco_dict_type, matrix_coefficients, distortion_coefficients):

    '''
    frame - Frame from the video stream
Tag not in dict
[ 7.06213432 -0.92753656  0.77495676  1.        ]
    matrix_coefficients - Intrinsic matrix of the calibrated camera
    distortion_coefficients - Distortion coefficients associated with your camera

    return:-
    frame - The frame with the axis drawn on it
    '''
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    cv2.aruco_dict = cv2.aruco.Dictionary_get(aruco_dict_type)
    parameters = cv2.aruco.DetectorParameters_create()


    corners, ids, rejected_img_points = cv2.aruco.detectMarkers(gray, cv2.aruco_dict,parameters=parameters,
        cameraMatrix=matrix_coefficients,
        distCoeff=distortion_coefficients)

    tvecs = []
    # If markers are detected
    if len(corners) > 0:
        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, markerPoints = cv2.aruco.estimatePoseSingleMarkers(corners[i], 0.1524, matrix_coefficients,
                                                                       distortion_coefficients)
            tvecs.append([[tvec[0][0][0], tvec[0][0][2], tvec[0][0][1]], ids[i]])
            # Draw a square around the markers
            cv2.aruco.drawDetectedMarkers(frame, corners) 

            # Draw Axis
            cv2.aruco.drawAxis(frame, matrix_coefficients, distortion_coefficients, rvec, tvec, 0.01)  

    return frame, tvecs

if __name__ == '__main__':
    NetworkTables.initialize(server=ip)
    sd = NetworkTables.getTable("SmartDashboard")
    ap = argparse.ArgumentParser()
    ap.add_argument("-k", "--K_Matrix", required=True, help="Path to calibration matrix (numpy file)")
    ap.add_argument("-d", "--D_Coeff", required=True, help="Path to distortion coefficients (numpy file)")
    ap.add_argument("-t", "--type", type=str, default="DICT_ARUCO_ORIGINAL", help="Type of ArUCo tag to detect")
    args = vars(ap.parse_args())
    # args["type"] = args["type"].upper()
    if ARUCO_DICT.get(args["type"], None) is None:
        print(f"ArUCo tag type '{args['type']}' is not supported")
        sys.exit(0)

    aruco_dict_type = ARUCO_DICT[args["type"]]
    calibration_matrix_path = args["K_Matrix"]
    distortion_coefficients_path = args["D_Coeff"]
    
    k = np.load(calibration_matrix_path)
    d = np.load(distortion_coefficients_path)

    video = cv2.VideoCapture("/dev/video0")
    time.sleep(2.0)

    adj = RelPoseTrans(phi)
    frames = 0
    startTime = time.perf_counter()
    elapsedTime = 0
    coneFrame = 0
    cubeFrame = 0
    while True:
        frames += 1
        ret, frame = video.read()
        undistortedFrame = cv2.undistort(frame, k, d)

        if not ret:
            break

        h, w, _ = frame.shape

        width=1000
        height = int(width*(h/w))
        frame = cv2.resize(frame, (width, height), interpolation=cv2.INTER_CUBIC)
        
        output, tvecs = pose_esitmation(frame, aruco_dict_type, k, d)
       
        
        absPose = [0, 0, 0]
        if len(tvecs) > 0:
            poseMap = [(singId[0], adj.transRelPose([-tvec[1], tvec[0], -tvec[1], 1])) for tvec, singId in tvecs]
            theta = sd.getNumber("angle", 0)
            absPose = field.getAbsPose(poseMap, 0) # maybe add rot
            
            # flipping variables
            absPose[0] += absPose[1]
            absPose[1] = absPose[0] - absPose[1]
            absPose[0] = absPose[0] - absPose[1]
            sd.putNumberArray("pose", absPose)
        '''
        if (cubeFrame == 5):
            cv2.imshow('maskCube', recognition.detectCube(undistortedFrame, absPose))
            cubeFrame = 0
        if (coneFrame == 5):
            cv2.imshow('maskCone', recognition.detectCone(undistortedFrame, absPose))
            coneFrame = 0
        '''
        recognition.detectCube(undistortedFrame, absPose)
        recognition.detectCone(undistortedFrame, absPose)
        coneFrame += 1
        cubeFrame += 1
        #key = cv2.waitKey(1) & 0xFF
        elapsedTime = time.perf_counter() - startTime
        if (elapsedTime >= 1):
            startTime = time.perf_counter()
            print(f'framerate: {frames}        ')
            frames = 0
       # if key == ord('q'):
       #     break
       
    video.release()
    cv2.destroyAllWindows()