Moving aruco marker tracking and pose estimation into a dedicated method.

1 files changed, 43 insertions, 31 deletions

diff --git a/src/argaze/ArScene.py b/src/argaze/ArScene.py
index f63f170..273dd7b 100644
--- a/src/argaze/ArScene.py
+++ b/src/argaze/ArScene.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 
-from typing import TypeVar
+from typing import TypeVar, Tuple
 from dataclasses import dataclass, field
 import json
 import os
@@ -13,6 +13,9 @@ import numpy
 ArSceneType = TypeVar('ArScene', bound="ArScene")
 # Type definition for type annotation convenience
 
+AOI2DSceneType = TypeVar('AOI2DScene', bound="AOI2DScene")
+# Type definition for type annotation convenience
+
 class PoseEstimationFailed(Exception):
 	"""Exception raised by ArScene project method when the pose can't be estimated due to unconsistencies."""
 
@@ -112,56 +115,65 @@ class ArScene():
 
 		return output
 
-	def project(self, frame, consistent_markers_number:int = 1, visual_hfov=0, pre_tracked_markers=False):
-		"""Project ArScene into frame."""
+	def estimate_pose(self, frame) -> Tuple[numpy.array, numpy.array, list]:
+		"""Estimate scene pose from ArUco markers into frame.
 
-		# Track markers with pose estimation if it not already done
-		if not pre_tracked_markers:
+		* **Returns:** 
+            - scene translation vector
+            - scene rotation matrix
+            - list of all tracked markers considered as consistent and used to estimate the pose
+        """
 
-			self.aruco_tracker.track(frame)
+		self.aruco_tracker.track(frame)
 
-		# When no marker is detected, no AOI scene projection can't be done
+		# When no marker is detected, no pose estimation can't be done
 		if len(self.aruco_tracker.tracked_markers) == 0:
 			
 			raise PoseEstimationFailed('No marker detected')
 
-		# Estimate set pose from tracked markers
+		# Estimate scene pose from tracked markers
 		tvec, rmat, success, consistent_markers, unconsistencies = self.aruco_scene.estimate_pose(self.aruco_tracker.tracked_markers)
 
-		# When pose estimation fails, ignore AOI scene projection
 		if not success:
-
 			raise PoseEstimationFailed('Unconsistent marker poses', unconsistencies)
 
-		# Consider pose estimation only if theer is a given number of consistent markers at least
-		elif len(consistent_markers) >= consistent_markers_number:
+		return tvec, rmat, consistent_markers
+
+	def project(self, tvec, rmat, visual_hfov=0) -> AOI2DSceneType:
+		"""Project AOI scene into frame according estimated pose."""
 
-			# Clip AOI out of the visual horizontal field of view (optional)
-			if visual_hfov > 0:
+		# 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, rmat)
+			# Transform scene into camera referential
+			aoi_scene_camera_ref = self.aoi_scene.transform(tvec, rmat)
 
-				# 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
+			# 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)
+			_, 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:
 
-				# Keep only aoi inside vision cone field
-				aoi_scene_copy = self.aoi_scene.copy(exclude=aoi_outside.keys())
+			aoi_scene_copy = self.aoi_scene.copy()
 
-			else:
+		# DON'T APPLY CAMERA DISTORSION : it projects points which are far from the frame into it
+		# This hack isn't realistic but as the gaze will mainly focus on centered AOI, where the distorsion is low, it is acceptable.
+		aoi_scene_projection = aoi_scene_copy.project(tvec, rmat, self.aruco_camera.K)
 
-				aoi_scene_copy = self.aoi_scene.copy()
+		# Warn user when the projected scene is empty
+		if len(aoi_scene_projection) == 0:
 
-			# DON'T APPLY CAMERA DISTORSION : it projects points which are far from the frame into it
-			# This hack isn't realistic but as the gaze will mainly focus on centered AOI, where the distorsion is low, it is acceptable.
-			aoi_scene_projection = aoi_scene_copy.project(tvec, rmat, self.aruco_camera.K)
+			raise SceneProjectionFailed('AOI projection is empty')
 
-			# Warn user when the projected scene is empty
-			if len(aoi_scene_projection) == 0:
+		return aoi_scene_projection
 
-				raise SceneProjectionFailed('AOI projection is empty')
+	def draw_axis(self, frame):
+		"""Draw scene axis into frame."""
 
-			return aoi_scene_projection
+		self.aruco_scene.draw_axis(frame, self.aruco_camera.K, self.aruco_camera.D)
+