in robots/LoCoBot/locobot_calibration/scripts/collect_calibration_data.py [0:0]
def collect_data(self, run_dir):
"""
Place the arm at a reasonable location, and take multiple images with
the camera.
:param run_dir: directory to store collected data
:type run_dir: string
"""
mkdir_if_missing(os.path.join(run_dir, "images"))
states, ar_trans, ar_quat = [], [], []
ar_img_loc = []
rgb_tf_tree_trans, rgb_tf_tree_quat, rgb_tf_tree_matrix = [], [], []
ar_tf_tree_trans, ar_tf_tree_quat, ar_tf_tree_matrix = [], [], []
arm_pose_ids = []
i = 0
N = len(ARM_POSES)
for arm_pose_id in range(N):
arm_pose = ARM_POSES[arm_pose_id]
self.move_arm(arm_pose)
# Code to compute camera poses.
camera_pan, camera_tilt = self.compute_camera_pose()
# Set pans and tilts.
pans = np.linspace(-PAN_LIM, PAN_LIM, FLAGS.num_angles)
pans = pans + camera_pan
pans = np.unique(pans)
tilts = np.linspace(-TILT_LIM, TILT_LIM, FLAGS.num_angles)
tilts = tilts + camera_tilt
tilts = np.minimum(tilts, MAX_TILT)
tilts = np.unique(tilts)
for pan in pans:
# Send pitch to be vertical
self.bot.camera.set_pan(pan)
for tilt in tilts:
self.bot.camera.set_tilt(tilt)
rospy.sleep(1.5)
states.append(
[self.bot.camera.get_pan(), self.bot.camera.get_tilt()]
)
img = self.bot.camera.get_rgb()
rgb_chain, T, TM = self.get_kinematic_chain(
"camera_color_optical_frame", "base_link", "base_link"
)
assert "camera_link" in rgb_chain
T_trans = [t[0] for t in T]
T_quat = [t[1] for t in T]
rgb_tf_tree_trans.append(T_trans)
rgb_tf_tree_quat.append(T_quat)
rgb_tf_tree_matrix.append(TM)
ar_chain, T, TM = self.get_kinematic_chain(
"ar_tag", "base_link", "base_link"
)
T_trans = [t[0] for t in T]
T_quat = [t[1] for t in T]
ar_tf_tree_trans.append(T_trans)
ar_tf_tree_quat.append(T_quat)
ar_tf_tree_matrix.append(TM)
position = [np.NaN for _ in range(3)]
orientation = [np.NaN for _ in range(4)]
ar_pose = self.bot.camera.get_ar_tag_pose()
img_pts = [[np.NaN for _ in range(2)] for __ in range(4)]
if ar_pose is not None and len(ar_pose.markers) == 1:
ar_pose_3d = ar_pose.markers[0].pose.pose
position = ar_pose_3d.position
position = [getattr(position, x) for x in ["x", "y", "z"]]
orientation = ar_pose_3d.orientation
orientation = [
getattr(orientation, x) for x in ["x", "y", "z", "w"]
]
img_pts = []
for _ in range(4):
pos_img = getattr(
ar_pose.markers[0], "pose_img_{:d}".format(_)
)
pos_img = pos_img.pose.position
img_pts.append([pos_img.x, pos_img.y])
ar_trans.append(position)
ar_img_loc.append(img_pts)
ar_quat.append(orientation)
file_name = os.path.join(
run_dir, "images", "img_{:04d}.png".format(i)
)
cv2.imwrite(file_name, img[:, :, ::-1])
arm_pose_ids.append(arm_pose_id)
i = i + 1
self.save_data(
run_dir,
rgb_chain,
ar_chain,
states,
ar_trans,
ar_quat,
ar_img_loc,
rgb_tf_tree_trans,
rgb_tf_tree_quat,
rgb_tf_tree_matrix,
ar_tf_tree_trans,
ar_tf_tree_quat,
ar_tf_tree_matrix,
arm_pose_ids,
)