5 from __future__
import print_function
9 import matplotlib.pyplot
as plt
11 from cv_bridge
import CvBridge
12 from jsk_topic_tools
import ConnectionBasedTransport
18 from dynamic_reconfigure.server
import Server
19 from jsk_perception.cfg
import SSDObjectDetectorConfig
as Config
21 from sensor_msgs.msg
import Image
22 from jsk_recognition_msgs.msg
import Rect, RectArray
23 from jsk_recognition_msgs.msg
import ClassificationResult
24 from jsk_recognition_msgs.msg
import Label
25 from jsk_recognition_msgs.msg
import LabelArray
26 from jsk_recognition_msgs.msg
import ClusterPointIndices
27 from pcl_msgs.msg
import PointIndices
29 import itertools, pkg_resources, sys
30 from distutils.version
import LooseVersion
31 if LooseVersion(pkg_resources.get_distribution(
"chainer").version) >= LooseVersion(
'7.0.0')
and \
32 sys.version_info.major == 2:
33 print(
'''Please install chainer < 7.0.0: 35 sudo pip install chainer==6.7.0 37 c.f https://github.com/jsk-ros-pkg/jsk_recognition/pull/2485 40 if [p
for p
in list(itertools.chain(*[pkg_resources.find_distributions(_)
for _
in sys.path]))
if "cupy-" in p.project_name ] == []:
41 print(
'''Please install CuPy 43 sudo pip install cupy-cuda[your cuda version] 45 sudo pip install cupy-cuda91 49 from chainercv.links
import SSD300
50 from chainercv.links
import SSD512
51 from chainercv.visualizations
import vis_bbox
54 chainer.config.cv_resize_backend =
'cv2' 60 super(SSDObjectDetector, self).
__init__()
61 self.
gpu = rospy.get_param(
"~gpu", -1)
68 rospy.loginfo(
"Loaded %d labels" % len(self.
label_names))
71 model_path = rospy.get_param(
"~model_path",
None)
72 model_name = rospy.get_param(
'~model',
'ssd300')
73 if model_name ==
'ssd300':
75 elif model_name ==
'ssd512':
78 rospy.logerr(
'Unsupported ~model: {}'.format(model_name))
81 pretrained_model=model_path)
83 chainer.cuda.get_device_from_id(self.
gpu).use()
85 rospy.loginfo(
"Loaded model: %s" % model_path)
91 self.
pub_labels = self.advertise(
"~output/labels", LabelArray,
93 self.
pub_indices = self.advertise(
"~output/cluster_indices", ClusterPointIndices,
95 self.
pub_rects = self.advertise(
"~output/rect", RectArray,
97 self.
pub_class = self.advertise(
"~output/class", ClassificationResult,
99 self.
pub_image = self.advertise(
"~output/image", Image,
104 queue_size=1, buff_size=2**26)
107 self.sub_image.unregister()
111 return self.pub_image.get_num_connections() > 0
114 label_names = rospy.get_param(
"~label_names", tuple())
117 from chainercv.datasets
import voc_detection_label_names
118 label_names = voc_detection_label_names
120 from chainercv.datasets
import voc_bbox_label_names
121 label_names = voc_bbox_label_names
122 elif isinstance(label_names, str):
123 with open(label_names,
"r") as f: 124 label_names = tuple(yaml.load(f)) 128 self.model.nms_thresh = config.nms_thresh
129 self.model.score_thresh = config.score_thresh
135 rospy.loginfo(
"callback start: incomming msg is %s msec behind" % ((rospy.Time.now() - msg.header.stamp).to_sec() * 1000.0))
139 img = self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding=
"rgb8")
140 img = np.asarray(img, dtype=np.float32)
141 img = img.transpose(2, 0, 1)
142 except Exception
as e:
143 rospy.logerr(
"Failed to convert image: %s" % str(e))
147 rospy.loginfo(
"%s: elapsed %f msec" % (
"convert", (tcur-tprev)*1000))
151 chainer.cuda.get_device_from_id(self.
gpu).use()
152 bboxes, labels, scores = self.model.predict([img])
153 bboxes, labels, scores = bboxes[0], labels[0], scores[0]
157 rospy.loginfo(
"%s: elapsed %f msec" % (
"predict", (tcur-tprev)*1000))
160 labels_msg = LabelArray(header=msg.header)
163 labels_msg.labels.append(Label(id=l, name=l_name))
167 rospy.loginfo(
"%s: elapsed %f msec" % (
"make labels msg", (tcur-tprev)*1000))
169 cluster_indices_msg = ClusterPointIndices(header=msg.header)
173 ymin = max(0,
int(np.floor(bbox[0])))
174 xmin = max(0,
int(np.floor(bbox[1])))
175 ymax = min(H,
int(np.ceil(bbox[2])))
176 xmax = min(W,
int(np.ceil(bbox[3])))
177 indices = [range(W*y+xmin, W*y+xmax)
for y
in range(ymin, ymax)]
178 indices = np.array(indices, dtype=np.int32).flatten()
179 indices_msg = PointIndices(header=msg.header, indices=indices)
180 cluster_indices_msg.cluster_indices.append(indices_msg)
184 rospy.loginfo(
"%s: elapsed %f msec" % (
"make cluster_indices msg", (tcur-tprev)*1000))
187 rect_msg = RectArray(header=msg.header)
189 rect = Rect(x=bbox[1], y=bbox[0],
190 width= bbox[3] - bbox[1],
191 height=bbox[2] - bbox[0])
192 rect_msg.rects.append(rect)
196 rospy.loginfo(
"%s: elapsed %f msec" % (
"make rect msg", (tcur-tprev)*1000))
199 cls_msg = ClassificationResult(
210 rospy.loginfo(
"%s: elapsed %f msec" % (
"make cls msg", (tcur-tprev)*1000))
213 self.pub_labels.publish(labels_msg)
214 self.pub_indices.publish(cluster_indices_msg)
215 self.pub_rects.publish(rect_msg)
216 self.pub_class.publish(cls_msg)
220 rospy.loginfo(
"%s: elapsed %f msec" % (
"publish msg", (tcur-tprev)*1000))
228 rospy.loginfo(
"%s: elapsed %f msec" % (
"callback end", (tcur-tprev)*1000))
232 vis_bbox(img, bbox, label, score,
236 w, h = fig.canvas.get_width_height()
237 img = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8)
239 img.shape = (h, w, 3)
242 msg = self.cv_bridge.cv2_to_imgmsg(img,
"rgb8")
243 except Exception
as e:
244 rospy.logerr(
"Failed to convert bbox image: %s" % str(e))
247 self.pub_image.publish(msg)
250 if __name__ ==
'__main__':
251 rospy.init_node(
"ssd_object_detector")
def publish_bbox_image(self, img, bbox, label, score, header)
def load_label_names(self)
def config_callback(self, config, level)