Add vision_msgs publisher

kini5gowda · Jun 13, 2019 · 5297028 · 5297028
1 parent ee57dcb
commit 5297028
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 9 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -12,8 +12,10 @@ find_package(catkin REQUIRED COMPONENTS
   geometry_msgs
   message_filters
   rospy
+  sensor_msgs
   std_msgs
   tf2
+  vision_msgs
 )
 
 ## System dependencies are found with CMake's conventions
@@ -106,7 +108,7 @@ catkin_python_setup()
 ## CATKIN_DEPENDS: catkin_packages dependent projects also need
 ## DEPENDS: system dependencies of this project that dependent projects also need
 catkin_package(
-  CATKIN_DEPENDS geometry_msgs std_msgs
+  CATKIN_DEPENDS geometry_msgs sensor_msgs std_msgs vision_msgs
 #  DEPENDS system_lib
 )
 

diff --git a/config/config_pose.yaml b/config/config_pose.yaml
@@ -26,6 +26,16 @@ dimensions: {
     "bleach": [10.267730712890625,26.625339508056641,7.5134143829345703]
         }
 
+class_ids: {
+    "cracker": 1,
+    "gelatin": 2,
+    "meat":    3,
+    "mustard": 4,
+    "soup":    5,
+    "sugar":   6,
+    "bleach":  7,
+}
+
 draw_colors: {
     "cracker": [13, 255, 128],  # green
     "gelatin": [255, 255, 255],  # while

diff --git a/nodes/dope b/nodes/dope
@@ -26,6 +26,7 @@ from dope.inference.detector import ModelData, ObjectDetector
 from geometry_msgs.msg import PoseStamped
 from sensor_msgs.msg import CameraInfo, Image as ImageSensor_msg
 from std_msgs.msg import String
+from vision_msgs.msg import Detection3D, Detection3DArray, ObjectHypothesisWithPose
 
 
 class Draw(object):
@@ -98,6 +99,7 @@ class DopeNode(object):
         self.pub_dimension = {}
         self.draw_colors = {}
         self.dimensions = {}
+        self.class_ids = {}
         self.cv_bridge = CvBridge()
 
         self.input_is_rectified = params['input_is_rectified']
@@ -128,6 +130,8 @@ class DopeNode(object):
                 tuple(params["draw_colors"][model])
             self.dimensions[model] = \
                 tuple(params["dimensions"][model])
+            self.class_ids[model] = params["class_ids"][model]
+
             self.pnp_solvers[model] = \
                 CuboidPNPSolver(
                     model,
@@ -146,13 +150,19 @@ class DopeNode(object):
                     queue_size=10
                 )
 
-        # Start ROS publisher
+        # Start ROS publishers
         self.pub_rgb_dope_points = \
             rospy.Publisher(
                 params['topic_publishing'] + "/rgb_points",
                 ImageSensor_msg,
                 queue_size=10
             )
+        self.pub_detections = \
+            rospy.Publisher(
+                '~detected_objects',
+                Detection3DArray,
+                queue_size=10
+            )
 
         # Start ROS subscriber
         image_sub = message_filters.Subscriber(
@@ -203,7 +213,10 @@ class DopeNode(object):
         im = Image.fromarray(img_copy)
         draw = Draw(im)
 
-        for m in self.models:
+        detection_array = Detection3DArray()
+        detection_array.header = image_msg.header
+
+        for model_id, m in enumerate(self.models):
             # Detect object
             results = ObjectDetector.detect_object_in_image(
                 self.models[m].net,
@@ -233,6 +246,19 @@ class DopeNode(object):
                 self.pubs[m].publish(pose_msg)
                 self.pub_dimension[m].publish(str(self.dimensions[m]))
 
+                # Add to Detection3DArray
+                detection = Detection3D()
+                hypothesis = ObjectHypothesisWithPose()
+                hypothesis.id = self.class_ids[result["name"]]
+                hypothesis.score = result["score"]
+                hypothesis.pose.pose = pose_msg.pose
+                detection.results.append(hypothesis)
+                detection.bbox.center = pose_msg.pose
+                detection.bbox.size.x = self.dimensions[m][0] / CONVERT_SCALE_CM_TO_METERS
+                detection.bbox.size.y = self.dimensions[m][1] / CONVERT_SCALE_CM_TO_METERS
+                detection.bbox.size.z = self.dimensions[m][2] / CONVERT_SCALE_CM_TO_METERS
+                detection_array.detections.append(detection)
+
                 # Draw the cube
                 if None not in result['projected_points']:
                     points2d = []
@@ -247,6 +273,7 @@ class DopeNode(object):
                 "bgr8"
             )
         )
+        self.pub_detections.publish(detection_array)
 
 
 def main():

diff --git a/package.xml b/package.xml
@@ -53,9 +53,11 @@
   <depend>geometry_msgs</depend>
   <depend>message_filters</depend>
   <depend>rospy</depend>
+  <depend>sensor_msgs</depend>
   <depend>std_msgs</depend>
   <depend>tf2</depend>
-
+  <depend>vision_msgs</depend>
+
 
   <!-- The export tag contains other, unspecified, tags -->
   <export>

diff --git a/src/dope/inference/detector.py b/src/dope/inference/detector.py
@@ -279,6 +279,7 @@ def find_object_poses(vertex2, aff, pnp_solver, config):
                 'quaternion': quaternion,
                 'cuboid2d': cuboid2d,
                 'projected_points': projected_points,
+                'score': obj[-1],
             })
 
         return detected_objects
@@ -432,11 +433,8 @@ def find_objects(vertex2, aff, config, numvertex=8):
 
                         # distance between vertexes
                         dist_point = np.linalg.norm(np.array(point) - np.array(center))
-
-                        if dist_angle < config.thresh_angle \
-                                and best_dist > 1000 \
-                                or dist_angle < config.thresh_angle \
-                                and best_dist > dist_point:
+
+                        if dist_angle < config.thresh_angle and (best_dist > 1000 or best_dist > dist_point):
                             i_best = i_obj
                             best_angle = dist_angle
                             best_dist = dist_point