diff options
| -rw-r--r-- | src/argaze/ArScene.py | 74 |
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) + |