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

#!/usr/bin/env python

""" """

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

from typing import TypeVar, Tuple
from dataclasses import dataclass, field
import json
import os
from collections import Counter
import time

from argaze.ArUcoMarkers import ArUcoMarkersDictionary, ArUcoMarker, ArUcoOpticCalibrator

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

ArUcoMarkerDictionaryType = TypeVar('ArUcoMarkerDictionary', bound="ArUcoMarkerDictionary")
# Type definition for type annotation convenience

ArUcoMarkerType = TypeVar('ArUcoMarker', bound="ArUcoMarker")
# Type definition for type annotation convenience

OpticParametersType = TypeVar('OpticParameters', bound="OpticParameters")
# Type definition for type annotation convenience

DetectorParametersType = TypeVar('DetectorParameters', bound="DetectorParameters")
# Type definition for type annotation convenience

ArUcoDetectorType = TypeVar('ArUcoDetector', bound="ArUcoDetector")
# Type definition for type annotation convenience

class DetectorParameters():
    """Wrapper class around ArUco marker detector parameters.

    .. note:: More details on [opencv page](https://docs.opencv.org/4.x/d1/dcd/structcv_1_1aruco_1_1DetectorParameters.html)
    """

    __parameters = aruco.DetectorParameters()
    __parameters_names = [
        'adaptiveThreshConstant',
        'adaptiveThreshWinSizeMax',
        'adaptiveThreshWinSizeMin',
        'adaptiveThreshWinSizeStep',
        'aprilTagCriticalRad',
        'aprilTagDeglitch',
        'aprilTagMaxLineFitMse',
        'aprilTagMaxNmaxima',
        'aprilTagMinClusterPixels',
        'aprilTagMinWhiteBlackDiff',
        'aprilTagQuadDecimate',
        'aprilTagQuadSigma',
        'cornerRefinementMaxIterations',
        'cornerRefinementMethod',
        'cornerRefinementMinAccuracy',
        'cornerRefinementWinSize',
        'markerBorderBits',
        'minMarkerPerimeterRate',
        'maxMarkerPerimeterRate',
        'minMarkerDistanceRate',
        'detectInvertedMarker',
        'errorCorrectionRate',
        'maxErroneousBitsInBorderRate',
        'minCornerDistanceRate',
        'minDistanceToBorder',
        'minOtsuStdDev',
        'perspectiveRemoveIgnoredMarginPerCell',
        'perspectiveRemovePixelPerCell',
        'polygonalApproxAccuracyRate'
    ]

    def __init__(self, **kwargs):

        for parameter, value in kwargs.items():

            setattr(self.__parameters, parameter, value)

        self.__dict__.update(kwargs)

    def __setattr__(self, parameter, value):

        setattr(self.__parameters, parameter, value)

    def __getattr__(self, parameter):

        return getattr(self.__parameters, parameter)

    @classmethod
    def from_json(self, json_filepath) -> DetectorParametersType:
        """Load detector parameters from .json file."""

        with open(json_filepath) as configuration_file:

            return DetectorParameters(**json.load(configuration_file))

    def __str__(self) -> str:
        """Detector parameters string representation."""

        return f'{self}'

    def __format__(self, spec) -> str:
        """Formated detector parameters string representation.

        Parameters:
            spec: 'modified' to get only modified parameters."""

        output = ''

        for parameter in self.__parameters_names:

            if parameter in self.__dict__.keys():

                output += f'\t*{parameter}: {getattr(self.__parameters, parameter)}\n'

            elif spec == "":

                output += f'\t{parameter}: {getattr(self.__parameters, parameter)}\n'

        return output

    @property
    def internal(self):
        return self.__parameters

@dataclass
class ArUcoDetector():
    """ArUco markers detector."""

    dictionary: ArUcoMarkersDictionary.ArUcoMarkersDictionary = field(default_factory=ArUcoMarkersDictionary.ArUcoMarkersDictionary)
    """ArUco markers dictionary to detect."""

    marker_size: float = field(default=0.)
    """Size of ArUco markers to detect in centimeter."""

    optic_parameters: ArUcoOpticCalibrator.OpticParameters = field(default_factory=ArUcoOpticCalibrator.OpticParameters)
    """Optic parameters to use for ArUco detection into image."""

    parameters: DetectorParameters = field(default_factory=DetectorParameters)
    """ArUco detector parameters."""

    def __post_init__(self):

        # Init detected markers data
        self.__detected_markers = {}
        self.__detected_markers_corners = []
        self.__detected_markers_ids = []

        # Init detected board data
        self.__board = None
        self.__board_corners_number = 0
        self.__board_corners = []
        self.__board_corners_ids = []

        # Init detect metrics data
        self.__detection_count = 0
        self.__detected_ids = []

    @classmethod
    def from_dict(self, aruco_detector_data: dict, working_directory: str = None) -> ArUcoDetectorType:
        """Load attributes from dictionary.

        Parameters:
            aruco_detector_data: dictionary with attributes to load
            working_directory: folder path where to load files when a dictionary value is a relative filepath.
        """

        # Load ArUco dictionary
        dictionary_value = aruco_detector_data.pop('dictionary')

        # str: dictionary name
        if type(dictionary_value) == str:

            new_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(dictionary_value)

        # dict:
        else:

            new_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(**dictionary_value)

        # Load ArUco marker size
        new_marker_size = aruco_detector_data.pop('marker_size')

        # Load optic parameters
        try:
            optic_parameters_value = aruco_detector_data.pop('optic_parameters')

            # str: relative path to .json file
            if type(optic_parameters_value) == str:

                optic_parameters_value = os.path.join(working_directory, optic_parameters_value)
                new_optic_parameters = ArUcoOpticCalibrator.OpticParameters.from_json(optic_parameters_value)

            # dict:
            else:

                new_optic_parameters = ArUcoOpticCalibrator.OpticParameters(**optic_parameters_value)

        except KeyError:

            new_optic_parameters = None

        # Load ArUco detector parameters
        try:

            # Check detector parameters value type
            parameters_value = aruco_detector_data.pop('parameters')

            # str: relative path to .json file
            if type(parameters_value) == str:

                parameters_value = os.path.join(working_directory, parameters_value)
                new_parameters = DetectorParameters.from_json(parameters_value)

            # dict:
            else:

                new_parameters = DetectorParameters(**parameters_value)

        except KeyError:

            new_parameters = DetectorParameters()

        # Create aruco detector
        return ArUcoDetector(new_dictionary, new_marker_size, new_optic_parameters, new_parameters)

    @classmethod
    def from_json(self, json_filepath: str) -> ArUcoDetectorType:
        """
        Load attributes from .json file.

        Parameters:
            json_filepath: path to json file
        """

        with open(json_filepath) as configuration_file:

            aruco_detector_data = json.load(configuration_file)
            working_directory = os.path.dirname(json_filepath)

            return ArUcoDetector.from_dict(aruco_detector_data, working_directory)

    def __str__(self) -> str:
        """String display"""

        output = f'\n\tDictionary: {self.dictionary}\n'
        output += f'\tMarker size: {self.marker_size} cm\n\n'
        output += f'\tOptic parameters:\n{self.optic_parameters}\n'
        output += f'\tDetection Parameters:\n{self.parameters}'

        return output

    def detect_markers(self, image: numpy.array) -> float:
        """Detect all ArUco markers into an image.

        .. danger:: DON'T MIRROR IMAGE
           It makes the markers detection to fail.

        Returns:
            - detection time: marker detection time in ms
        """

        # Reset detected markers data
        self.__detected_markers, self.__detected_markers_corners, self.__detected_markers_ids = {}, [], []

        # Store marker detection start date
        detection_start = time.perf_counter()

        # Detect markers into gray picture
        self.__detected_markers_corners, self.__detected_markers_ids, _ = aruco.detectMarkers(cv.cvtColor(image, cv.COLOR_BGR2GRAY), self.dictionary.markers, parameters = self.parameters.internal)
        
        # Assess marker detection time in ms
        detection_time = (time.perf_counter() - detection_start) * 1e3

        # Is there detected markers ?
        if len(self.__detected_markers_corners) > 0:

            # Transform markers ids array into list
            self.__detected_markers_ids = self.__detected_markers_ids.T[0]

            # Gather detected markers data and update metrics
            self.__detection_count += 1

            for i, marker_id in enumerate(self.__detected_markers_ids):

                marker = ArUcoMarker.ArUcoMarker(self.dictionary, marker_id, self.marker_size)

                marker.corners = self.__detected_markers_corners[i]

                # No pose estimation: call estimate_markers_pose to get one
                marker.translation = numpy.empty([0])
                marker.rotation = numpy.empty([0])
                marker.points = numpy.empty([0])

                self.__detected_markers[marker_id] = marker

                self.__detected_ids.append(marker_id)

        return detection_time

    def estimate_markers_pose(self, markers_ids: list = []):
        """Estimate pose of current detected markers or of given markers id list."""

        # Is there detected markers ?
        if len(self.__detected_markers_corners) > 0:

            # Is there a marker selection ?
            if len(markers_ids) > 0:

                selected_markers_corners = tuple()
                selected_markers_ids = []

                for i, marker_id in enumerate(self.__detected_markers_ids):

                    if marker_id in markers_ids:

                        selected_markers_corners += (self.__detected_markers_corners[i],)
                        selected_markers_ids.append(marker_id)

            # Otherwise, estimate pose of all markers
            else:

                selected_markers_corners = self.__detected_markers_corners
                selected_markers_ids = self.__detected_markers_ids

            # Estimate pose of selected markers
            if len(selected_markers_corners) > 0:

                markers_rvecs, markers_tvecs, markers_points = aruco.estimatePoseSingleMarkers(selected_markers_corners, self.marker_size, numpy.array(self.optic_parameters.K), numpy.array(self.optic_parameters.D)) 

                for i, marker_id in enumerate(selected_markers_ids):

                    marker = self.__detected_markers[marker_id]

                    marker.translation = markers_tvecs[i][0]
                    marker.rotation, _ = cv.Rodrigues(markers_rvecs[i][0])
                    marker.points = markers_points.reshape(4, 3)

    @property
    def detected_markers(self) -> dict[ArUcoMarkerType]:
        """Access to detected markers dictionary."""

        return self.__detected_markers

    @property
    def detected_markers_number(self) -> int:
        """Return detected markers number."""

        return len(list(self.__detected_markers.keys()))

    def draw_detected_markers(self, image: numpy.array, draw_marker: dict = None):
        """Draw detected markers.

        Parameters:
            image: image where to draw
            draw_marker: ArucoMarker.draw parameters (if None, no marker drawn)
        """

        if draw_marker is not None:

            for marker_id, marker in self.__detected_markers.items():

                marker.draw(image, self.optic_parameters.K, self.optic_parameters.D, **draw_marker)
            
    def detect_board(self, image: numpy.array, board, expected_markers_number):
        """Detect ArUco markers board in image setting up the number of detected markers needed to agree detection.

        .. danger:: DON'T MIRROR IMAGE
           It makes the markers detection to fail.
        """
        
        # detect markers from gray picture
        gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
        self.__detected_markers_corners, self.__detected_markers_ids, _ = aruco.detectMarkers(gray, self.dictionary.markers, parameters = self.parameters.internal)

         # if all board markers are detected
        if len(self.__detected_markers_corners) == expected_markers_number:

            self.__board = board
            self.__board_corners_number, self.__board_corners, self.__board_corners_ids = aruco.interpolateCornersCharuco(self.__detected_markers_corners, self.__detected_markers_ids, gray, self.__board.model)

        else:

            self.__board = None
            self.__board_corners_number = 0
            self.__board_corners = []
            self.__board_corners_ids = []

    def draw_board(self, image: numpy.array):
        """Draw detected board corners in image."""

        if self.__board != None:

            cv.drawChessboardCorners(image, ((self.__board.size[0] - 1 ), (self.__board.size[1] - 1)), self.__board_corners, True)

    def reset_detection_metrics(self):
        """Enable marker detection metrics."""

        self.__detection_count = 0
        self.__detected_ids = []

    @property
    def detection_metrics(self) -> Tuple[int, dict]:
        """Get marker detection metrics.
        Returns:
                number of detect function call
                dict with number of detection for each marker identifier"""

        return self.__detection_count, Counter(self.__detected_ids)

    @property
    def board_corners_number(self) -> int:
        """Get detected board corners number."""

        return self.__board_corners_number

    @property
    def board_corners_identifier(self) -> list[int]:
        """Get detected board corners identifier."""

        return self.__board_corners_ids

    @property
    def board_corners(self) -> list:
        """Get detected board corners."""

        return self.__board_corners