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
from collections import Counter

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, print_all=False) -> str:
        """Detector paremeters string representation."""

        output = ''

        for parameter in self.__parameters_names:

            if parameter in self.__dict__.keys():

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

            elif print_all:

                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 frame."""

    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_json(self, json_filepath: str) -> ArUcoDetectorType:
        """Load ArUcoDetector setup from .json file."""

        with open(json_filepath) as configuration_file:

            data = json.load(configuration_file)

            new_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(**data.pop('dictionary'))
            new_marker_size = data.pop('marker_size')
            new_optic_parameters = ArUcoOpticCalibrator.OpticParameters(**data.pop('optic_parameters'))
            new_parameters = DetectorParameters(**data.pop('parameters'))

            return ArUcoDetector(new_dictionary, new_marker_size, new_optic_parameters, new_parameters)

    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, frame: numpy.array):
        """Detect all ArUco markers into a frame.

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

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

        # Detect markers into gray picture
        self.__detected_markers_corners, self.__detected_markers_ids, _ = aruco.detectMarkers(cv.cvtColor(frame, cv.COLOR_BGR2GRAY), self.dictionary.markers, parameters = self.parameters.internal)

        # 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_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)

    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, frame: numpy.array):
        """Draw traked markers."""

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

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

        .. danger:: DON'T MIRROR FRAME
           It makes the markers detection to fail.
        """
        
        # detect markers from gray picture
        gray = cv.cvtColor(frame, 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, frame: numpy.array):
        """Draw detected board corners in frame."""

        if self.__board != None:

            cv.drawChessboardCorners(frame, ((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