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

T0 = numpy.array([0., 0., 0.])
"""Define no translation vector."""

R0 = numpy.array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]])
"""Define no rotation matrix."""

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.ndarray
    """Position in scene referential."""

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

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

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

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

        self.__dictionary = dictionary
        self.__marker_size = marker_size

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

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

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

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

            self.__load_places_from_obj(data)

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

            for name, place in data.items():

                tvec = numpy.array(place['translation']).astype(numpy.float32)
                rmat = self.__make_rotation_matrix(*place.pop('rotation')).astype(numpy.float32)
                marker = ArUcoMarker.ArUcoMarker(self.__dictionary, place['marker'], self.__marker_size)

                self.__places[name] = Place(tvec, rmat, marker)

        # Init pose data
        self._translation = numpy.zeros(3)
        self._rotation = numpy.zeros(3)

        # 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 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.__places[A_name].rotation
            B = self.__places[B_name].rotation

            if numpy.array_equal(A, B):

                angle = 0.

            else:

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

                # 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.__places[A_name].translation
            B = self.__places[B_name].translation

            # 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\tPlaces:'
        for name, place in self.__places.items():
            output += f'\n\t\t- {name}:'
            output += f'\n{place.translation}'
            output += f'\n{place.rotation}'

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

        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 belonging to the scene."""

        return list(self.__identifier_cache.keys())

    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 __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 = {}
            normals = {}
            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 matrix
                    elif key == 'normal':

                        normals[name] = tuple([float(match.group(1)), float(match.group(2)), float(match.group(3))])

                    # 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():

                    # Gather place corners from counter clockwise ordered face vertices
                    corners = numpy.array([ vertices[i-1] for i in face ])

                    # Edit translation (Tp) allowing to move world axis (W) at place axis (P)
                    Tp = corners.mean(axis=0)

                    # Edit place axis from corners positions
                    place_x_axis = corners[1:3].mean(axis=0) - Tp
                    place_x_axis = place_x_axis / numpy.linalg.norm(place_x_axis)

                    place_y_axis = corners[2:4].mean(axis=0) - Tp
                    place_y_axis = place_y_axis / numpy.linalg.norm(place_y_axis)

                    place_z_axis = normals[name]

                    # Edit rotation (Rp) allowing to transform world axis (W) into place axis (P)
                    W = numpy.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
                    P = numpy.array([place_x_axis, place_y_axis, place_z_axis])
                    Rp = W.dot(P.T)

                    self.__places[name] = Place(Tp, Rp, markers[name])

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

    def filter_markers(self, detected_markers) -> Tuple[dict, dict]:
        """Sort markers belonging to the scene from given detected markers dict (cf ArUcoDetector.detect_markers()).

        * **Returns:**
            - dict of markers belonging to this scene
            - dict of remaining markers not belonging to this scene
        """

        scene_markers = {}
        remaining_markers = {}

        for (marker_id, marker) in detected_markers.items():

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

            except KeyError:
                
                remaining_markers[marker_id] = marker

        return scene_markers, remaining_markers

    def check_markers_consistency(self, scene_markers, angle_tolerance: float, distance_tolerance: float) -> Tuple[dict, dict]:
        """Evaluate if given markers configuration match related places configuration.

        * **Returns:**
            - dict of consistent markers
            - dict of unconsistent markers
            - dict of identified distance or angle unconsistencies and out-of-bounds values
        """

        consistent_markers = {}
        unconsistencies = {}

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

            # Rotation matrix from A marker to B marker
            AB = B_marker.rotation.dot(A_marker.rotation.T)

            # 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=angle_tolerance)
            consistent_distance = math.isclose(distance, expected_distance, abs_tol=distance_tolerance)

            if consistent_angle and consistent_distance:

                if A_name not in consistent_markers.keys():

                    # Remember this marker is already validated
                    consistent_markers[A_name] = A_marker

                if B_name not in consistent_markers.keys():

                    # Remember this marker is already validated
                    consistent_markers[B_name] = B_marker

            else:

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

        # Gather unconsistent markers
        unconsistent_markers = {}

        for name, marker in scene_markers.items():

            if name not in consistent_markers:

                unconsistent_markers[name] = marker

        return consistent_markers, unconsistent_markers, unconsistencies

    def estimate_pose_from_axis_markers(self, axis_markers) -> Tuple[numpy.array, numpy.array]:
        """Calculate tranformations (rotation and translation) from a list of 3 (name, marker) tuples defining an orthogonal axis."""

        O_name, O_marker = axis_markers[0]
        A_name, A_marker = axis_markers[1]
        B_name, B_marker = axis_markers[2]

        O_place = self.__places[O_name]
        A_place = self.__places[A_name]
        B_place = self.__places[B_name]

        # Place axis
        OA = A_place.translation - O_place.translation
        OA = OA / numpy.linalg.norm(OA)

        OB = B_place.translation - O_place.translation
        OB = OB / numpy.linalg.norm(OB)

        # Detect and correct bad place axis orientation
        X_sign = numpy.sign(OA)[0]
        Y_sign = numpy.sign(OB)[1]

        P = numpy.array([OA*X_sign, OB*Y_sign, numpy.cross(OA*X_sign, OB*Y_sign)])

        OA = A_marker.translation - O_marker.translation
        OA = OA / numpy.linalg.norm(OA)

        OB = B_marker.translation - O_marker.translation
        OB = OB / numpy.linalg.norm(OB)

        # Detect and correct bad place axis orientation
        X_sign = numpy.sign(OA)[0]
        Y_sign = -numpy.sign(OB)[1]

        M = numpy.array([OA*X_sign, OB*Y_sign, numpy.cross(OA*X_sign, OB*Y_sign)])

        # Then estimate ArUcoScene rotation
        self._rotation = P.dot(M.T)

        # Consider ArUcoScene translation as the translation of the marker at axis origin
        self._translation = O_marker.translation - O_place.translation.dot(O_place.rotation).dot(O_marker.rotation.T)

        return self._translation, self._rotation

    def estimate_pose_from_any_markers(self, consistent_markers) -> Tuple[numpy.array, numpy.array]:
        """Calculate transformations (rotation and translation) that move each marker to its related place."""

        rotations = []
        translations = []

        for (name, marker) in consistent_markers.items():

            place = self.__places[name]

            # Rotation matrix that transform marker to related place
            R = marker.rotation.dot(place.rotation.T)

            # Translation vector that transform marker to related place
            T = marker.translation - place.translation.dot(place.rotation).dot(marker.rotation.T)

            rotations.append(R)
            translations.append(T)

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

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

        return self._translation, self._rotation

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

        return self._translation

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

        self._translation = tvec

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

        return self._translation

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

        self._rotation = rmat

    def draw_axis(self, frame, K, D, consistency=2):
        """Draw scene axis according a consistency score."""

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

        # Select color according consistency score
        n = 95 * consistency if consistency < 2 else 0
        f = 159 * consistency if consistency < 2 else 255

        # 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), 6) # X (red)
        cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 6) # Y (green)
        cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 6) # Z (blue)

    def draw_places(self, frame, K, D, consistency=2):
        """Draw scene places and their axis according a consistency score."""

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

        # Select color according consistency score
        n = 95 * consistency if consistency < 2 else 0
        f = 159 * consistency if consistency < 2 else 255

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

            try:

                T = self.__places[name].translation
                R = self.__places[name].rotation

                # Draw place axis
                axisPoints = (T + numpy.float32([R.dot([l/2, 0, 0]), R.dot([0, l/2, 0]), R.dot([0, 0, l/2]), R.dot([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), 6) # X (red)
                cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[1].ravel()), (n,f,n), 6) # Y (green)
                cv.line(frame, tuple(axisPoints[3].ravel()), tuple(axisPoints[2].ravel()), (f,n,n), 6) # Z (blue)

                # Draw place
                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), 3)
                cv.line(frame, tuple(placePoints[1].ravel()), tuple(placePoints[2].ravel()), (f,f,f), 3)
                cv.line(frame, tuple(placePoints[2].ravel()), tuple(placePoints[3].ravel()), (f,f,f), 3)
                cv.line(frame, tuple(placePoints[3].ravel()), tuple(placePoints[0].ravel()), (f,f,f), 3)

            except cv.error as e:

                print('ArUcoScene.draw_places: ', e)
                print('T: ', T)
                print('R: ', R)