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

""" 

This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by the Free Software
Foundation, either version 3 of the License, or (at your option) any later
version.
This program is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
You should have received a copy of the GNU General Public License along with
this program. If not, see <https://www.gnu.org/licenses/>.
"""

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

import json
from collections import Counter
from typing import Self

import cv2 as cv
import numpy
from cv2 import aruco

from argaze import DataFeatures
from argaze.ArUcoMarker import ArUcoMarkerDictionary, ArUcoMarker, ArUcoOpticCalibrator


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',
		'useAruco3Detection'
	]

	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(cls, json_filepath) -> Self:
		"""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) -> 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


class ArUcoDetector(DataFeatures.PipelineStepObject):
	"""OpenCV ArUco library wrapper."""

	# noinspection PyMissingConstructor
	@DataFeatures.PipelineStepInit
	def __init__(self, **kwargs):
		"""Initialize ArUcoDetector."""

		# Init private attributes
		self.__dictionary = None
		self.__optic_parameters = None
		self.__parameters = None

		# Init detected markers data
		self.__detected_markers = {}

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

	@property
	def dictionary(self) -> ArUcoMarkerDictionary.ArUcoMarkerDictionary:
		"""ArUco markers dictionary to detect."""
		return self.__dictionary

	@dictionary.setter
	@DataFeatures.PipelineStepAttributeSetter
	def dictionary(self, dictionary: ArUcoMarkerDictionary.ArUcoMarkerDictionary):

		self.__dictionary = dictionary

	@property
	def optic_parameters(self) -> ArUcoOpticCalibrator.OpticParameters:
		"""Optic parameters to use for ArUco detection into image."""
		return self.__optic_parameters

	@optic_parameters.setter
	@DataFeatures.PipelineStepAttributeSetter
	def optic_parameters(self, optic_parameters: ArUcoOpticCalibrator.OpticParameters):

		self.__optic_parameters = optic_parameters

	@property
	def parameters(self) -> DetectorParameters:
		"""ArUco detector parameters."""
		return self.__parameters

	@parameters.setter
	@DataFeatures.PipelineStepAttributeSetter
	def parameters(self, parameters: DetectorParameters):

		self.__parameters = parameters

	@DataFeatures.PipelineStepMethod
	def detect_markers(self, image: numpy.array):
		"""Detect all ArUco markers into an image.

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

		!!! danger "DON'T UNDISTORTED IMAGE"
			Camera intrinsic parameters and distortion coefficients are used later during pose estimation.
		"""

		# Reset detected markers data
		self.__detected_markers, detected_markers_corners, detected_markers_ids = {}, [], []

		# Detect markers into gray picture
		detected_markers_corners, detected_markers_ids, _ = aruco.detectMarkers(cv.cvtColor(image, cv.COLOR_BGR2GRAY),
																				self.__dictionary.markers,
																				parameters=self.__parameters.internal)

		# Is there detected markers ?
		if len(detected_markers_corners) > 0:

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

			for i, marker_id in enumerate(detected_markers_ids):
				marker = ArUcoMarker.ArUcoMarker(self.__dictionary, marker_id)
				marker.corners = detected_markers_corners[i][0]

				# 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

	def estimate_markers_pose(self, size: float, ids: list = []):
		"""Estimate pose detected markers pose considering a marker size.

		Parameters:
			size: size of markers in centimeters.
			ids: markers id list to select detected markers.
		"""

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

			# Select all markers by default
			if len(ids) == 0:
				ids = self.__detected_markers.keys()

			# Prepare data for aruco.estimatePoseSingleMarkers function
			selected_markers_corners = tuple()
			selected_markers_ids = []

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

				if marker_id in ids:
					selected_markers_corners += (marker.corners,)
					selected_markers_ids.append(marker_id)

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

				markers_rvecs, markers_tvecs, markers_points = aruco.estimatePoseSingleMarkers(selected_markers_corners,
																							   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.size = size
					marker.points = markers_points.reshape(4, 3).dot(marker.rotation) - marker.translation

	def detected_markers(self) -> dict[int, ArUcoMarker.ArUcoMarker]:
		"""Access to detected markers' dictionary."""

		return self.__detected_markers

	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)
		detected_markers_corners, detected_markers_ids, _ = aruco.detectMarkers(gray, self.__dictionary.markers,
																				parameters=self.__parameters.internal)

		# if all board markers are detected
		if len(detected_markers_corners) == expected_markers_number:

			self.__board = board
			self.__board_corners_number, self.__board_corners, self.__board_corners_ids = aruco.interpolateCornersCharuco(
				detected_markers_corners, 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 is not None:
			cv.drawChessboardCorners(image, ((self.__board.size[0] - 1), (self.__board.size[1] - 1)),
									 self.__board_corners, True)

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

		return self.__board_corners_number

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

		return self.__board_corners_ids

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

		return self.__board_corners


class Observer():
	"""Define ArUcoDetector observer to count how many times detection succeeded and how many times markers are detected."""

	def __init__(self):
		"""Initialize marker detection metrics."""

		self.__try_count = 0
		self.__success_count = 0
		self.__detected_ids = []

	@property
	def metrics(self) -> tuple[int, int, dict]:
		"""Get marker detection metrics.

		Returns:
			number of detect function call
			dict with number of detection for each marker identifier
		"""

		return self.__try_count, self.__success_count, Counter(self.__detected_ids)

	def reset(self):
		"""Reset marker detection metrics."""

		self.__try_count = 0
		self.__success_count = 0
		self.__detected_ids = []

	def on_detect_markers(self, timestamp, aruco_detector, exception):
		"""Update ArUco markers detection metrics."""

		self.__try_count += 1
		detected_markers_list = list(aruco_detector.detected_markers().keys())

		if len(detected_markers_list):
			self.__success_count += 1
			self.__detected_ids.extend(detected_markers_list)