path: root/src/argaze/ArUcoMarkers/ArUcoCube.py

#!/usr/bin/env python

from dataclasses import dataclass, field
import json
import math
import itertools

from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker

import numpy
import cv2 as cv
import cv2.aruco as aruco

@dataclass
class ArUcoCubeFace():
    """Define cube face pose and marker."""

    translation: numpy.array
    """Position in cube referential."""

    rotation: numpy.array
    """Rotation in cube referential."""

    marker: dict
    """ArUco marker linked to the face """

@dataclass
class ArUcoCube():
    """Define a cube with ArUco markers on each face and estimate its pose."""

    dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary
    """ArUco dictionary of cube markers."""

    marker_size: int = field(init=False)
    """Size of markers in centimeter."""

    edge_size: int = field(init=False)
    """Size of the cube edges in centimeter."""

    faces: dict = field(init=False, default_factory=dict)
    """All named faces of the cube and their ArUco markers."""

    angle_tolerance: float = field(init=False)
    """Angle error tolerance allowed to validate face pose in degree."""

    distance_tolerance: float = field(init=False)
    """Distance error tolerance allowed to validate face pose in centimeter."""

    def __init__(self, configuration_filepath):
        """Define cube from a .json  file."""

        with open(configuration_filepath) as configuration_file:

            # Deserialize .json
            # TODO find a better way
            configuration = json.load(configuration_file)

            # Load dictionary
            self.dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(configuration['dictionary'])

            # Load marker size
            self.marker_size = configuration['marker_size']

            # Load edge size
            self.edge_size = configuration['edge_size']

            # Load faces
            self.faces = {}
            for name, face in configuration['faces'].items():
                marker = ArUcoMarker.ArUcoMarker(self.dictionary, face['marker'], self.marker_size)
                self.faces[name] = ArUcoCubeFace(numpy.array(face['translation']).astype(numpy.float32), numpy.array(face['rotation']).astype(numpy.float32), marker)

            # Load angle tolerance
            self.angle_tolerance = configuration['angle_tolerance']

            # Load distance tolerance
            self.distance_tolerance = configuration['distance_tolerance']

            # Init pose data
            self.__translation = numpy.zeros(3)
            self.__rotation = numpy.zeros(3)
            self.__succeded = False
            self.__validity = 0

            # Process markers ids to speed up further calculations
            self.__identifier_cache = {}
            for name, face in self.faces.items():
                self.__identifier_cache[face.marker.identifier] = name

            # Process each face pose to speed up further calculations
            self.__translation_cache = {}
            for name, face in self.faces.items():
                self.__translation_cache[name] = face.translation * self.edge_size / 2

            # Process each face rotation matrix to speed up further calculations
            self.__rotation_cache = {}
            for name, face in self.faces.items():

                # Create intrinsic rotation matrix
                R = self.__make_rotation_matrix(*face.rotation)

                assert(self.__is_rotation_matrix(R))

                # Store rotation matrix
                self.__rotation_cache[name] = R

            # Process each axis-angle face combination to speed up further calculations
            self.__angle_cache = {}
            for (A_name, A_face), (B_name, B_face) in itertools.combinations(self.faces.items(), 2):

                A = self.__rotation_cache[A_name]
                B = self.__rotation_cache[B_name]

                # Rotation matrix from A face to B face
                AB = B.dot(A.T)

                assert(self.__is_rotation_matrix(AB))

                # Calculate axis-angle representation of AB rotation matrix
                angle = numpy.rad2deg(numpy.arccos((numpy.trace(AB) - 1) / 2))

                try:
                    self.__angle_cache[A_name][B_name] = angle
                except:
                    self.__angle_cache[A_name] = {B_name: angle}

                try:
                    self.__angle_cache[B_name][A_name] = angle
                except:
                    self.__angle_cache[B_name] = {A_name: angle}

            # Process distance between face combination to speed up further calculations
            self.__distance_cache = numpy.linalg.norm(numpy.array([0, 0, self.edge_size/2]) - numpy.array([0, self.edge_size/2, 0]))

    def print_cache(self):
        """Print pre-processed data."""

        print('\nIdentifier cache:')
        for i, name in self.__identifier_cache.items():
            print(f'- {i}: {name}')

        print('\nTranslation cache:')
        for name, item in self.__translation_cache.items():
            print(f'- {name}: {item}')

        print('\nRotation cache:')
        for name, item in self.__rotation_cache.items():
            print(f'- {name}:\n{item}')

        print('\nAngle cache:')
        for A_name, A_angle_cache in self.__angle_cache.items():
            for B_name, angle in A_angle_cache.items():
                print(f'- {A_name}/{B_name}: {angle:3f}')

        print(f'\nDistance cache: {self.__distance_cache}')

    def __make_rotation_matrix(self, x, y, z):

        # Create rotation matrix around x axis
        c = numpy.cos(numpy.deg2rad(x))
        s = numpy.sin(numpy.deg2rad(x))
        Rx = numpy.array([[1, 0, 0], [0, c, -s], [0, s, c]])

        # Create rotation matrix around y axis
        c = numpy.cos(numpy.deg2rad(y))
        s = numpy.sin(numpy.deg2rad(y))
        Ry = numpy.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])

        # Create rotation matrix around z axis
        c = numpy.cos(numpy.deg2rad(z))
        s = numpy.sin(numpy.deg2rad(z))
        Rz = numpy.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])

        # Return intrinsic rotation matrix
        return Rx.dot(Ry.dot(Rz))
    
    def __is_rotation_matrix(self, R):
        """Checks if a matrix is a valid rotation matrix."""

        I = numpy.identity(3, dtype = R.dtype)
        return numpy.linalg.norm(I - numpy.dot(R.T, R)) < 1e-6

    def __normalise_face_pose(self, name, face, F):

        # Transform face rotation into cube rotation vector
        R = self.__rotation_cache[name]
        rvec, _ = cv.Rodrigues(F.dot(R))

        #print(f'{name} rotation vector: {rvec[0][0]:3f} {rvec[1][0]:3f} {rvec[2][0]:3f}')

        # Transform face translation into cube translation vector
        OF = face.translation
        T = self.__translation_cache[name]
        FC = F.dot(R.dot(T))

        tvec = OF + FC

        #print(f'{name} translation vector: {tvec[0]:3f} {tvec[1]:3f} {tvec[2]:3f}')

        return rvec, tvec

    def estimate_pose(self, tracked_markers):

        # Init pose data
        self.__translation = numpy.zeros(3)
        self.__rotation = numpy.zeros(3)
        self.__succeded = False
        self.__validity = 0

        # Don't try to estimate pose if there is no tracked markers
        if len(tracked_markers) == 0:

            return self.get_pose()

        # Look for faces related to tracked markers
        tracked_faces = {}
        for (marker_id, marker) in tracked_markers.items():

            try:
                name = self.__identifier_cache[marker_id]
                tracked_faces[name] = marker

            except KeyError:
                continue

        #print('-------------- ArUcoCube pose estimation --------------')

        # Pose validity checking is'nt possible when only one face of the cube is tracked
        if len(tracked_faces.keys()) == 1:

            # Get arcube pose from to the unique face pose
            name, face = tracked_faces.popitem()
            F, _ = cv.Rodrigues(face.rotation)

            self.__rotation, self.__translation = self.__normalise_face_pose(name,face, F)
            self.__succeded = True
            self.__validity = 1

            #print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
            #print(f'arcube rotation vector: {self.__rotation[0][0]:3f} {self.__rotation[1][0]:3f} {self.__rotation[2][0]:3f}')
            #print(f'arcube translation vector: {self.__translation[0]:3f} {self.__translation[1]:3f} {self.__translation[2]:3f}')

        # Pose validity checking processes faces two by two
        else:

            valid_faces = []
            valid_rvecs = []
            valid_tvecs = []

            for (A_name, A_face), (B_name, B_face) in itertools.combinations(tracked_faces.items(), 2):

                #print(f'** {A_name} > {B_name}')

                # Get face rotation estimation
                # Use rotation matrix instead of rotation vector
                A, _ = cv.Rodrigues(A_face.rotation)
                B, _ = cv.Rodrigues(B_face.rotation)

                # Rotation matrix from A face to B face
                AB = B.dot(A.T)

                assert(self.__is_rotation_matrix(AB))

                # Calculate axis-angles representation of AB rotation matrix
                angle = numpy.rad2deg(numpy.arccos((numpy.trace(AB) - 1) / 2))

                #print('rotation angle:')
                #print(angle)

                expected_angle = self.__angle_cache[A_name][B_name]

                #print('expected angle:')
                #print(expected_angle)

                # Calculate distance between A face center and B face center
                distance = numpy.linalg.norm(A_face.translation - B_face.translation)
                expected_distance = self.__distance_cache

                # Check angle and distance according given tolerance then normalise face pose
                valid_angle = math.isclose(angle, expected_angle, abs_tol=self.angle_tolerance)
                valid_distance = math.isclose(distance, expected_distance, abs_tol=self.distance_tolerance)

                if valid_angle and valid_distance:

                    if A_name not in valid_faces:

                        # Remember this face is already validated
                        valid_faces.append(A_name)

                        rvec, tvec = self.__normalise_face_pose(A_name, A_face, A)

                        # Store normalised face pose
                        valid_rvecs.append(rvec)
                        valid_tvecs.append(tvec)
                        
                    if B_name not in valid_faces:

                        # Remember this face is already validated
                        valid_faces.append(B_name)

                        rvec, tvec = self.__normalise_face_pose(B_name, B_face, B)

                        # Store normalised face pose
                        valid_rvecs.append(rvec)
                        valid_tvecs.append(tvec)

            if len(valid_faces) > 1:

                # Consider arcube rotation as the mean of all valid translations
                # !!! WARNING !!! This is a bad hack : processing rotations average is a very complex problem that needs to well define the distance calculation method before.
                self.__rotation = numpy.mean(numpy.array(valid_rvecs), axis=0)

                # Consider arcube translation as the mean of all valid translations
                self.__translation = numpy.mean(numpy.array(valid_tvecs), axis=0)

                #print(':::::::::::::::::::::::::::::::::::::::::::::::::::')
                #print(f'arcube rotation vector: {self.__rotation[0][0]:3f} {self.__rotation[1][0]:3f} {self.__rotation[2][0]:3f}')
                #print(f'arcube translation vector: {self.__translation[0]:3f} {self.__translation[1]:3f} {self.__translation[2]:3f}')

                self.__succeded = True
                self.__validity = len(valid_faces)

        #print('----------------------------------------------------')

        return self.get_pose()

    def get_pose(self):

        return self.__translation, self.__rotation, self.__succeded, self.__validity

    def set_pose(self, tvec = numpy.array([]), rvec = numpy.array([])):

        if tvec.size == 3:
            self.__translation = tvec

        if rvec.size == 3:
            self.__rotation = rvec

        self.__succeded = True
        self.__validity = 0

    def draw(self, frame, K, D):

        l = self.edge_size / 2
        ll = self.edge_size

        # Select color according validity score
        n = 95 * self.__validity if self.__validity < 2 else 0
        f = 159 * self.__validity if self.__validity < 2 else 255

        try:

            # Draw axis
            axisPoints = numpy.float32([[ll, 0, 0], [0, ll, 0], [0, 0, ll], [0, 0, 0]]).reshape(-1, 3)
            axisPoints, _ = cv.projectPoints(axisPoints, self.__rotation, self.__translation, K, D)
            axisPoints = axisPoints.astype(int)

            frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[0].ravel()), (n,n,f), 5) # X (red)
            frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 5) # Y (green)
            frame = cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 5) # Z (blue)

            # Draw left face
            leftPoints = numpy.float32([[-l, l, l], [-l, -l, l], [-l, -l, -l], [-l, l, -l]]).reshape(-1, 3)
            leftPoints, _ = cv.projectPoints(leftPoints, self.__rotation, self.__translation, K, D)
            leftPoints = leftPoints.astype(int)
            
            frame = cv.line(frame, tuple(leftPoints[0].ravel()), tuple(leftPoints[1].ravel()), (n,n,f), 2)
            frame = cv.line(frame, tuple(leftPoints[1].ravel()), tuple(leftPoints[2].ravel()), (n,n,f), 2)
            frame = cv.line(frame, tuple(leftPoints[2].ravel()), tuple(leftPoints[3].ravel()), (n,n,f), 2)
            frame = cv.line(frame, tuple(leftPoints[3].ravel()), tuple(leftPoints[0].ravel()), (n,n,f), 2)

            # Draw top face
            topPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, l, -l], [l, l, -l]]).reshape(-1, 3)
            topPoints, _ = cv.projectPoints(topPoints, self.__rotation, self.__translation, K, D)
            topPoints = topPoints.astype(int)

            frame = cv.line(frame, tuple(topPoints[0].ravel()), tuple(topPoints[1].ravel()), (n,f,n), 2)
            frame = cv.line(frame, tuple(topPoints[1].ravel()), tuple(topPoints[2].ravel()), (n,f,n), 2)
            frame = cv.line(frame, tuple(topPoints[2].ravel()), tuple(topPoints[3].ravel()), (n,f,n), 2)
            frame = cv.line(frame, tuple(topPoints[3].ravel()), tuple(topPoints[0].ravel()), (n,f,n), 2)

            # Draw front face
            frontPoints = numpy.float32([[l, l, l], [-l, l, l], [-l, -l, l], [l, -l, l]]).reshape(-1, 3)
            frontPoints, _ = cv.projectPoints(frontPoints, self.__rotation, self.__translation, K, D)
            frontPoints = frontPoints.astype(int)
            
            frame = cv.line(frame, tuple(frontPoints[0].ravel()), tuple(frontPoints[1].ravel()), (f,n,n), 2)
            frame = cv.line(frame, tuple(frontPoints[1].ravel()), tuple(frontPoints[2].ravel()), (f,n,n), 2)
            frame = cv.line(frame, tuple(frontPoints[2].ravel()), tuple(frontPoints[3].ravel()), (f,n,n), 2)
            frame = cv.line(frame, tuple(frontPoints[3].ravel()), tuple(frontPoints[0].ravel()), (f,n,n), 2)

        except Exception as e:

            print(e)
            print(self.__translation)
            print(self.__rotation)
            print(self.__succeded)
            print(self.__validity)
            print(axisPoints)

        return frame