general refactoring

This commit is contained in:
Simon Pirkelmann 2020-10-24 21:08:20 +02:00
parent dd162018d8
commit 7ded3bee79

View File

@ -9,6 +9,7 @@ from queue import Queue
import aruco
class ArucoEstimator:
corner_marker_ids = {
'a': 0,
@ -17,8 +18,6 @@ class ArucoEstimator:
'd': 3
}
angles = []
corner_estimates = {
'a': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
'b': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
@ -26,14 +25,16 @@ class ArucoEstimator:
'd': {'pixel_coordinate': None, 'x': None, 'y': None, 'n_estimates': 0},
}
def __init__(self, robot_marker_ids=None, use_realsense=True, grid_columns=8, grid_rows=8, draw_marker_coordinate_system=False):
def __init__(self, robot_marker_ids=None, use_realsense=True, grid_columns=8, grid_rows=8,
draw_marker_coordinate_system=False):
self.grid_columns = grid_columns
self.grid_rows = grid_rows
if robot_marker_ids is None:
robot_marker_ids = []
self.robot_marker_ids = robot_marker_ids
self.robot_marker_estimates = dict([(id, None) for id in self.robot_marker_ids])
self.robot_marker_estimates = dict([(marker_id, {'t': None, 'x': None, 'y': None, 'angle': None})
for marker_id in self.robot_marker_ids])
self.event_queue = Queue()
self.draw_marker_coordinate_system = draw_marker_coordinate_system
@ -165,8 +166,8 @@ class ArucoEstimator:
aruco.CvDrawingUtils.draw3dAxis(color_image, self.camparam, marker.Rvec, marker.Tvec, .1)
# store data
for id, data in detected_marker_data.items():
self.update_estimate(id, data['marker_center'], data['Rvec'], data['Tvec'], t_image)
for marker_id, data in detected_marker_data.items():
self.update_estimate(marker_id, data['marker_center'], data['Rvec'], data['Tvec'], t_image)
# draw grid
color_image = self.draw_grid_lines(color_image, detected_marker_data)
@ -186,14 +187,14 @@ class ArucoEstimator:
# Stop streaming
self.pipeline.stop()
def update_estimate(self, id, pixel_coord_center, rvec, tvec, t_image):
def update_estimate(self, marker_id, pixel_coord_center, rvec, tvec, t_image):
# update the marker estimate with new data
if id in self.corner_marker_ids.values():
if marker_id in self.corner_marker_ids.values():
# for corner markers we compute the mean of all measurements s.t. the position stabilizes over time
# (we assume the markers do not move)
# get corresponding corner to the detected marker
corner = next(filter(lambda key: self.corner_marker_ids[key] == id, self.corner_marker_ids.keys()))
corner = next(filter(lambda key: self.corner_marker_ids[key] == marker_id, self.corner_marker_ids.keys()))
old_estimate_x = self.corner_estimates[corner]['x']
old_estimate_y = self.corner_estimates[corner]['y']
n_estimates = self.corner_estimates[corner]['n_estimates']
@ -212,7 +213,7 @@ class ArucoEstimator:
self.corner_estimates[corner]['n_estimates'] = n_estimates + 1
self.corner_estimates[corner]['pixel_coordinate'] = pixel_coord_center
elif id in self.robot_marker_ids:
elif marker_id in self.robot_marker_ids:
# for robot markers we extract x and y position as well as the angle
# here we could also implement a filter
x = tvec[0][0]
@ -224,8 +225,7 @@ class ArucoEstimator:
_, _, _, _, _, _, euler_angles = cv2.decomposeProjectionMatrix(pose_mat)
angle = -euler_angles[2][0] * np.pi / 180.0
self.angles.append(angle)
self.robot_marker_estimates[id] = {'t': t_image, 'x': x, 'y': y, 'angle': angle}
self.robot_marker_estimates[marker_id] = {'t': t_image, 'x': x, 'y': y, 'angle': angle}
def all_corners_detected(self):
# checks if all corner markers have been detected at least once
@ -233,7 +233,7 @@ class ArucoEstimator:
def all_robots_detected(self):
# checks if all robot markers have been detected at least once
return not any([estimate is None for estimate in self.robot_marker_estimates.values()])
return not any([estimate['t'] is None for estimate in self.robot_marker_estimates.values()])
def get_pos_from_grid_point(self, x, y, orientation=None):
"""
@ -244,8 +244,8 @@ class ArucoEstimator:
:return: numpy array with corresponding real world x- and y-position
if orientation was specified the array also contains the matching angle for the orientation
"""
assert x >= 0 and x < self.grid_columns
assert y >= 0 and y < self.grid_rows
assert 0 <= x < self.grid_columns
assert 0 <= y < self.grid_rows
assert self.all_corners_detected()
# compute column line
@ -346,29 +346,30 @@ class ArucoEstimator:
return frame
def get_robot_state_estimate(self, id):
if id in self.robot_marker_estimates:
if self.robot_marker_estimates[id] is not None:
return np.array([self.robot_marker_estimates[id]['x'], self.robot_marker_estimates[id]['y'],
self.robot_marker_estimates[id]['angle']])
def get_robot_state_estimate(self, marker_id):
if marker_id in self.robot_marker_estimates:
if self.robot_marker_estimates[marker_id]['t'] is not None:
return np.array([self.robot_marker_estimates[marker_id]['x'],
self.robot_marker_estimates[marker_id]['y'],
self.robot_marker_estimates[marker_id]['angle']])
else:
print(f"error: no estimate available for robot {id}")
print(f"error: no estimate available for robot {marker_id}")
return None
else:
print(f"error: invalid robot id {id}")
print(f"error: invalid robot id {marker_id}")
return None
def draw_robot_pos(self, frame, detected_marker_data):
# draws information about the robot positions onto the given frame
robot_corners_pixel_coords = {}
for id, estimate in self.robot_marker_estimates.items():
if id in detected_marker_data.keys():
robot_corners_pixel_coords[id] = tuple(detected_marker_data[id]['center'])
for marker_id, estimate in self.robot_marker_estimates.items():
if marker_id in detected_marker_data.keys():
robot_corners_pixel_coords[marker_id] = tuple(detected_marker_data[marker_id]['center'])
for id, coord in robot_corners_pixel_coords.items():
x = self.robot_marker_estimates[id]['x']
y = self.robot_marker_estimates[id]['y']
angle = self.robot_marker_estimates[id]['angle']
for marker_id, coord in robot_corners_pixel_coords.items():
x = self.robot_marker_estimates[marker_id]['x']
y = self.robot_marker_estimates[marker_id]['y']
angle = self.robot_marker_estimates[marker_id]['angle']
frame = cv2.putText(frame, "pos = ({:5.3f}, {:5.3f}), ang = {:5.3f}".format(x, y, angle), coord,
cv2.FONT_HERSHEY_SIMPLEX, 0.50, (0, 255, 0))
return frame