Defining proper exceptions during projection function. Allowing to pre-process tracking.

1 files changed, 27 insertions, 12 deletions

diff --git a/src/argaze/ArScene.py b/src/argaze/ArScene.py
index 29c8e95..1fea998 100644
--- a/src/argaze/ArScene.py
+++ b/src/argaze/ArScene.py
@@ -13,6 +13,22 @@ import numpy
 ArSceneType = TypeVar('ArScene', bound="ArScene")
 # 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."""
+
+	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 environnement thanks to ArUco markers and project it onto incoming frames."""
@@ -92,19 +108,22 @@ class ArScene():
 		output += f'\n\nArUcoTracker tracking data: {self.aruco_tracker.tracking_data}'
 		output += f'\n\nArUcoScene: {self.aruco_scene}'
 		output += f'\n\nAOIScene: {self.aoi_scene}'
+		output += '\n'
 
 		return output
 
-	def project(self, frame, consistent_markers_number:int = 1, visual_hfov=0):
+	def project(self, frame, consistent_markers_number:int = 1, visual_hfov=0, pre_tracked_markers=False):
 		"""Project ArScene into frame."""
 
-		# Track markers with pose estimation and draw them
-		self.aruco_tracker.track(frame)
+		# Track markers with pose estimation if it not already done
+		if not pre_tracked_markers:
+
+			self.aruco_tracker.track(frame)
 
 		# When no marker is detected, no AOI scene projection can't be done
 		if len(self.aruco_tracker.tracked_markers) == 0:
 			
-			raise UserWarning('No marker detected')
+			raise PoseEstimationFailed('No marker detected')
 
 		# Estimate set pose from tracked markers
 		tvec, rvec, success, consistent_markers, unconsistencies = self.aruco_scene.estimate_pose(self.aruco_tracker.tracked_markers)
@@ -112,11 +131,7 @@ class ArScene():
 		# When pose estimation fails, ignore AOI scene projection
 		if not success:
 
-			# DEBUG: print unconsistencies distances or angles
-			for key, value in unconsistencies.items():
-				print(f'Unconsistent {key}: {value}')
-
-			raise UserWarning('Pose estimation fails')
+			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:
@@ -144,9 +159,9 @@ class ArScene():
 			# 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, rvec, self.aruco_camera.K)
 
-			# Warn user when the merged scene is empty
-			if len(aoi_scene_projection.keys()) == 0:
+			# Warn user when the projected scene is empty
+			if len(aoi_scene_projection) == 0:
 
-				raise UserWarning('AOI projection is empty')
+				raise SceneProjectionFailed('AOI projection is empty')
 
 			return aoi_scene_projection