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

#!/usr/bin/env python

from typing import TypeVar, Tuple
from dataclasses import dataclass, field
import json
import math
import itertools
import re

from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoCamera

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

ArUcoSceneType = TypeVar('ArUcoScene', bound="ArUcoScene")
# Type definition for type annotation convenience

@dataclass(frozen=True)
class Place():
    """Define a place as a pose and a marker."""

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

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

    marker: dict
    """ArUco marker linked to the place."""

class ArUcoScene():
    """Define abstract class to handle group of ArUco markers as one unique spatial entity and estimate its pose."""

    angle_tolerance: float
    """Angle error tolerance allowed to validate place pose in degree."""

    distance_tolerance: float
    """Distance error tolerance allowed to validate place pose in centimeter."""

    def __init__(self, dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary, marker_size: float, angle_tolerance: float, distance_tolerance: float, places: dict | str = None):
        """Define scene attributes."""

        self.__dictionary = dictionary
        self.__marker_size = marker_size

        self.angle_tolerance = angle_tolerance
        self.distance_tolerance = distance_tolerance

        # NEVER USE {} as default function argument
        self.places = places

    @property
    def places(self) -> dict:
        """All named places of the scene and their ArUco markers."""

        return self.__places
    
    @places.setter
    def places(self, places: dict | str):

        # str: path to .obj file
        if type(places) == str:

            self.__load_places_from_obj(places)

        # dict: all places
        else:
  
            self.__places = {}

            for name, place in places.items():
                marker = ArUcoMarker.ArUcoMarker(self.__dictionary, place['marker'], self.__marker_size)
                self.__places[name] = Place(numpy.array(place['translation']).astype(numpy.float32), numpy.array(place['rotation']).astype(numpy.float32), marker)

        # Init pose data
        self._translation = numpy.zeros(3)
        self._rotation = numpy.zeros(3)
        self._succeded = False
        self._consistent_markers = 0

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

        # Process each place pose to speed up further calculations
        self.__translation_cache = {}
        for name, place in self.__places.items():
            self.__translation_cache[name] = place.translation

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

            # Create intrinsic rotation matrix
            R = self.__euler_vector_to_rotation_matrix(*place.rotation)

            assert(self.__is_rotation_matrix(R))

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

        # Process axis-angle between place combination to speed up further calculations
        self.__angle_cache = {}
        for (A_name, A_place), (B_name, B_place) in itertools.combinations(self.__places.items(), 2):

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

            if numpy.array_equal(A, B):

                angle = 0.

            else:

                # Rotation matrix from A place to B place
                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 each place combination to speed up further calculations
        self.__distance_cache = {}
        for (A_name, A_place), (B_name, B_place) in itertools.combinations(self.__places.items(), 2):

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

            # Calculate axis-angle representation of AB rotation matrix
            distance = numpy.linalg.norm(B - A)

            try:
                self.__distance_cache[A_name][B_name] = distance
            except:
                self.__distance_cache[A_name] = {B_name: distance}

            try:
                self.__distance_cache[B_name][A_name] = distance
            except:
                self.__distance_cache[B_name] = {A_name: distance}

    @classmethod
    def from_json(self, json_filepath) -> ArUcoSceneType:
        """Load ArUco scene from .json file."""

        with open(json_filepath) as configuration_file:

            return ArUcoScene(**json.load(configuration_file))
    
    def __str__(self) -> str:
        """String display"""

        output = f'\n\n\tDictionary: {self.__dictionary.name}'

        output += '\n\n\tIdentifier cache:'
        for i, name in self.__identifier_cache.items():
            output += f'\n\t\t- {i}: {name}'

        output += '\n\n\tTranslation cache:'
        for name, item in self.__translation_cache.items():
            output += f'\n\t\t- {name}: {item}'

        output += '\n\n\tRotation cache:'
        for name, item in self.__rotation_cache.items():
            output += f'\n\t\t- {name}:\n{item}'

        output += '\n\n\tAngle cache:'
        for A_name, A_angle_cache in self.__angle_cache.items():
            for B_name, angle in A_angle_cache.items():
                output += f'\n\t\t- {A_name}/{B_name}: {angle:3f}'

        output += '\n\n\tDistance cache:'
        for A_name, A_distance_cache in self.__distance_cache.items():
            for B_name, distance in A_distance_cache.items():
                output += f'\n\t\t- {A_name}/{B_name}: {distance:3f}'

        return output

    @property
    def identifiers(self) -> list:
        """List all makers identifier."""

        return list(self.__identifier_cache.keys())

    def __load_places_from_obj(self, obj_filepath: str) -> dict:
        """Load places from .obj file.

        .. warning:: 'o' tag string format should be DICTIONARY#IDENTIFIER_NAME
        """

        self.__places = {}

        # Regex rules for .obj file parsing
        OBJ_RX_DICT = {
            'object': re.compile(r'o (.*)#([0-9]+)_(.*)\n'),
            'vertice': re.compile(r'v ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+)\n'),
            'normal': re.compile(r'vn ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+) ([+-]?[0-9]*[.]?[0-9]+)\n'),
            'face': re.compile(r'f ([0-9]+)//([0-9]+) ([0-9]+)//([0-9]+) ([0-9]+)//([0-9]+) ([0-9]+)//([0-9]+)\n'),
            'comment': re.compile(r'#(.*)\n') # keep comment regex after object regex because the # is used in object string too
        }

        # Regex .obj line parser
        def __parse_obj_line(line):

            for key, rx in OBJ_RX_DICT.items():
                match = rx.search(line)
                if match:
                    return key, match

            # If there are no matches
            return None, None
        
        # Start parsing
        try:

            name = None
            vertices = []
            markers = {}
            rotations = {}
            faces = {}

            # Open the file and read through it line by line
            with open(obj_filepath, 'r') as file:

                line = file.readline()

                while line:

                    # At each line check for a match with a regex
                    key, match = __parse_obj_line(line)

                    # Extract comment
                    if key == 'comment':
                        pass

                    # Extract marker dictionary and identifier
                    elif key == 'object':

                        dictionary = str(match.group(1))
                        identifier = int(match.group(2))
                        last = str(match.group(3))

                        # Check that marker dictionary is like the scene dictionary
                        if dictionary == self.__dictionary.name:

                            name = f'{dictionary}#{identifier}' # ignore last part
                            markers[name] = ArUcoMarker.ArUcoMarker(self.__dictionary, identifier, self.__marker_size)

                        else:

                            raise NameError(f'Marker#{identifier} dictionary is not {self.__dictionary.name}')

                    # Fill vertices array
                    elif key == 'vertice':

                        vertices.append(tuple([float(match.group(1)), float(match.group(2)), float(match.group(3))]))

                    # Extract normal to calculate rotation vectore
                    elif key == 'normal':

                        R = self.__normal_vector_to_rotation_matrix(float(match.group(1)), float(match.group(2)), float(match.group(3)))
                        rvecs, _ = cv.Rodrigues(R)
                        rvecs = rvecs.reshape(1, 3)[0]
                        rvecs = numpy.array([numpy.rad2deg(rvecs[0]), numpy.rad2deg(rvecs[1]), numpy.rad2deg(rvecs[2])])

                        print(name)
                        print(R)
                        print(rvecs)

                        rotations[name] = rvecs

                    # Extract vertice ids
                    elif key == 'face':

                        faces[name] = [int(match.group(1)), int(match.group(3)), int(match.group(5)), int(match.group(7))]

                    # Go to next line
                    line = file.readline()

                file.close()

                # Retreive marker vertices thanks to face vertice ids
                for name, face in faces.items():

                    translation = numpy.array([ vertices[i-1] for i in face ]).mean(axis=0)
                    rotation = rotations[name]
                    marker = markers[name]

                    self.__places[name] = Place(translation, rotation, marker) 

        except IOError:
            raise IOError(f'File not found: {obj_filepath}')

    def __normal_vector_to_rotation_matrix(self, x, y, z):
        # Applying formula found at https://math.stackexchange.com/questions/1956699/getting-a-transformation-matrix-from-a-normal-vector

        xy_dist = math.sqrt(x**2 + y**2)

        if xy_dist > 0:

            return numpy.array([[y/xy_dist, -x/xy_dist, 0], [x*z/xy_dist, y*z/xy_dist, -xy_dist], [x, y, z]])

        else:

            return numpy.array([[1, 0, 0], [0, 1, 0], [x, y, z]])

    def __euler_vector_to_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_marker_pose(self, name, place, F):

        # Transform place rotation into scene 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 place translation into scene translation vector
        OF = place.translation
        T = self.__translation_cache[name]
        FC = R.dot(F.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) -> Tuple[numpy.array, numpy.array, bool, int, dict]:
        """Estimate scene pose from tracked markers (cf ArUcoTracker.track()) and validate its consistency according expected scene places.
        
        * **Returns:** 
            - scene translation vector
            - scene rotation vector
            - scene pose estimation success status
            - all tracked markers considered as consistent and used to estimate the pose
            - dict of identified distance or angle unconsistencies and out-of-bounds values
        """

        # Init pose data
        self._translation = numpy.zeros(3)
        self._rotation = numpy.zeros(3)
        self._succeded = False
        self._consistent_markers = {}
        self._unconsistencies = {}

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

            return self._translation, self._rotation, self._succeded, self._consistent_markers, self._unconsistencies

        # Look for tracked markers which belong to the scene and store them by name
        scene_markers = {}
        for (marker_id, marker) in tracked_markers.items():

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

            except KeyError:
                continue

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

        # Pose consistency checking is'nt possible when only one marker is tracked
        if len(scene_markers) == 1:

            # Get scene pose from to the unique marker
            name, marker = scene_markers.popitem()
            F, _ = cv.Rodrigues(marker.rotation)

            self._rotation, self._translation = self.__normalise_marker_pose(name, marker, F)
            self._succeded = True
            self._consistent_markers[name] = marker

            #print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
            #print(f'ArUcoScene rotation vector: {self._rotation[0][0]:3f} {self._rotation[1][0]:3f} {self._rotation[2][0]:3f}')
            #print(f'ArUcoScene translation vector: {self._translation[0]:3f} {self._translation[1]:3f} {self._translation[2]:3f}')

        # Pose consistency checking processes markers two by two
        else:

            consistent_markers = []
            consistent_rvecs = []
            consistent_tvecs = []

            for (A_name, A_marker), (B_name, B_marker) in itertools.combinations(scene_markers.items(), 2):

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

                # Get marker rotation estimation
                # Use rotation matrix instead of rotation vector
                A, _ = cv.Rodrigues(A_marker.rotation)
                B, _ = cv.Rodrigues(B_marker.rotation)

                # Rotation matrix from A marker to B marker
                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))
                expected_angle = self.__angle_cache[A_name][B_name]

                # Calculate distance between A marker center and B marker center
                distance = numpy.linalg.norm(A_marker.translation - B_marker.translation)
                expected_distance = self.__distance_cache[A_name][B_name]

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

                if consistent_angle and consistent_distance:

                    if A_name not in consistent_markers:

                        # Remember this marker is already validated
                        consistent_markers.append(A_name)

                        rvec, tvec = self.__normalise_marker_pose(A_name, A_marker, A)

                        # Store normalised marker pose
                        consistent_rvecs.append(rvec)
                        consistent_tvecs.append(tvec)
                        
                    if B_name not in consistent_markers:

                        # Remember this marker is already validated
                        consistent_markers.append(B_name)

                        rvec, tvec = self.__normalise_marker_pose(B_name, B_marker, B)

                        # Store normalised marker pose
                        consistent_rvecs.append(rvec)
                        consistent_tvecs.append(tvec)

                else:

                    if not consistent_angle:
                        self._unconsistencies[f'{A_name}/{B_name} angle'] = angle
                    
                    if not consistent_distance:
                        self._unconsistencies[f'{A_name}/{B_name} distance'] = distance

            if len(consistent_markers) > 1:

                # Consider ArUcoScene rotation as the mean of all consistent 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(consistent_rvecs), axis=0)

                # Consider ArUcoScene translation as the mean of all consistent translations
                self._translation = numpy.mean(numpy.array(consistent_tvecs), axis=0)

                #print(':::::::::::::::::::::::::::::::::::::::::::::::::::')
                #print(f'ArUcoScene rotation vector: {self._rotation[0][0]:3f} {self._rotation[1][0]:3f} {self._rotation[2][0]:3f}')
                #print(f'ArUcoScene translation vector: {self._translation[0]:3f} {self._translation[1]:3f} {self._translation[2]:3f}')

                self._succeded = True
                self._consistent_markers = consistent_markers

            else:

                unconsistent_rvecs = []
                unconsistent_tvecs = []

                # Gather unconsistent pose estimations
                for name, marker in scene_markers.items():

                    if name not in consistent_markers:

                        R, _ = cv.Rodrigues(marker.rotation)
                        rvec, tvec = self.__normalise_marker_pose(name, marker, R)

                        unconsistent_rvecs = [rvec]
                        unconsistent_tvecs = [tvec]

                # Consider ArUcoScene rotation as the mean of all unconsistent 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(unconsistent_rvecs), axis=0)

                # Consider ArUcoScene translation as the mean of all unconsistent translations
                self._translation = numpy.mean(numpy.array(unconsistent_tvecs), axis=0)

                #print(':::::::::::::::::::::::::::::::::::::::::::::::::::')
                #print(f'ArUcoScene rotation vector: {self._rotation[0][0]:3f} {self._rotation[1][0]:3f} {self._rotation[2][0]:3f}')
                #print(f'ArUcoScene translation vector: {self._translation[0]:3f} {self._translation[1]:3f} {self._translation[2]:3f}')

                self._succeded = False
                self._consistent_markers = {}

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

        return self._translation, self._rotation, self._succeded, self._consistent_markers, self._unconsistencies

    @property
    def translation(self) -> numpy.array:
        """Access to scene translation vector.

        .. warning::
           Setting scene translation vector implies succeded status to be True and consistency score to be 0."""

        return self._translation

    @translation.setter
    def translation(self, tvec):

        self._translation = tvec
        self._succeded = True
        self._consistent_markers = 0

    @property
    def rotation(self) -> numpy.array:
        """Access to scene rotation vector.

        .. warning::
           Setting scene rotation vector implies succeded status to be True and consistency score to be 0."""

        return self._translation

    @rotation.setter
    def rotation(self, rvec):

        self._rotation = rvec
        self._succeded = True
        self._consistent_markers = 0

    @property
    def succeded(self) -> bool:
        """Access to scene pose estimation succeded status."""

        return self._succeded

    @property
    def consistency(self) -> int:
        """Access to scene pose estimation consistency score."""

        return len(self._consistent_markers)

    def draw(self, frame, K, D, draw_places=True):
        """Draw scene axis and places."""

        l = self.__marker_size / 2
        ll = self.__marker_size

        # Select color according consistency score
        n = 95 * self.consistency if self.consistency < 2 else 0
        f = 159 * self.consistency if self.consistency < 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, numpy.array(K), numpy.array(D))
            axisPoints = axisPoints.astype(int)

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

            # Draw places (optional)
            if draw_places:

                for name, place in self.__places.items():

                    if name != "top":
                        continue

                    T = self.__translation_cache[name]
                    R = self.__rotation_cache[name]

                    placePoints = (T + numpy.float32([R.dot([-l, -l, 0]), R.dot([l, -l, 0]), R.dot([l, l, 0]), R.dot([-l, l, 0])])).reshape(-1, 3)
                    placePoints, _ = cv.projectPoints(placePoints, self._rotation, self._translation, numpy.array(K), numpy.array(D))
                    placePoints = placePoints.astype(int)
                    
                    cv.line(frame, tuple(placePoints[0].ravel()), tuple(placePoints[1].ravel()), (f,f,f), 2)
                    cv.line(frame, tuple(placePoints[1].ravel()), tuple(placePoints[2].ravel()), (f,f,f), 2)
                    cv.line(frame, tuple(placePoints[2].ravel()), tuple(placePoints[3].ravel()), (f,f,f), 2)
                    cv.line(frame, tuple(placePoints[3].ravel()), tuple(placePoints[0].ravel()), (f,f,f), 2)

        except Exception as e:

            print(e)
            print(self._translation)
            print(self._rotation)
            print(self._succeded)
            print(self._consistent_markers)
            print(axisPoints)