path: root/src/argaze/ArFeatures.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
import os

from argaze import DataStructures
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 
	"""Environment name"""

	aruco_detector: ArUcoDetector.ArUcoDetector = field(default_factory=ArUcoDetector.ArUcoDetector)
	"""ArUco detector"""

	scenes: dict = field(default_factory=dict)
	"""All environment scenes"""

	def __post_init__(self):

		# Setup scenes environment after environment creation
		for name, scene in self.scenes.items():
			scene._environment = self

	@classmethod
	def from_json(self, json_filepath: str) -> ArSceneType:
		"""Load ArEnvironment from .json file."""

		with open(json_filepath) as configuration_file:

			data = json.load(configuration_file)
			working_directory = os.path.dirname(json_filepath)

			new_name = data.pop('name')
			new_detector_data = data.pop('aruco_detector')

			new_aruco_dictionary = ArUcoMarkersDictionary.ArUcoMarkersDictionary(**new_detector_data.pop('dictionary'))
			new_marker_size = new_detector_data.pop('marker_size')

			# Check aruco_camera value type
			aruco_camera_value = new_detector_data.pop('camera')

			# str: relative path to .json file
			if type(aruco_camera_value) == str:

				aruco_camera_value = os.path.join(working_directory, aruco_camera_value)
				new_aruco_camera = ArUcoCamera.ArUcoCamera.from_json(aruco_camera_value)

			# dict:
			else:

				new_aruco_camera = ArUcoCamera.ArUcoCamera(**aruco_camera_value)

			new_aruco_detecor_parameters = ArUcoDetector.DetectorParameters(**new_detector_data.pop('parameters'))
			new_aruco_detector = ArUcoDetector.ArUcoDetector(new_aruco_dictionary, new_marker_size, new_aruco_camera, new_aruco_detecor_parameters)

			new_scenes = {}
			for scene_name, scene_data in data.pop('scenes').items():

				new_aruco_scene = None
				new_aoi_scene = None
				
				# Check aruco_scene value type
				aruco_scene_value = scene_data.pop('aruco_scene')

				# str: relative path to .obj file
				if type(aruco_scene_value) == str:

					aruco_scene_value = os.path.join(working_directory, aruco_scene_value)
					new_aruco_scene = ArUcoScene.ArUcoScene.from_obj(aruco_scene_value)

				# dict:
				else:

					new_aruco_scene = ArUcoScene.ArUcoScene(**aruco_scene_value)

				# Check aoi_scene value type
				aoi_scene_value = scene_data.pop('aoi_scene')

				# str: relative path to .obj file
				if type(aoi_scene_value) == str:

					obj_filepath = os.path.join(working_directory, aoi_scene_value)
					new_aoi_scene = AOI3DScene.AOI3DScene.from_obj(obj_filepath)

				# dict:
				else:

					new_aoi_scene = AOI3DScene.AOI3DScene(aoi_scene_value)

				new_scenes[scene_name] = ArScene(new_aruco_scene, new_aoi_scene, **scene_data)

			return ArEnvironment(new_name, new_aruco_detector, new_scenes)

	def __str__(self) -> str:
		"""String display"""

		output = f'Name:\n{self.name}\n'
		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 to_json(self, json_filepath):
		"""Save environment to .json file."""

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

			json.dump(self, file, ensure_ascii=False, indent=4, cls=DataStructures.JsonEncoder)

class PoseEstimationFailed(Exception):
	"""Exception raised by ArScene estimate_pose 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 with ArUco markers and AOI scenes."""

	aruco_scene: ArUcoScene.ArUcoScene = field(default_factory=ArUcoScene.ArUcoScene)
	"""ArUco markers 3D scene description used to estimate scene pose from detected markers: see `estimate_pose` function below."""

	aoi_scene: AOI3DScene.AOI3DScene = field(default_factory=AOI3DScene.AOI3DScene)
	"""AOI 3D scene description that will be projected onto estimated scene once its pose will be estimated : see `project` function below."""

	aruco_axis: dict = field(default_factory=dict)
	"""Optional dictionary to define orthogonal axis where each axis is defined by list of 3 markers identifier (first is origin). \
	This pose estimation strategy is used by `estimate_pose` function when at least 3 markers are detected."""

	aruco_aoi: dict = field(default_factory=dict)
	"""Optional dictionary of AOI defined by list of markers identifier and markers corners index tuples: see `build_aruco_aoi_scene` function below."""

	angle_tolerance: float = field(default=0.)
	"""Optional angle error tolerance to validate marker pose in degree used into `estimate_pose` function."""

	distance_tolerance: float = field(default=0.)
	"""Optional distance error tolerance to validate marker pose in centimeter used into `estimate_pose` function."""

	def __post_init__(self):

		# Define environment attribute: it will be setup by parent environment later
		self._environment = None

		# Preprocess orthogonal projection to speed up further aruco aoi processings
		self.__orthogonal_projection_cache = self.orthogonal_projection

	def __str__(self) -> str:
		"""String display"""

		output = f'ArEnvironment:\n{self._environment.name}\n'
		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 whole AOI 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
			- pose estimation strategy
			- 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:

			marker_id, marker = scene_markers.popitem()
			tvec, rmat = self.aruco_scene.estimate_pose_from_single_marker(marker)
			
			return tvec, rmat, 'estimate_pose_from_single_marker', {marker_id: marker}

		# 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, axis_markers in self.aruco_axis.items():

				try:

					origin_marker = scene_markers[axis_markers['origin_marker']]
					horizontal_axis_marker = scene_markers[axis_markers['horizontal_axis_marker']]
					vertical_axis_marker = scene_markers[axis_markers['vertical_axis_marker']]

					tvec, rmat = self.aruco_scene.estimate_pose_from_axis_markers(origin_marker, horizontal_axis_marker, vertical_axis_marker)

					return tvec, rmat, 'estimate_pose_from_axis_markers', {origin_marker.identifier: origin_marker, horizontal_axis_marker.identifier: horizontal_axis_marker, vertical_axis_marker.identifier: vertical_axis_marker}

				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 consistent markers pose
		tvec, rmat = self.aruco_scene.estimate_pose_from_markers(consistent_markers)

		return tvec, rmat, 'estimate_pose_from_markers', 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._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._environment.aruco_detector.camera.K, self._environment.aruco_detector.camera.D)

	def draw_places(self, frame):
		"""Draw scene places into frame."""

		self.aruco_scene.draw_places(frame, self._environment.aruco_detector.camera.K, self._environment.aruco_detector.camera.D)