def __init__(self, name):
"""
Initialize the super class with common parameters available to all classes.
@param name The name of this particular data writer. Used to look
up values on the parameter server.
"""
self._name = name
# Look up where the data should go.
output_param = '~' + self._name + '/output_folder'
self._output_location = ParameterLookup.lookup(parameter=output_param)
# This is used as a base path, so make sure it ends in a '/'
if self._output_location[-1] != '/':
self._output_location += '/'
self._createDirIfNotFound(self._output_location)
# Count how many images have been written for use in the names.
self._count = 0
# Create a TF listener to transform key points.
self._tf_buffer = tf2_ros.Buffer()
self._tf_listener = tf2_ros.TransformListener(buffer=self._tf_buffer)
timeout_param = '~' + self._name + '/tf_timeout'
self._tf_timeout = ParameterLookup.lookupWithDefault(
parameter=timeout_param, default=2.0)
# Look up the camera frame ID to project points.
self._camera_frame_id = ParameterLookup.lookupWithDefault(
parameter='~camera_frame_id', default='camera/base_link')
def __init__(self, name):
"""!
The init function instantiates the class and loads all of the relevant
information from the parameter server.
@param nameThe name of this robot. Also includes how to find information
about it on the parameter server.
"""
# Set the name
self._name = name
# Look up the class number and associated ID
self._class_id = ParameterLookup.lookup(
'~' + self._resolveString('class'))
self._id = ParameterLookup.lookup('~' + self._resolveString("id"))
# Also look up the keypoints used for this robot
self._keypoints = self._initializeKeypoints(
'~' + self._resolveString("keypoints"))
# Look up the bounding shape, which is treated mathematically
# the same as keypoints.
self._bounding_shape = self._initializeKeypoints(
'~' + self._resolveString('bounding_shape'))
# Determine what name to use to find the robot in the TF tree.
self._frame_id = ParameterLookup.lookupWithDefault(
'~' + self._resolveString("frame_id"), "base_link")
# Set the pose to the origin.
transform = Transform()
transform.rotation.w = 1.0
self.recordTransform(transform)
def initializeBackgroundSubtractor(robot_list, gazebo_set_pose_client):
"""!
Create the background subtractor. Then, manipulate the Gazebo environment to allow
the capture of a background image. This uses the MOG2 algorithm.
@param robot_list A list of the robots that shouldn't be part of the background.
@param gazebo_set_pose_client A server client used to move every robot within Gazebo.
@return The background subtractor used in the script.
"""
# Set a large history to maximize the detection of differences from background.
background_history = ParameterLookup.lookupWithDefault(
parameter='~background_subtractor/history_length', default=100)
# Explore how shadow detection impacts the results.
var_threshold = ParameterLookup.lookupWithDefault(
parameter='~background_subtractor/var_threshold', default=500)
detect_shadows = ParameterLookup.lookupWithDefault(
parameter='~background_subtractor/detect_shadows', default=False)
background_subtractor = cv2.createBackgroundSubtractorMOG2(
history=background_history, varThreshold=var_threshold, detectShadows=detect_shadows)
# Create a client to determine where each robot is located within the Gazebo environment.
# This is only done at initialization to move every robot out of the way. It is assumed the
# robots start in a stable location. So the system is safe to just move them up in the air,
# but leave their X/Y the same. Given Gazebo's propensity to send colliding objects flying,
# this seems like the best course of action. The user is unlikely to start the node if
# the simulation freaks out right away. Using the bag file first location is possible, but
# would require passing the tf_buffer into this function. After initialization, the bag file
# positions are used for the rest of the node.
gazebo_get_pose_name = '/gazebo/get_model_state'
rospy.loginfo('Waiting to find robot poses...')
rospy.wait_for_service(gazebo_get_pose_name)
gazebo_get_pose_client = rospy.ServiceProxy(
name=gazebo_get_pose_name, service_class=GetModelState)
# Move all the robots way up in the sky, presumably outside of the view of the camera.
rospy.loginfo('Moving robots and capturing background...')
for robot in robot_list:
# Create each message via the robot object.
robot_current_state = gazebo_get_pose_client(robot.getName(), '')
robot_transform_msg = Transform()
robot_transform_msg.translation.x = robot_current_state.pose.position.x
robot_transform_msg.translation.y = robot_current_state.pose.position.y
robot_transform_msg.translation.z = robot_current_state.pose.position.z
robot_transform_msg.rotation = robot_current_state.pose.orientation
robot.recordTransform(robot_transform_msg)
robot_new_pose_request = robot.createSetModelStateRequest()
robot_new_pose_request.pose.position.z = 1e6
moveRobot(gazebo_set_pose_client, robot_new_pose_request)
# Capture N images and apply to subtractor
rospy.sleep(2.0)
for _ in range(background_history):
# Sleep long enough for any noise to update
(image, _) = captureImage(sleep_duration=0.25)
background_subtractor.apply(image)
return background_subtractor
def _initializeKeypoints(self, param_name):
"""!
This is an internal function to look up and process the list of keypoints
that is provided for the robot. It converts them to a list of
geometry_msgs/Point. All parameters are assumed within the robot's namespace
(which is also within the node's namespace)
@param param_name The parameter to lookup, without the robot's name.
@return The keypoints as a list of geometry_msgs/Point.
@throw rospy.InitException Thrown if the keypoints parameter is not
found.
"""
keypoint_param = ParameterLookup.lookup(param_name)
# The specifications are imported as a list of lists, so iterate
# through the outer one and convert the inner into the message
# format.
keypoints = []
for keypoint_list in keypoint_param:
# Verify that each point has three entitites for the 3 dimensions.
if len(keypoint_list) != 3:
error_string = 'Keypoints must have all three dimensions: [x, y, z]'
rospy.logfatal(error_string)
raise rospy.ROSInitException(error_string)
keypoint_message = Point()
keypoint_message.x = keypoint_list[0]
keypoint_message.y = keypoint_list[1]
keypoint_message.z = keypoint_list[2]
keypoints.append(keypoint_message)
return keypoints
def initializeRobots():
"""!
This function determines the names of each robot that should be controlled. It then creates
a @ref Robot object for each robot. Those objects perform the correct initialization of needed
parameters.
@return A list of all created Robot objects
@throw rospy.ROSInitException Thrown if no robot_list parameter is found or if no robots are
specified.
"""
# Determine the list of robots and error if not found.
robot_names = ParameterLookup.lookup('/robot_list')
# Throw an error if there are no robots specified.
if len(robot_names) < 1:
error_string = 'Must specify at least one robot in /robot_list'
rospy.logfatal(error_string)
raise rospy.ROSInitException(error_string)
# There is a chance the user specifies only a single robot without brackets, which would make
# this parameter as a string. Check for that and convert it to a list for use later.
if type(robot_names) is not list:
single_robot_name = robot_names
robot_names = [single_robot_name]
# Create each robot object and add it to a list.
robot_list = []
for name in robot_names:
robot_list.append(Robot(name))
return robot_list
def initializeReplay(global_frame_id, robot_list):
"""!
@brief Load the TF tree from the provided bag file.
This initialization takes a bag file, as provided by the user on the parameter
server and reads the entire TF tree into a buffer for use in replaying.
@param global_frame_id The global refernce frame used by Gazebo, often 'map'
@param robot_list The list of robots in the simulation. Used to determine the start
and end times that have valid TF transforms for all robots.
@return A tuple containing:
1. The tf buffer
2. A rospy.Time of the start time of the bag file
3. A rospy.Time of the end time of the bag file
4. A rospy.Duration of the period used during processing
@throw rospy.ROSInitException Thrown if the bag file can't be opened or found.
"""
rospy.loginfo('Loading bag file...')
# Determine where the bag file is located via parameter server.
bag_file = ParameterLookup.lookup(parameter='~bag_file')
if 'example_replay.bag' in bag_file:
rospy.logwarn(
msg='You might be using the example bag file. Did you mean to provide your own using the bag_file parameter?')
# Open the bag file. This might throw an error, so catch it and send a warning
# to the user.
try:
bag = rosbag.Bag(f=bag_file)
except Exception as ex:
error_string = 'Unable to open {bag}: {error}'.format(
bag=bag_file, error=ex)
rospy.logfatal(msg=error_string)
raise rospy.ROSInitException(error_string)
# If this point is reached, the bag file is loaded and ready.
# Now look up the simulated rate with the same method as the bag parameter.
rate = ParameterLookup.lookupWithDefault(
parameter='~simulated_rate', default=30.0)
period = rospy.Duration(secs=(1.0 / rate))
# Now, parse the bag file into a tf buffer
start_time = rospy.Time(bag.get_start_time())
end_time = rospy.Time(bag.get_end_time())
tf_buffer = tf2_py.BufferCore((end_time - start_time))
for topic, msg, _ in bag.read_messages(topics=['/tf', '/tf_static']):
for msg_tf in msg.transforms:
if topic == '/tf_static':
tf_buffer.set_transform_static(msg_tf, 'default_authority')
else:
tf_buffer.set_transform(msg_tf, 'default_authority')
# Now that everything is loaded, we don't need the bag file.
bag.close()
# Find the first start_time that has a valid transform for every single robot.
good_start_time = False
# This only finishes if every robot has a viable transform at this time step
while not good_start_time and start_time < end_time:
good_start_time = True
for robot in robot_list:
(is_good, _) = tf_buffer.can_transform_core(
global_frame_id, robot.getFullFrame(), start_time)
good_start_time = good_start_time and is_good
if not good_start_time:
start_time += period
# If the start time is now greater than the end time, then there were no
# time frames that had valid transforms for every robot. This is a problem.
if not start_time < end_time:
error_string = 'Unable to find valid transforms for all robots in bag file.'
rospy.logfatal(error_string)
raise rospy.exceptions.ROSInitException(error_string)
# Now do the same thing for the end time, but decrement instead. The end time
# is only a gate on when the whole simulation should end, so it does not have to
# be a whole number of increments from start time. Therefore, just start at the
# end of the bag file.
good_end_time = False
# You technically need to watch for the end > start since there may be a specific
# setup that only has frame interpolation in a very narrow time range and if period
# is larger than that range, end_time is liable to never exist within that range.
while not good_end_time and end_time > start_time:
good_end_time = True
for robot in robot_list:
(is_good, _) = tf_buffer.can_transform_core(
global_frame_id, robot.getFullFrame(), end_time)
good_end_time = good_end_time and is_good
if not good_end_time:
end_time -= period
if not end_time > start_time:
error_string = 'Unable to find valid transforms for all robots in bag file.'
rospy.logfatal(error_string)
raise rospy.exceptions.ROSInitException(error_string)
return (tf_buffer, start_time, end_time, period)
if __name__ == "__main__":
rospy.init_node(name='collect_data')
# Set up the different data writers.
data_writers = [
label_writers.BoundingShape('bounding_shape'),
label_writers.YOLO('yolo'),
label_writers.STVNet('stvnet')
]
# Flag if an instance mask is needed for each robot.
mask_required = False
for writer in data_writers:
mask_required = mask_required or writer.requireInstanceMask()
# Setup camera
camera_frame_id = ParameterLookup.lookupWithDefault(
parameter='~camera_frame_id', default='camera/optical_link')
# Look up what frame to use as the reference point for all placement.
global_frame_id = ParameterLookup.lookupWithDefault(
parameter='~global_frame', default='map')
# Load each robot
robot_list = initializeRobots()
# Load the buffer containing the motion record from the bag file.
(tf_buffer, start_time, end_time, period) = initializeReplay(
global_frame_id, robot_list)
current_time = start_time
# Create the client to tell each robot to move in Gazebo.
gazebo_client_name = '/gazebo/set_model_state'
rospy.loginfo('Waiting for Gazebo to start...')
rospy.wait_for_service(gazebo_client_name)
gazebo_client = rospy.ServiceProxy(
name=gazebo_client_name, service_class=SetModelState)
def __init__(self, name):
super(BoundingShape, self).__init__(name)
self._output_frame = ParameterLookup.lookup(
parameter='~' + self._name + '/output_frame')