1 files changed, 476 insertions, 0 deletions
diff --git a/src/argaze/ArUcoMarker/ArUcoMarkerGroup.py b/src/argaze/ArUcoMarker/ArUcoMarkerGroup.py
new file mode 100644
index 0000000..b013829
--- /dev/null
+++ b/src/argaze/ArUcoMarker/ArUcoMarkerGroup.py
@@ -0,0 +1,476 @@
+""" 
+
+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')