Merge pull request #130 from mqcmd196/PR/publish_skel

publish HumanSkeletonArray in human_pose_estimator

python/coral_usb/human_pose_estimator.py

@@ -27,11 +27,15 @@
 from geometry_msgs.msg import Point
 from geometry_msgs.msg import Pose
 from jsk_recognition_msgs.msg import ClassificationResult
+from jsk_recognition_msgs.msg import HumanSkeleton
+from jsk_recognition_msgs.msg import HumanSkeletonArray
 from jsk_recognition_msgs.msg import PeoplePose
 from jsk_recognition_msgs.msg import PeoplePoseArray
 from jsk_recognition_msgs.msg import Rect
 from jsk_recognition_msgs.msg import RectArray
+from jsk_recognition_msgs.msg import Segment
 from sensor_msgs.msg import CompressedImage
+from std_msgs.msg import Int32
 
 
 class EdgeTPUHumanPoseEstimator(EdgeTPUNodeBase):
@@ -59,6 +63,8 @@ def __init__(self, model_file=None, namespace='~'):
         self.label_names = ['human']
 
         # publishers
+        self.pub_skel = self.advertise(
+            namespace + 'output/skeletons', HumanSkeletonArray, queue_size=1)
         self.pub_pose = self.advertise(
             namespace + 'output/poses', PeoplePoseArray, queue_size=1)
         self.pub_rects = self.advertise(
@@ -163,25 +169,35 @@ def image_cb(self, msg):
         points, key_names, visibles, bboxes, labels, scores \
             = self._estimate(img)
 
+        skels_msg = HumanSkeletonArray(header=msg.header)
         poses_msg = PeoplePoseArray(header=msg.header)
         rects_msg = RectArray(header=msg.header)
         for point, key_name, visible, bbox, label, score in zip(
                 points, key_names, visibles, bboxes, labels, scores):
             pose_msg = PeoplePose()
+            skel_msg = HumanSkeleton(header=msg.header)
             for pt, key_nm, vs, sc in zip(point, key_name, visible, score):
                 if vs:
                     key_y, key_x = pt
                     pose_msg.limb_names.append(key_nm)
                     pose_msg.scores.append(sc)
                     pose_msg.poses.append(
                         Pose(position=Point(x=key_x, y=key_y)))
+            for i in range(len(pose_msg.poses)):
+                bone_name = "{}->{}".format(pose_msg.limb_names[i - 1],
+                                            pose_msg.limb_names[i])
+                skel_msg.bone_names.append(bone_name)
+                seg = Segment()
+                seg.start_point = pose_msg.poses[i - 1].position
+                seg.end_point = pose_msg.poses[i].position
+                skel_msg.bones.append(seg)
+            skels_msg.skeletons.append(skel_msg)
             poses_msg.poses.append(pose_msg)
             y_min, x_min, y_max, x_max = bbox
             rect = Rect(
                 x=x_min, y=y_min,
                 width=x_max - x_min, height=y_max - y_min)
             rects_msg.rects.append(rect)
-
         cls_msg = ClassificationResult(
             header=msg.header,
             classifier=self.classifier_name,
@@ -204,6 +220,9 @@ def image_cb(self, msg):
         self.pub_pose.publish(poses_msg)
         self.pub_rects.publish(rects_msg)
         self.pub_class.publish(cls_msg)
+        for i in range(len(skels_msg.skeletons)):
+            skels_msg.human_ids.append(Int32(data=i))
+        self.pub_skel.publish(skels_msg)
 
     def visualize_cb(self, event):
         if (not self.visualize or self.img is None or self.encoding is None