aboutsummaryrefslogtreecommitdiff
path: root/src/argaze/ArUcoMarkers/ArUcoScene.py
blob: 997ad40c2e1465f067a63415ae3e6dbf9e6c0e98 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
#!/usr/bin/env python

"""ArScene based of ArUco markers technology."""

__author__ = "Théo de la Hogue"
__credits__ = []
__copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
__license__ = "BSD"

from typing import TypeVar, Tuple
import json
import os

from argaze import ArFeatures, DataFeatures
from argaze.ArUcoMarkers import ArUcoMarkersGroup
from argaze.AreaOfInterest import AOI2DScene

import cv2
import numpy

ArUcoSceneType = TypeVar('ArUcoScene', bound="ArUcoScene")
# Type definition for type annotation convenience

class ArUcoScene(ArFeatures.ArScene):
	"""
	Define an ArScene based on an ArUcoMarkersGroup description.
	"""
	
	def __init__(self, aruco_markers_group: ArUcoMarkersGroup.ArUcoMarkersGroup, **kwargs):
		""" Initialize ArUcoScene

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

		# Init parent classes
		super().__init__(**kwargs)

		# Init private attribute
		self.__aruco_markers_group = aruco_markers_group

	@property
	def aruco_markers_group(self) -> ArUcoMarkersGroup.ArUcoMarkersGroup:
		"""Get ArUco scene markers group object."""
		return self.__aruco_markers_group

	@classmethod
	def from_dict(self, aruco_scene_data: dict, working_directory: str = None) -> ArUcoSceneType:
		"""
		Load ArUcoScene from dictionary.

		Parameters:
			aruco_scene_data: dictionary
			working_directory: folder path where to load files when a dictionary value is a relative filepath.
		"""

		# Load aruco markers group
		try:

			# Check aruco_markers_group value type
			aruco_markers_group_value = aruco_scene_data.pop('aruco_markers_group')

			# str: relative path to description file
			if type(aruco_markers_group_value) == str:

				filepath = os.path.join(working_directory, aruco_markers_group_value)
				file_format = filepath.split('.')[-1]

				# JSON file format
				if file_format == 'json':

					new_aruco_markers_group = ArUcoMarkersGroup.ArUcoMarkersGroup.from_json(filepath)

				# OBJ file format
				elif file_format == 'obj':

					new_aruco_markers_group = ArUcoMarkersGroup.ArUcoMarkersGroup.from_obj(filepath)

			# dict:
			else:

				new_aruco_markers_group = ArUcoMarkersGroup.ArUcoMarkersGroup(**aruco_markers_group_value)

		except KeyError:

			new_aruco_markers_group = None

		# Load temporary scene from aruco_scene_data then export it as dict
		temp_scene_data = ArFeatures.ArScene.from_dict(aruco_scene_data, working_directory).as_dict()

		# Create new aruco scene using temporary ar scene values
		return ArUcoScene( \
			aruco_markers_group = new_aruco_markers_group, \
			**temp_scene_data \
			)
		
	def __str__(self) -> str:
		"""
		Returns:
			String representation
		"""

		output = output = super().__str__()
		output += f'ArUcoMarkersGroup:\n{self.__aruco_markers_group}\n'

		return output

	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
			dict of markers used to estimate the pose
		"""

		# Pose estimation fails when no marker is detected
		if len(detected_markers) == 0:

			raise ArFeatures.PoseEstimationFailed('No marker detected')

		scene_markers, _ = self.__aruco_markers_group.filter_markers(detected_markers)

		# Pose estimation fails when no marker belongs to the scene
		if len(scene_markers) == 0:

			raise ArFeatures.PoseEstimationFailed('No marker belongs to the scene')

		# Pose estimation fails if only one marker belongs to the scene
		if len(scene_markers) == 1:

			raise ArFeatures.PoseEstimationFailed('Only one marker belongs to the scene')

		# Estimate pose from a markers corners
		success, tvec, rmat = self.__aruco_markers_group.estimate_pose_from_markers_corners(scene_markers, self.parent.aruco_detector.optic_parameters.K, self.parent.aruco_detector.optic_parameters.D)

		if not success:

			raise ArFeatures.PoseEstimationFailed('Can\'t estimate pose from markers corners positions')

		return tvec, rmat, scene_markers

	def draw(self, image: numpy.array, draw_aruco_markers_group: dict = None):
		"""
		Draw scene into image.
		
		Parameters:
			image: where to draw
			draw_aruco_markers_group: ArUcoMarkersGroup.draw parameters (if None, no group drawn)
		"""

		# Draw group if required
		if draw_aruco_markers_group is not None:

			self.__aruco_markers_group.draw(image, self.parent.aruco_detector.optic_parameters.K, self.parent.aruco_detector.optic_parameters.D, **draw_aruco_markers_group)