selective_search.py
Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 
00003 import rospy
00004 
00005 import cv_bridge
00006 from jsk_topic_tools import ConnectionBasedTransport
00007 from jsk_recognition_msgs.msg import RectArray, Rect
00008 from sensor_msgs.msg import Image
00009 import dynamic_reconfigure.server
00010 from jsk_perception.cfg import SelectiveSearchConfig
00011 import dlib
00012 import cv2
00013 
00014 class SelectiveSearch(ConnectionBasedTransport):
00015     def __init__(self):
00016         super(SelectiveSearch, self).__init__()
00017         dynamic_reconfigure.server.Server(
00018             SelectiveSearchConfig, self._cb_dyn_reconfig)
00019         self.pub_ = self.advertise('~output', RectArray, queue_size=10)
00020         self.debug_pub_ = self.advertise('~debug', Image, queue_size=1)
00021 
00022     def subscribe(self):
00023         self.sub_ = rospy.Subscriber('~input', Image, self._apply, queue_size=1)
00024 
00025     def unsubscribe(self):
00026         self.sub_.unregister()
00027     def _cb_dyn_reconfig(self, config, level):
00028         self.min_size = config.min_size
00029         self.max_size = config.max_size
00030         return config
00031     def _apply(self, img_msg):
00032         bridge = cv_bridge.CvBridge()
00033         img_bgr = bridge.imgmsg_to_cv2(img_msg, desired_encoding='bgr8')
00034         img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
00035         rects = []
00036         dlib.find_candidate_object_locations(img_rgb, rects,
00037                                              min_size=self.min_size)
00038         ros_rect_array = RectArray()
00039         ros_rect_array.header = img_msg.header
00040         for d in rects:
00041             if (d.right() - d.left()) * (d.bottom() - d.top()) > self.max_size:
00042                 continue
00043             cv2.rectangle(img_bgr, (d.left(), d.top()), (d.right(), d.bottom()), (255, 0, 0), 3)
00044             ros_rect_array.rects.append(Rect(x=d.left(), y=d.top(), width=d.right() - d.left(), height=d.bottom() - d.top()))
00045         imgmsg = bridge.cv2_to_imgmsg(img_bgr, encoding='bgr8')
00046         self.debug_pub_.publish(imgmsg)
00047         self.pub_.publish(ros_rect_array)
00048 
00049 
00050 if __name__ == '__main__':
00051     rospy.init_node('selective_search')
00052     selective_search = SelectiveSearch()
00053     rospy.spin()


jsk_perception
Author(s): Manabu Saito, Ryohei Ueda
autogenerated on Sun Oct 8 2017 02:43:23