$search
00001 #!/usr/bin/env python 00002 00003 PKG = 'pr2_image_snapshot_recorder' 00004 import roslib; roslib.load_manifest(PKG) 00005 00006 import rospy, actionlib 00007 00008 import rosbag 00009 00010 from sensor_msgs.msg import Image 00011 from pr2_image_snapshot_recorder.msg import ImageSnapshotAction, ImageSnapshotGoal, ImageSnapshotResult 00012 00013 import time 00014 import threading 00015 00016 def get_bagname(topic_name): 00017 if topic_name.startswith('/'): 00018 topic_name = topic_name[1:] 00019 return topic_name.replace('/', '_') + '_' + time.strftime('%Y%m%d_%I%M%S', time.localtime()) + '.bag' 00020 00021 00022 class ImageSnapshotter(object): 00023 def __init__(self): 00024 self._reset() 00025 self._mutex = threading.Lock() 00026 00027 self._server = actionlib.SimpleActionServer('image_snapshot', ImageSnapshotAction, self._execute) 00028 00029 def _reset(self): 00030 self._image_sub = None 00031 self._max_image_count = -1 00032 self._filename = -1 00033 self._images_received = 0 00034 self._bag = None 00035 self._topic_name = None 00036 self._done = False 00037 00038 def _cb(self, msg): 00039 with self._mutex: 00040 # Check if we're full 00041 if self._max_image_count > 0 and self._images_received >= self._max_image_count: 00042 self._done = True 00043 return 00044 00045 # Add msg to bag file 00046 self._bag.write(self._topic_name, msg) 00047 self._images_received += 1 00048 00049 00050 def _execute(self, goal): 00051 # Open bag 00052 filename = goal.output_file_name if goal.output_file_name else get_bagname(goal.topic_name) 00053 00054 rospy.loginfo('Snapshot recording of %d images on \"%s\" to file \"%s\"' % \ 00055 (goal.num_images, goal.topic_name, filename)) 00056 00057 if not filename.endswith('.bag'): 00058 filename += '.bag' 00059 00060 self._bag = rosbag.Bag(filename, 'w') 00061 00062 self._topic_name = goal.topic_name 00063 self._max_image_count = goal.num_images 00064 00065 # Set up image subscription 00066 self._image_sub = rospy.Subscriber(goal.topic_name, Image, self._cb) 00067 00068 # Poll until we're done 00069 while not self._done and not rospy.is_shutdown(): 00070 if self._server.is_preempt_requested() or \ 00071 self._server.is_new_goal_available(): 00072 break 00073 rospy.sleep(0.05) 00074 00075 self._image_sub.unregister() 00076 00077 self._bag.close() 00078 00079 if self._done: 00080 self._server.set_succeeded() 00081 else: 00082 self._server.set_preempted() 00083 self._reset() 00084 00085 rospy.loginfo('Finished snapshot recording of images') 00086 00087 if __name__ == '__main__': 00088 rospy.init_node('image_snapshotter') 00089 00090 ImageSnapshotter() 00091 00092 rospy.spin() 00093 00094