path: root/src/argaze/ArFeatures.py

#!/usr/bin/env python

"""Manage AR environement assets."""

__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
import importlib

from argaze import DataStructures, GazeFeatures
from argaze.ArUcoMarkers import *
from argaze.AreaOfInterest import *
from argaze.GazeAnalysis import *

import numpy
import cv2

ArEnvironmentType = TypeVar('ArEnvironment', bound="ArEnvironment")
# Type definition for type annotation convenience

ArSceneType = TypeVar('ArScene', bound="ArScene")
# Type definition for type annotation convenience

ArScreenType = TypeVar('ArScreen', bound="ArScreen")
# Type definition for type annotation convenience

@dataclass
class ArEnvironment():
	"""
	Define Augmented Reality environment based on ArUco marker detection.

	Parameters:
		name: Environment name
		aruco_detector: ArUco detector
		scenes: All environment scenes
	"""

	name: str 
	aruco_detector: ArUcoDetector.ArUcoDetector = field(default_factory=ArUcoDetector.ArUcoDetector)
	scenes: dict = field(default_factory=dict)

	def __post_init__(self):

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

		# Init AOI scene projections
		self.__aoi_2d_scenes = {}

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

		Parameters:
			json_filepath: path to 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')

			try:
				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 optic_parameters value type
				optic_parameters_value = new_detector_data.pop('optic_parameters')

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

					optic_parameters_value = os.path.join(working_directory, optic_parameters_value)
					new_optic_parameters = ArUcoOpticCalibrator.OpticParameters.from_json(optic_parameters_value)

				# dict:
				else:

					new_optic_parameters = ArUcoOpticCalibrator.OpticParameters(**optic_parameters_value)

				# Check detector parameters value type
				detector_parameters_value = new_detector_data.pop('parameters')

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

					detector_parameters_value = os.path.join(working_directory, detector_parameters_value)
					new_aruco_detector_parameters = ArUcoDetector.DetectorParameters.from_json(detector_parameters_value)

				# dict:
				else:

					new_aruco_detector_parameters = ArUcoDetector.DetectorParameters(**detector_parameters_value)
				
				new_aruco_detector = ArUcoDetector.ArUcoDetector(new_aruco_dictionary, new_marker_size, new_optic_parameters, new_aruco_detector_parameters)

			except KeyError:

				new_aruco_detector = None

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

				new_aruco_scene = None
				new_aoi_scene = None
				
				try:

					# 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)

				except KeyError:

						new_aruco_scene = None

				# Check aoi_3d_scene value type
				aoi_3d_scene_value = scene_data.pop('aoi_3d_scene')

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

					obj_filepath = os.path.join(working_directory, aoi_3d_scene_value)
					new_aoi_3d_scene = AOI3DScene.AOI3DScene.from_obj(obj_filepath)

				# dict:
				else:

					new_aoi_3d_scene = AOI3DScene.AOI3DScene(aoi_3d_scene_value)

				# Build screens
				new_screens = {}
				for screen_name, screen_data in scene_data.pop('screens').items():

					new_screen_size = screen_data.pop('size')

					# Load background image
					try:

						new_screen_background_value = screen_data.pop('background')
						new_screen_background = cv2.imread(os.path.join(working_directory, new_screen_background_value))
						new_screen_background = cv2.resize(new_screen_background, dsize=(new_screen_size[0], new_screen_size[1]), interpolation=cv2.INTER_CUBIC)

					except:

						new_screen_background = numpy.zeros((new_screen_size[1], new_screen_size[0], 3)).astype(numpy.uint8)

					# Load gaze movement identifier
					try:

						gaze_movement_identifier_value = screen_data.pop('gaze_movement_identifier')

						gaze_movement_identifier_type = gaze_movement_identifier_value['type']
						gaze_movement_identifier_parameters = gaze_movement_identifier_value['parameters']

						gaze_movement_identifier_module = importlib.import_module(f'argaze.GazeAnalysis.{gaze_movement_identifier_type}')
						gaze_movement_identifier = gaze_movement_identifier_module.GazeMovementIdentifier(**gaze_movement_identifier_parameters)

					except:

						gaze_movement_identifier = None

					# Append new screen
					new_screens[screen_name] = ArScreen.from_scene(new_aoi_3d_scene, screen_name, new_screen_size, new_screen_background, gaze_movement_identifier, **screen_data)

				# Append new scene
				new_scenes[scene_name] = ArScene(new_aruco_scene, new_aoi_3d_scene, new_screens, **scene_data)

			return ArEnvironment(new_name, new_aruco_detector, new_scenes)

	def __str__(self) -> str:
		"""
		Returns:
			String representation
		"""

		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

	@property
	def image(self):
		"""Get ArUco detection visualisation and scenes projections."""

		# Draw detected markers
		self.aruco_detector.draw_detected_markers(self.__image)

		# Draw each AOI scene
		for scene_name, aoi_2d_scene in self.__aoi_2d_scenes.items():

			# Draw AOI scene projection
			aoi_2d_scene.draw(self.__image, color=(255, 255, 255))

		return self.__image

	@property
	def screens(self):
		"""Iterate over all environment screens"""

		# For each scene
		for scene_name, scene in self.scenes.items():

			# For each screen
			for screen_name, screen in scene.screens.items():

				yield scene_name, screen_name, screen

	def detect_and_project(self, image: numpy.array) -> dict:
		"""Detect environment aruco markers from image and project scenes."""

		self.__image = image

		# Detect aruco markers
		self.aruco_detector.detect_markers(self.__image)
		
		# Project each AOI scene
		self.__aoi_2d_scenes = {}
		for scene_name, scene in self.scenes.items():

			# Project scene
			try:

				# Try to build AOI scene from detected ArUco marker corners
				self.__aoi_2d_scenes[scene_name] = scene.build_aruco_aoi_scene(self.aruco_detector.detected_markers)

			except:

				# Estimate scene markers poses
				self.aruco_detector.estimate_markers_pose(scene.aruco_scene.identifiers)
				
				# Estimate scene pose from detected scene markers
				tvec, rmat, _, _ = scene.estimate_pose(self.aruco_detector.detected_markers)
				
				# Project AOI scene into video image according estimated pose
				self.__aoi_2d_scenes[scene_name] = scene.project(tvec, rmat)

	def look(self, timestamp: int|float, gaze_position: GazeFeatures.GazePosition):
		"""Project gaze position into environment at particular time."""

		# For each aoi scene projection
		for scene_name, scene in self.scenes.items():

			try:

				aoi_2d_scene = self.__aoi_2d_scenes[scene_name]

				# For each scene screens
				for screen_name, screen in scene.screens.items():

					# TODO: Add option to use gaze precision circle
					if aoi_2d_scene[screen.name].contains_point(gaze_position.value):

						inner_x, inner_y = self.__aoi_2d_scenes[scene_name][screen.name].clockwise().inner_axis(gaze_position.value)

						# QUESTION: How to project gaze precision?
						inner_gaze_position = GazeFeatures.GazePosition((inner_x, inner_y))

						screen.look(timestamp, inner_gaze_position * screen.size)

			# Ignore missing aoi scene projection 
			except KeyError:

				pass

	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.

	Parameters:
		aruco_scene: ArUco markers 3D scene description used to estimate scene pose from detected markers: see [estimate_pose][argaze.ArFeatures.ArScene.estimate_pose] function below.

		aoi_3d_scene: AOI 3D scene description that will be projected onto estimated scene once its pose will be estimated : see [project][argaze.ArFeatures.ArScene.project] function below.

		screens: All scene screens

		aruco_axis: 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][argaze.ArFeatures.ArScene.estimate_pose] function when at least 3 markers are detected.

		aruco_aoi: Optional dictionary of AOI defined by list of markers identifier and markers corners index tuples: see [build_aruco_aoi_scene][argaze.ArFeatures.ArScene.build_aruco_aoi_scene] function below.

		angle_tolerance: Optional angle error tolerance to validate marker pose in degree used into [estimate_pose][argaze.ArFeatures.ArScene.estimate_pose] function.

		distance_tolerance: Optional distance error tolerance to validate marker pose in centimeter used into [estimate_pose][argaze.ArFeatures.ArScene.estimate_pose] function.
	"""

	aruco_scene: ArUcoScene.ArUcoScene = field(default_factory=ArUcoScene.ArUcoScene)
	aoi_3d_scene: AOI3DScene.AOI3DScene = field(default_factory=AOI3DScene.AOI3DScene)
	screens: dict = field(default_factory=dict)
	aruco_axis: dict = field(default_factory=dict)
	aruco_aoi: dict = field(default_factory=dict)
	angle_tolerance: float = field(default=0.)
	distance_tolerance: float = field(default=0.)

	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.aoi_3d_scene.orthogonal_projection

		# Setup screens scene after screen creation
		for name, screen in self.screens.items():
			screen._scene = self

	def __str__(self) -> str:
		"""
		Returns:
			String representation
		"""

		output = f'ArEnvironment:\n{self._environment.name}\n'
		output += f'ArUcoScene:\n{self.aruco_scene}\n'
		output += f'AOI3DScene:\n{self.aoi_3d_scene}\n'

		return output

	def estimate_pose(self, detected_markers) -> Tuple[numpy.array, numpy.array, str, 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: float = 0.) -> AOI2DScene.AOI2DScene:
		"""Project AOI scene according estimated pose and optional horizontal field of view clipping angle.	

		Parameters:
			tvec: translation vector
			rvec: rotation vector
			visual_hfov: 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_3d_scene_camera_ref = self.aoi_3d_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_3d_scene_camera_ref.vision_cone(cone_vision_radius_cm, cone_vision_height_cm)

			# Keep only aoi inside vision cone field
			aoi_3d_scene_copy = self.aoi_3d_scene.copy(exclude=aoi_outside.keys())

		else:

			aoi_3d_scene_copy = self.aoi_3d_scene.copy()

		aoi_2d_scene = aoi_3d_scene_copy.project(tvec, rvec, self._environment.aruco_detector.optic_parameters.K)

		# Warn user when the projected scene is empty
		if len(aoi_2d_scene) == 0:

			raise SceneProjectionFailed('AOI projection is empty')

		return aoi_2d_scene

	def build_aruco_aoi_scene(self, detected_markers) -> AOI2DScene.AOI2DScene:
		"""
		Build AOI scene from detected ArUco markers as defined in aruco_aoi dictionary.

		Returns:
			built AOI 2D scene
		"""

		# 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_3d_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, image: numpy.array):
		"""
		Draw scene axis into image.
		
		Parameters:
			image: where to draw
		"""

		self.aruco_scene.draw_axis(image, self._environment.aruco_detector.optic_parameters.K, self._environment.aruco_detector.optic_parameters.D)

	def draw_places(self, image: numpy.array):
		"""
		Draw scene places into image.

		Parameters:
			image: where to draw
		"""

		self.aruco_scene.draw_places(image, self._environment.aruco_detector.optic_parameters.K, self._environment.aruco_detector.optic_parameters.D)

@dataclass
class ArScreen():
	"""
	Define Augmented Reality screen as an AOI2DScene made from a projected then reframed parent AOI3DScene.

	Parameters:
		name: name of the screen
		size: screen dimension in pixel.
		background: image to draw behind
		aoi_2d_scene: AOI 2D scene description ... : see [orthogonal_projection][argaze.ArFeatures.ArScene.orthogonal_projection] and [reframe][argaze.AreaOfInterest.AOI2DScene.reframe] functions.
	"""

	name: str
	size: tuple[int] = field(default=(1, 1))
	background: numpy.array = field(default_factory=numpy.array)
	aoi_2d_scene: AOI2DScene.AOI2DScene = field(default_factory=AOI2DScene.AOI2DScene)
	gaze_movement_identifier: GazeFeatures.GazeMovementIdentifier = field(default_factory=GazeFeatures.GazeMovementIdentifier)
	scan_path: GazeFeatures.ScanPath = field(default_factory=GazeFeatures.ScanPath)
	aoi_scan_path: GazeFeatures.AOIScanPath = field(default_factory=GazeFeatures.AOIScanPath)
	heatmap: AOIFeatures.Heatmap = field(default_factory=AOIFeatures.Heatmap)

	def __post_init__(self):

		# Define scene attribute: it will be setup by parent scene later
		self._scene = None

		# Init gaze data
		self.__gaze_position = GazeFeatures.UnvalidGazePosition()

		if self.heatmap:

			self.heatmap.init()

	@classmethod
	def from_scene(self, aoi_3d_scene, aoi_name, size, background, gaze_movement_identifier, scan_path: bool = False, aoi_scan_path: bool = False, heatmap: bool = False) -> ArScreenType:

		aoi_2d_scene = aoi_3d_scene.orthogonal_projection.reframe(aoi_name, size)

		return ArScreen(aoi_name, \
						size, \
						background, \
						aoi_2d_scene, \
						gaze_movement_identifier, \
						GazeFeatures.ScanPath() if scan_path else None, \
						GazeFeatures.AOIScanPath(aoi_2d_scene.keys()) if aoi_scan_path else None, \
						AOIFeatures.Heatmap(size) if heatmap else None \
						)

	@property
	def current_gaze_position(self):
		"""Get current gaze position on screen."""

		return self.__gaze_position

	@property
	def current_gaze_movement(self):
		"""Get current gaze movement on screen."""

		# Check current screen fixation
		current_fixation = self.gaze_movement_identifier.current_fixation

		if current_fixation.valid:

			return current_fixation

		# Check current screen saccade
		current_saccade = self.gaze_movement_identifier.current_saccade

		if current_saccade.valid:

			return current_saccade

		return GazeFeatures.UnvalidGazeMovement()

	def look(self, timestamp: int|float, inner_gaze_position: GazeFeatures.GazePosition):
		""" 
		
		GazeFeatures.AOIScanStepError
		"""

		self.__gaze_position = inner_gaze_position

		# Identify gaze movement
		if self.gaze_movement_identifier:

			# Identify gaze movement
			gaze_movement = self.gaze_movement_identifier.identify(timestamp, self.__gaze_position)

			# QUESTION: How to notify new gaze movement?

			if GazeFeatures.is_fixation(gaze_movement):

				# Does the fixation match an AOI?
				look_at = self.name
				for name, aoi in self.aoi_2d_scene.items():

					_, _, circle_ratio = aoi.circle_intersection(gaze_movement.focus, self.gaze_movement_identifier.deviation_max_threshold)

					if circle_ratio > 0.25:

						if name != self.name:

							look_at = name
							break

				# Append fixation to scan path
				if self.scan_path:

					self.scan_path.append_fixation(timestamp, gaze_movement)

				# Append fixation to aoi scan path
				if self.aoi_scan_path:

					self.__aoi_scan_step = self.aoi_scan_path.append_fixation(timestamp, gaze_movement, look_at)

					# QUESTION: How to notify new step?

			elif GazeFeatures.is_saccade(gaze_movement):

				# Append saccade to scan path
				if self.scan_path:
					
					self.__scan_step = self.scan_path.append_saccade(timestamp, gaze_movement)

					# QUESTION: How to notify new step?

				# Append saccade to aoi scan path
				if self.aoi_scan_path:

					self.aoi_scan_path.append_saccade(timestamp, gaze_movement)	

		# Update heatmap
		if self.heatmap:

			self.heatmap.update(self.__gaze_position.value, sigma=0.05)