Move processing to image callback

This has the advantage that the processing is event-driven, i.e., an image is processed as soon as it is received instead of being throttled by the processing loop frequency. Additionally, it ensures that each input image is processed at most once, whereas the old code repeatedly processed the same input image until a new image arrived.
kini5gowda · Jun 13, 2019 · aeb89bb · aeb89bb
1 parent 2894d78
commit aeb89bb
Showing 1 changed file with 56 additions and 65 deletions.
diff --git a/nodes/dope b/nodes/dope
@@ -28,7 +28,6 @@ from std_msgs.msg import String
 ### Global Variables
 g_path2package = rospkg.RosPack().get_path('dope')
 g_bridge = CvBridge()
-g_img = None
 
 
 class Draw(object):
@@ -169,76 +168,68 @@ class DopeNode(object):
         rospy.Subscriber(
             topic_cam,
             ImageSensor_msg,
-            self.image_callback
+            self.image_callback,
+            queue_size=1
         )
 
         print("Running DOPE...  (Listening to camera topic: '{}')".format(topic_cam))
         print("Ctrl-C to stop")
 
     def image_callback(self, msg):
         """Image callback"""
-        global g_img
-        g_img = g_bridge.imgmsg_to_cv2(msg, "rgb8")
-        # cv2.imwrite('img.png', cv2.cvtColor(g_img, cv2.COLOR_BGR2RGB))  # for debugging
-
-    def run_dope_node(self, freq=5):
-        global g_img
-
-        rate = rospy.Rate(freq)
-
-        while not rospy.is_shutdown():
-            if g_img is not None:
-                # Copy and draw image
-                img_copy = g_img.copy()
-                im = Image.fromarray(img_copy)
-                draw = Draw(im)
-
-                for m in self.models:
-                    # Detect object
-                    results = ObjectDetector.detect_object_in_image(
-                        self.models[m].net,
-                        self.pnp_solvers[m],
-                        g_img,
-                        self.config_detect
-                    )
-
-                    # Publish pose and overlay cube on image
-                    for i_r, result in enumerate(results):
-                        if result["location"] is None:
-                            continue
-                        loc = result["location"]
-                        ori = result["quaternion"]
-                        msg = PoseStamped()
-                        msg.header.frame_id = self.params["frame_id"]
-                        msg.header.stamp = rospy.Time.now()
-                        CONVERT_SCALE_CM_TO_METERS = 100
-                        msg.pose.position.x = loc[0] / CONVERT_SCALE_CM_TO_METERS
-                        msg.pose.position.y = loc[1] / CONVERT_SCALE_CM_TO_METERS
-                        msg.pose.position.z = loc[2] / CONVERT_SCALE_CM_TO_METERS
-                        msg.pose.orientation.x = ori[0]
-                        msg.pose.orientation.y = ori[1]
-                        msg.pose.orientation.z = ori[2]
-                        msg.pose.orientation.w = ori[3]
-
-                        # Publish
-                        self.pubs[m].publish(msg)
-                        self.pub_dimension[m].publish(str(self.params['dimensions'][m]))
-
-                        # Draw the cube
-                        if None not in result['projected_points']:
-                            points2d = []
-                            for pair in result['projected_points']:
-                                points2d.append(tuple(pair))
-                            draw.draw_cube(points2d, self.draw_colors[m])
-
-                # Publish the image with results overlaid
-                self.pub_rgb_dope_points.publish(
-                    CvBridge().cv2_to_imgmsg(
-                        np.array(im)[..., ::-1],
-                        "bgr8"
-                    )
-                )
-            rate.sleep()
+        img = g_bridge.imgmsg_to_cv2(msg, "rgb8")
+        # cv2.imwrite('img.png', cv2.cvtColor(img, cv2.COLOR_BGR2RGB))  # for debugging
+
+        # Copy and draw image
+        img_copy = img.copy()
+        im = Image.fromarray(img_copy)
+        draw = Draw(im)
+
+        for m in self.models:
+            # Detect object
+            results = ObjectDetector.detect_object_in_image(
+                self.models[m].net,
+                self.pnp_solvers[m],
+                img,
+                self.config_detect
+            )
+
+            # Publish pose and overlay cube on image
+            for i_r, result in enumerate(results):
+                if result["location"] is None:
+                    continue
+                loc = result["location"]
+                ori = result["quaternion"]
+                msg = PoseStamped()
+                msg.header.frame_id = self.params["frame_id"]
+                msg.header.stamp = rospy.Time.now()
+                CONVERT_SCALE_CM_TO_METERS = 100
+                msg.pose.position.x = loc[0] / CONVERT_SCALE_CM_TO_METERS
+                msg.pose.position.y = loc[1] / CONVERT_SCALE_CM_TO_METERS
+                msg.pose.position.z = loc[2] / CONVERT_SCALE_CM_TO_METERS
+                msg.pose.orientation.x = ori[0]
+                msg.pose.orientation.y = ori[1]
+                msg.pose.orientation.z = ori[2]
+                msg.pose.orientation.w = ori[3]
+
+                # Publish
+                self.pubs[m].publish(msg)
+                self.pub_dimension[m].publish(str(self.params['dimensions'][m]))
+
+                # Draw the cube
+                if None not in result['projected_points']:
+                    points2d = []
+                    for pair in result['projected_points']:
+                        points2d.append(tuple(pair))
+                    draw.draw_cube(points2d, self.draw_colors[m])
+
+        # Publish the image with results overlaid
+        self.pub_rgb_dope_points.publish(
+            CvBridge().cv2_to_imgmsg(
+                np.array(im)[..., ::-1],
+                "bgr8"
+            )
+        )
 
 
 def main():
@@ -265,7 +256,7 @@ def main():
     n = DopeNode(params)
 
     try:
-        n.run_dope_node()
+        rospy.spin()
     except rospy.ROSInterruptException:
         pass