1 files changed, 343 insertions, 0 deletions
diff --git a/src/argaze/ArFeatures.py b/src/argaze/ArFeatures.py
new file mode 100644
index 0000000..25fd10d
--- /dev/null
+++ b/src/argaze/ArFeatures.py
@@ -0,0 +1,343 @@
+#!/usr/bin/env python
+
+from typing import TypeVar, Tuple
+from dataclasses import dataclass, field
+import json
+import os
+
+from argaze.ArUcoMarkers import *
+from argaze.AreaOfInterest import *
+
+import numpy
+
+ArEnvironmentType = TypeVar('ArEnvironment', bound="ArEnvironment")
+# Type definition for type annotation convenience
+
+ArSceneType = TypeVar('ArScene', bound="ArScene")
+# Type definition for type annotation convenience
+
+AOI2DSceneType = TypeVar('AOI2DScene', bound="AOI2DScene")
+# Type definition for type annotation convenience
+
+@dataclass
+class ArEnvironment():
+	"""Define an Augmented Reality environment based ArUco marker detection."""
+
+	name: str 
+	"""Environement name."""
+
+	aruco_detector: ArUcoDetector.ArUcoDetector = field(init=False, default_factory=ArUcoDetector.ArUcoDetector)
+	"""ArUco detecor."""
+
+	def __init__(self, **kwargs):
+
+		self.name = kwargs.pop('name')
+
+		self.aruco_detector = ArUcoDetector.ArUcoDetector(**kwargs.pop('aruco_detector'))
+
+		self.__scenes = {}
+		for name, scene_kwargs in kwargs.items():
+
+			self.__scenes[name] = ArScene(self, **scene_kwargs)
+
+	def __getitem__(self, name) -> ArSceneType:
+		"""Get an ArScene of the environment."""
+
+		return self.__scenes[name]
+
+	@classmethod
+	def from_json(self, json_filepath: str) -> ArSceneType:
+		"""Load ArEnvironment from .json file."""
+
+		with open(json_filepath) as configuration_file:
+
+			return ArEnvironment(**json.load(configuration_file))
+
+	def __str__(self) -> str:
+		"""String display"""
+
+		output = f'ArUcoDetector:\n{self.aruco_detector}\n'
+
+		for name, scene in self.__scenes.items():
+			output += f'\"{name}\" ArScene:\n{scene}\n'
+
+		return output
+
+	def items(self) -> Tuple[str, ArSceneType]:
+		"""Iterate over scenes."""
+
+		return self.__scenes.items()
+
+	def keys(self) -> list[str]:
+		"""Get scenes name."""
+
+		return self.__scenes.keys()
+
+class PoseEstimationFailed(Exception):
+	"""Exception raised by ArScene project method when the pose can't be estimated due to unconsistencies."""
+
+	def __init__(self, message, unconsistencies=None):  
+
+		super().__init__(message)
+
+		self.unconsistencies = unconsistencies
+
+class SceneProjectionFailed(Exception):
+	"""Exception raised by ArScene project method when the scene can't be projected."""
+
+	def __init__(self, message):  
+
+		super().__init__(message)
+
+@dataclass
+class ArScene():
+	"""Define an Augmented Reality scene based ArUco markers and AOI scenes."""
+
+	aruco_scene: ArUcoScene.ArUcoScene = field(init=False, default_factory=ArUcoScene.ArUcoScene)
+	"""ArUco scene ..."""
+
+	aoi_scene: AOI3DScene.AOI3DScene = field(init=False, default_factory=AOI3DScene.AOI3DScene)
+	"""AOI 3D scene ..."""
+
+	angle_tolerance: float
+	"""Angle error tolerance allowed to validate marker pose in degree."""
+
+	distance_tolerance: float
+	"""Distance error tolerance allowed to validate marker pose in centimeter."""
+
+	aruco_axis: dict
+	"""Dictionary of orthogonal axis where each axis is defined by list of 3 markers identifier (first is origin)."""
+
+	aruco_aoi: dict
+	"""Dictionary of AOI defined by list of markers identifier and markers corners index tuples."""
+
+	def __init__(self, ar_environment: ArEnvironment, **kwargs):
+
+		self.__ar_environment = ar_environment
+
+		# Check aruco_scene value type
+		aruco_scene_value = kwargs.pop('aruco_scene')
+
+		# str: relative path to .obj file
+		if type(aruco_scene_value) == str:
+
+			aruco_scene_value = os.path.join(os.getcwd(), aruco_scene_value)
+
+		self.aruco_scene = ArUcoScene.ArUcoScene(self.__ar_environment.aruco_detector.dictionary, self.__ar_environment.aruco_detector.marker_size, aruco_scene_value)
+
+		# Check aoi_scene value type
+		aoi_scene_value = kwargs.pop('aoi_scene')
+
+		# str: relative path to .obj file
+		if type(aoi_scene_value) == str:
+
+			obj_filepath = os.path.join(os.getcwd(), aoi_scene_value)
+			self.aoi_scene = AOI3DScene.AOI3DScene.from_obj(obj_filepath)
+
+		# dict: all AOI
+		else:
+			self.aoi_scene = AOI3DScene.AOI3DScene(aoi_scene_value)
+
+		# Init aruco axis
+		self.aruco_axis = {}
+
+		# Init aruco aoi
+		self.aruco_aoi = {}
+
+		# Update all attributes from arguments
+		self.__dict__.update(kwargs)
+
+		# Convert aruco axis markers identifier into ARUCO_DICT_NAME#ID string
+		aruco_axis_string = {}
+		for axis_name, markers_id in self.aruco_axis.items():
+
+			# Estimate pose from axis markers
+			aruco_axis_names = []
+			for marker_id in markers_id:
+				aruco_axis_names.append(f'{ar_environment.aruco_detector.dictionary.name}#{marker_id}')
+
+			aruco_axis_string[axis_name] = aruco_axis_names
+
+		self.aruco_axis = aruco_axis_string
+
+		# Preprocess orthogonal projection to speed up further aruco aoi processings
+		self.__orthogonal_projection_cache = self.orthogonal_projection
+
+	@classmethod
+	def from_json(self, json_filepath: str) -> ArSceneType:
+		"""Load ArScene from .json file."""
+
+		with open(json_filepath) as configuration_file:
+
+			return ArScene(**json.load(configuration_file))
+
+	def __str__(self) -> str:
+		"""String display"""
+
+		output = f'ArUcoScene:\n{self.aruco_scene}\n'
+		output += f'AOIScene:\n{self.aoi_scene}\n'
+
+		return output
+
+	@property
+	def orthogonal_projection(self) -> AOI2DSceneType:
+		"""Orthogonal projection of the aoi whole scene."""
+
+		scene_size = self.aoi_scene.size
+
+		# Center, step back and rotate pose to get whole scene into field of view
+		tvec = self.aoi_scene.center*[-1, 1, 0] + [0, 0, scene_size[1]]
+		rvec = numpy.array([[-numpy.pi, 0.0, 0.0]])
+
+		# Edit intrinsic camera parameter to capture whole scene
+		K = numpy.array([[scene_size[1]/scene_size[0], 0.0, 0.5], [0.0, 1., 0.5], [0.0, 0.0, 1.0]])
+
+		return self.aoi_scene.project(tvec, rvec, K)
+
+	def estimate_pose(self, detected_markers) -> Tuple[numpy.array, numpy.array, dict]:
+		"""Estimate scene pose from detected ArUco markers.
+
+		* **Returns:** 
+			- scene translation vector
+			- scene rotation matrix
+			- dict of markers used to estimate the pose
+		"""
+
+		# Pose estimation fails when no marker is detected
+		if len(detected_markers) == 0:
+			
+			raise PoseEstimationFailed('No marker detected')
+
+		scene_markers, _ = self.aruco_scene.filter_markers(detected_markers)
+
+		# Pose estimation fails when no marker belongs to the scene
+		if len(scene_markers) == 0:
+
+			raise PoseEstimationFailed('No marker belongs to the scene')
+
+		# Estimate scene pose from unique marker transformations
+		elif len(scene_markers) == 1:
+
+			tvec, rmat = self.aruco_scene.estimate_pose_from_any_markers(scene_markers)
+
+			return tvec, rmat, scene_markers
+
+		# Try to estimate scene pose from 3 markers defining an orthogonal axis
+		elif len(scene_markers) >= 3 and len(self.aruco_axis) > 0:
+
+			for axis_name, markers_names in self.aruco_axis.items():
+
+				try:
+
+					axis_markers = []
+					for name in markers_names:
+						axis_markers.append((name, scene_markers[name]))
+
+					tvec, rmat = self.aruco_scene.estimate_pose_from_axis_markers(axis_markers)
+
+					return tvec, rmat, axis_markers
+
+				except:
+					pass
+
+			raise PoseEstimationFailed('No marker axis')
+
+		# Otherwise, check markers consistency
+		consistent_markers, unconsistent_markers, unconsistencies = self.aruco_scene.check_markers_consistency(scene_markers, self.angle_tolerance, self.distance_tolerance)
+
+		# Pose estimation fails when no marker passes consistency checking
+		if len(consistent_markers) == 0:
+
+			raise PoseEstimationFailed('Unconsistent marker poses', unconsistencies)
+
+		# Otherwise, estimate scene pose from all markers transformations
+		tvec, rmat = self.aruco_scene.estimate_pose_from_any_markers(consistent_markers)
+
+		return tvec, rmat, consistent_markers
+
+	def project(self, tvec: numpy.array, rvec: numpy.array, visual_hfov=0) -> AOI2DSceneType:
+		"""Project AOI scene according estimated pose and optional horizontal field of view clipping angle.	
+
+		* **Arguments:** 
+			- translation vector
+			- rotation vector
+			- horizontal field of view clipping angle
+		"""
+
+		# Clip AOI out of the visual horizontal field of view (optional)
+		if visual_hfov > 0:
+
+			# Transform scene into camera referential
+			aoi_scene_camera_ref = self.aoi_scene.transform(tvec, rvec)
+
+			# Get aoi inside vision cone field 
+			cone_vision_height_cm = 200 # cm
+			cone_vision_radius_cm = numpy.tan(numpy.deg2rad(visual_hfov / 2)) * cone_vision_height_cm
+
+			_, aoi_outside = aoi_scene_camera_ref.vision_cone(cone_vision_radius_cm, cone_vision_height_cm)
+
+			# Keep only aoi inside vision cone field
+			aoi_scene_copy = self.aoi_scene.copy(exclude=aoi_outside.keys())
+
+		else:
+
+			aoi_scene_copy = self.aoi_scene.copy()
+
+		aoi_scene_projection = aoi_scene_copy.project(tvec, rvec, self.__ar_environment.aruco_detector.camera.K)
+
+		# Warn user when the projected scene is empty
+		if len(aoi_scene_projection) == 0:
+
+			raise SceneProjectionFailed('AOI projection is empty')
+
+		return aoi_scene_projection
+
+	def build_aruco_aoi_scene(self, detected_markers) -> AOI2DSceneType:
+		"""Build AOI scene from ArUco markers into frame as defined in aruco_aoi dictionary."""
+
+		# AOI projection fails when no marker is detected
+		if len(detected_markers) == 0:
+			
+			raise SceneProjectionFailed('No marker detected')
+
+		aruco_aoi_scene = {}
+
+		for aruco_aoi_name, aoi in self.aruco_aoi.items():
+
+			# Each aoi's corner is defined by a marker's corner
+			aoi_corners = []
+			for corner in ["upper_left_corner", "upper_right_corner", "lower_right_corner", "lower_left_corner"]:
+
+				marker_identifier = aoi[corner]["marker_identifier"]
+
+				try:
+
+					aoi_corners.append(detected_markers[marker_identifier].corners[0][aoi[corner]["marker_corner_index"]])
+
+				except Exception as e:
+					
+					raise SceneProjectionFailed(f'Missing marker #{e} to build ArUco AOI scene')
+
+			aruco_aoi_scene[aruco_aoi_name] = AOIFeatures.AreaOfInterest(aoi_corners)
+
+			# Then each inner aoi is projected from the current aruco aoi
+			for inner_aoi_name, inner_aoi in self.aoi_scene.items():
+
+				if aruco_aoi_name != inner_aoi_name:
+
+					aoi_corners = [numpy.array(aruco_aoi_scene[aruco_aoi_name].outter_axis(inner)) for inner in self.__orthogonal_projection_cache[inner_aoi_name]]
+					aruco_aoi_scene[inner_aoi_name] = AOIFeatures.AreaOfInterest(aoi_corners)
+
+		return AOI2DScene.AOI2DScene(aruco_aoi_scene)
+
+	def draw_axis(self, frame):
+		"""Draw scene axis into frame."""
+
+		self.aruco_scene.draw_axis(frame, self.__ar_environment.aruco_detector.camera.K, self.__ar_environment.aruco_detector.camera.D)
+
+	def draw_places(self, frame):
+		"""Draw scene places into frame."""
+
+		self.aruco_scene.draw_places(frame, self.__ar_environment.aruco_detector.camera.K, self.__ar_environment.aruco_detector.camera.D)
+
+