counter in 3d

src/rob2002_tutorial/rob2002_tutorial/counter_3d.py

@@ -0,0 +1,65 @@
+import rclpy
+from rclpy.node import Node
+from rclpy import qos
+import math
+
+from std_msgs.msg import Header
+from geometry_msgs.msg import PoseStamped, PoseArray
+
+class Counter3D(Node):
+    detection_threshold = 0.2 # in meters    
+
+    def __init__(self):      
+        super().__init__('counter_3d')
+
+        self.detected_objects = [] # list of all detected objects
+
+        # subscribe to object detector
+        self.subscriber = self.create_subscription(PoseStamped, '/object_location', 
+                                                   self.counter_callback,
+                                                   qos_profile=qos.qos_profile_sensor_data)
+
+        # publish all detected object as an array of poses
+        self.publisher = self.create_publisher(PoseArray, '/object_count_array',
+                                                      qos.qos_profile_parameters)
+
+    def counter_callback(self, data):
+        new_object = data.pose
+        # check if the new object is away from all detected objects so far
+
+        object_exists = False
+        for object in self.detected_objects:
+            # calculate the distance between the new_object and each in the list
+            pos_a = object.position
+            pos_b = new_object.position
+            d = math.sqrt((pos_a.x - pos_b.x) ** 2 + (pos_a.y - pos_b.y) ** 2 + (pos_a.z - pos_b.z) ** 2)
+            if d < self.detection_threshold: # found a close neighbour in the already existing list, so this one won't be added
+                object_exists = True
+                break
+
+        if not object_exists: # new object!
+            self.detected_objects.append(new_object)
+
+        # publish a PoseArray of object poses for visualisation in rviz
+        parray = PoseArray(header=Header(frame_id=data.header.frame_id))
+        for object in self.detected_objects:
+            parray.poses.append(object)
+        self.publisher.publish(parray)            
+
+        # print to the console
+        print(f'total count {len(self.detected_objects)}')
+        for object in self.detected_objects:
+            print(object.position)
+
+def main(args=None):
+    rclpy.init(args=args)
+    counter_3d = Counter3D()
+
+    rclpy.spin(counter_3d)
+
+    counter_3d.destroy_node()
+    rclpy.shutdown()
+
+
+if __name__ == '__main__':
+    main()

src/rob2002_tutorial/rob2002_tutorial/detector_3d.py

@@ -19,7 +19,7 @@
 from geometry_msgs.msg import Pose, PoseStamped, Point, Quaternion
 
 class Detector3D(Node):
-    # use real robot?
+    # use the real robot?
     real_robot = False
 
     ccamera_model = None
@@ -32,13 +32,13 @@ class Detector3D(Node):
     # aspect ratio between color and depth cameras
     # calculated as (color_horizontal_FOV/color_width) / (depth_horizontal_FOV/depth_width)
     color2depth_aspect = 1.0 # simulated camera
-    camera_frame = 'depth_link' # for simulated robot
+    camera_frame = 'depth_link' 
 
     if real_robot:
-        color2depth_aspect = (71.0/640) / (67.9/400) # real camera
-        camera_frame = 'camera_link' # for simulated robot
+        color2depth_aspect = (71.0/640) / (67.9/640) # real camera
+        camera_frame = 'camera_color_optical_frame' 
 
-    visualisation = True
+    visualisation = False
 
     def __init__(self):    
         super().__init__('Detector3D')
@@ -137,7 +137,7 @@ def image_color_callback(self, data):
                                                 self.tf_buffer.lookup_transform(self.global_frame, self.camera_frame, rclpy.time.Time())) 
 
                 # publish so we can see that in rviz
-                self.object_location_pub.publish(PoseStamped(header=Header(frame_id=self.camera_frame),
+                self.object_location_pub.publish(PoseStamped(header=Header(frame_id=self.global_frame),
                                               pose=global_pose))        
 
                 print(f'--- object id {num} ---')
@@ -151,8 +151,9 @@ def image_color_callback(self, data):
 
         if self.visualisation:
             #resize and adjust for visualisation
-            # image_color = cv2.resize(image_color, (0,0), fx=0.5, fy=0.5)
             self.image_depth *= 1.0/10.0 # scale for visualisation (max range 10.0 m)
+            # self.image_color = cv2.resize(self.image_color, (0,0), fx=0.5, fy=0.5)
+            # self.image_depth = cv2.resize(self.image_depth, (0,0), fx=0.5, fy=0.5)
             cv2.imshow("image color", self.image_color)
             cv2.imshow("image depth", self.image_depth)
             cv2.waitKey(1)

src/rob2002_tutorial/setup.py

@@ -31,6 +31,7 @@
             'detector_better = rob2002_tutorial.detector_dblcounting:main',
             'detector_3d = rob2002_tutorial.detector_3d:main',
             'counter_basic = rob2002_tutorial.counter_basic:main',
+            'counter_3d = rob2002_tutorial.counter_3d:main',
         ],
     },
 )