forked from Telos4/RoboRally
general refactoring
This commit is contained in:
parent
dd162018d8
commit
7ded3bee79
|
@ -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
|
||||
|
|
Loading…
Reference in New Issue
Block a user