path: root/src/argaze/ArUcoMarker/ArUcoMarkerGroup.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 math
import re
from dataclasses import dataclass
from typing import Self

import cv2
import numpy

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

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


def make_rotation_matrix(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(mat):
	rt = numpy.transpose(mat)
	should_be_identity = numpy.dot(rt, mat)
	i = numpy.identity(3, dtype=mat.dtype)
	n = numpy.linalg.norm(i - should_be_identity)

	return n < 1e-3


@dataclass(frozen=True)
class Place:
	"""Define a place as list of corners position and a marker.

	Parameters:
		corners: 3D corners position in group referential.
		marker: ArUco marker linked to the place.
	"""

	corners: numpy.array
	marker: ArUcoMarker.ArUcoMarker


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

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

		# Init private attributes
		self.marker_size = None
		self.__dictionary = None
		self.__places = {}
		self.__translation = numpy.zeros(3)
		self.__rotation = numpy.zeros(3)

	@property
	def dictionary(self) -> ArUcoMarkerDictionary.ArUcoMarkerDictionary:
		"""Expected dictionary of all markers in the group."""
		return self.__dictionary

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

		self.__dictionary = dictionary

	@property
	def places(self) -> dict:
		"""Expected markers place."""
		return self.__places

	@places.setter
	def places(self, places: dict):

		# Normalize places data
		new_places = {}

		for identifier, data in places.items():

			# Convert string identifier to int value
			if type(identifier) is str:

				identifier = int(identifier)

				# Get translation vector
				tvec = numpy.array(data.pop('translation')).astype(numpy.float32)

				# Check rotation value shape
				rvalue = numpy.array(data.pop('rotation')).astype(numpy.float32)

				# Rotation matrix
				if rvalue.shape == (3, 3):

					rmat = rvalue

				# Rotation vector (expected in degree)
				elif rvalue.shape == (3,):

					rmat = make_rotation_matrix(rvalue[0], rvalue[1], rvalue[2]).astype(numpy.float32)

				else:

					raise ValueError(f'Bad rotation value: {rvalue}')

				assert (is_rotation_matrix(rmat))

				# Get marker size
				size = float(numpy.array(data.pop('size')).astype(numpy.float32))

				new_marker = ArUcoMarker.ArUcoMarker(self.__dictionary, identifier, size)

				# Build marker corners thanks to translation vector and rotation matrix
				place_corners = numpy.array([[-size / 2, size / 2, 0], [size / 2, size / 2, 0], [size / 2, -size / 2, 0], [-size / 2, -size / 2, 0]])
				place_corners = place_corners.dot(rmat) + tvec

				new_places[identifier] = Place(place_corners, new_marker)

			# else places are configured using detected markers estimated points
			elif isinstance(data, ArUcoMarker.ArUcoMarker):

				new_places[identifier] = Place(data.points, data)

			# else places are already at expected format
			elif (type(identifier) is int) and isinstance(data, Place):

				new_places[identifier] = data

		self.__places = new_places

	@property
	def identifiers(self) -> list:
		"""List place marker identifiers belonging to the group."""
		return list(self.__places.keys())

	@property
	def translation(self) -> numpy.array:
		"""Get ArUco marker group translation vector."""
		return self.__translation

	@translation.setter
	def translation(self, tvec):
		"""Set ArUco marker group translation vector."""
		self.__translation = tvec

	@property
	def rotation(self) -> numpy.array:
		"""Get ArUco marker group rotation matrix."""
		return self.__translation

	@rotation.setter
	def rotation(self, rmat):
		"""Set ArUco marker group rotation matrix."""
		self.__rotation = rmat

	def as_dict(self) -> dict:
		"""Export ArUco marker group properties as dictionary."""

		return {
			**DataFeatures.PipelineStepObject.as_dict(self),
			"dictionary": self.__dictionary,
			"places": self.__places
		}

	@classmethod
	def from_obj(cls, obj_filepath: str) -> Self:
		"""Load ArUco markers group from .obj file.

		!!! note
		    Expected object (o) name format: <DICTIONARY>#<IDENTIFIER>_Marker

		!!! note
		    All markers have to belong to the same dictionary.
		"""

		new_dictionary = None
		new_places = {}

		# Regex rules for .obj file parsing
		obj_rx_dict = {
			'object': re.compile(r'o (.*)#([0-9]+)_(.*)\n'),
			'vertices': re.compile(r'v ([+-]?[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]+)\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(ln):

			for k, rx in obj_rx_dict.items():
				m = rx.search(ln)
				if m:
					return k, m

			# If there are no matches
			return None, None

		# Start parsing
		try:

			identifier = None
			vertices = []
			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))

						# Init new group dictionary with first dictionary name
						if new_dictionary is None:

							new_dictionary = ArUcoMarkerDictionary.ArUcoMarkerDictionary(dictionary)

						# Check all others marker dictionary are equal to new group dictionary 
						elif dictionary != new_dictionary.name:

							raise NameError(f'Marker {identifier} dictionary is not {new_dictionary.name}')

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

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

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

						faces[identifier] = [int(match.group(1)), int(match.group(2)), int(match.group(3)), int(match.group(4))]

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

				file.close()

				# Retrieve marker vertices thanks to face vertices ids
				for identifier, face in faces.items():

					# Gather place corners in clockwise order
					cw_corners = numpy.array([vertices[i - 1] for i in reversed(face)])

					# Edit place axis from corners positions
					place_x_axis = cw_corners[2] - cw_corners[3]
					place_x_axis_norm = numpy.linalg.norm(place_x_axis)

					place_y_axis = cw_corners[0] - cw_corners[3]
					place_y_axis_norm = numpy.linalg.norm(place_y_axis)

					# Check axis size: they should be almost equal
					if math.isclose(place_x_axis_norm, place_y_axis_norm, rel_tol=1e-3):

						new_marker_size = place_x_axis_norm

					else:

						raise ValueError(f'{new_dictionary}#{identifier}_Marker is not a square.')

					# Create a new place related to a new marker
					new_marker = ArUcoMarker.ArUcoMarker(new_dictionary, identifier, new_marker_size)
					new_places[identifier] = Place(cw_corners, new_marker)

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

		# Instantiate ArUco markers group
		data = {
			'dictionary': new_dictionary,
			'places': new_places
		}

		return ArUcoMarkerGroup(**data)

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

		Returns:
			dict of markers belonging to this group
			dict of remaining markers not belonging to this group
		"""

		group_markers = {}
		remaining_markers = {}

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

			if marker_id in self.__places.keys():

				group_markers[marker_id] = marker

			else:

				remaining_markers[marker_id] = marker

		return group_markers, remaining_markers

	def estimate_pose_from_markers_corners(self, markers: dict, k: numpy.array, d: numpy.array) -> tuple[
		bool, numpy.array, numpy.array]:
		"""Estimate pose from markers corners and places corners.

		Parameters:
			markers: detected markers to use for pose estimation.
			k: intrinsic camera parameters
			d: camera distortion matrix

		Returns:
			success: True if the pose estimation succeeded
			tvec: scene translation vector
			rvec: scene rotation vector
		"""

		markers_corners_2d = []
		places_corners_3d = []

		for identifier, marker in markers.items():

			try:

				place = self.__places[identifier]

				for marker_corner in marker.corners:
					markers_corners_2d.append(list(marker_corner))

				for place_corner in place.corners:
					places_corners_3d.append(list(place_corner))

			except KeyError:

				raise ValueError(f'Marker {marker.identifier} doesn\'t belong to the group.')

		# SolvPnP using cv2.SOLVEPNP_SQPNP flag
		# TODO: it works also with cv2.SOLVEPNP_EPNP flag so we need to test which is the faster.
		# About SolvPnP flags: https://docs.opencv.org/4.x/d5/d1f/calib3d_solvePnP.html
		success, rvec, tvec = cv2.solvePnP(numpy.array(places_corners_3d), numpy.array(markers_corners_2d), numpy.array(k), numpy.array(d), flags=cv2.SOLVEPNP_SQPNP)

		# Refine pose estimation using Gauss-Newton optimisation
		if success:
			rvec, tvec = cv2.solvePnPRefineVVS(numpy.array(places_corners_3d), numpy.array(markers_corners_2d), numpy.array(k), numpy.array(d), rvec, tvec)

		self.__translation = tvec.T
		self.__rotation = rvec.T

		return success, self.__translation, self.__rotation

	def draw_axes(self, image: numpy.array, k: numpy.array, d: numpy.array, thickness: int = 0, length: float = 0):
		"""Draw group axes."""

		try:
			axis_points = numpy.float32([[length, 0, 0], [0, length, 0], [0, 0, length], [0, 0, 0]]).reshape(-1, 3)
			axis_points, _ = cv2.projectPoints(axis_points, self.__rotation, self.__translation, numpy.array(k), numpy.array(d))
			axis_points = axis_points.astype(int)

			cv2.line(image, tuple(axis_points[3].ravel()), tuple(axis_points[0].ravel()), (0, 0, 255), thickness)  # X (red)
			cv2.line(image, tuple(axis_points[3].ravel()), tuple(axis_points[1].ravel()), (0, 255, 0), thickness)  # Y (green)
			cv2.line(image, tuple(axis_points[3].ravel()), tuple(axis_points[2].ravel()), (255, 0, 0), thickness)  # Z (blue)

		# Ignore errors due to out of field axis: their coordinate are larger than int32 limitations.
		except cv2.error:
			pass

	def draw_places(self, image: numpy.array, k: numpy.array, d: numpy.array, color: tuple = None, border_size: int = 0):
		"""Draw group places."""

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

			try:

				place_points, _ = cv2.projectPoints(place.corners, self.__rotation, self.__translation, numpy.array(k), numpy.array(d))
				place_points = place_points.astype(int)

				cv2.line(image, tuple(place_points[0].ravel()), tuple(place_points[1].ravel()), color, border_size)
				cv2.line(image, tuple(place_points[1].ravel()), tuple(place_points[2].ravel()), color, border_size)
				cv2.line(image, tuple(place_points[2].ravel()), tuple(place_points[3].ravel()), color, border_size)
				cv2.line(image, tuple(place_points[3].ravel()), tuple(place_points[0].ravel()), color, border_size)

			# Ignore errors due to out of field places: their coordinate are larger than int32 limitations.
			except cv2.error:
				pass

	def draw(self, image: numpy.array, k: numpy.array, d: numpy.array, draw_axes: dict = None, draw_places: dict = None):
		"""Draw group axes and places.
	
		Parameters:
			image: where to draw.
			k: intrinsic camera parameters
			d: camera distortion matrix
			draw_axes: draw_axes parameters (if None, no axes drawn)
			draw_places: draw_places parameters  (if None, no places drawn)
		"""

		# Draw axes if required
		if draw_axes is not None:
			self.draw_axes(image, k, d, **draw_axes)

		# Draw places if required
		if draw_places is not None:
			self.draw_places(image, k, d, **draw_places)

	def to_obj(self, obj_filepath):
		"""Save group to .obj file."""

		with open(obj_filepath, 'w', encoding='utf-8') as file:

			file.write('# ArGaze OBJ File\n')
			file.write('# https://achil.recherche.enac.fr/features/eye/argaze/\n')

			v_count = 0

			for p, (identifier, place) in enumerate(self.__places.items()):

				file.write(f'o {self.__dictionary.name}#{identifier}_Marker\n')

				vertices = ''

				# Write vertices in reverse order
				for v in [3, 2, 1, 0]:
					file.write(f'v {" ".join(map(str, place.corners[v]))}\n')
					v_count += 1

					vertices += f' {v_count}'

				# file.write('s off\n')
				file.write(f'f{vertices}\n')