add inference

_tc/ros/yolov5_tc/package.xml

@@ -11,6 +11,7 @@
   <depend>cv_bridge</depend>
   <depend>perception_msgs</depend>
   <depend>python-transforms3d-pip</depend>
+  <depend>python3-ultralytics-pip</depend>
   <depend>rclpy</depend>
   <depend>sensor_msgs</depend>
 

_tc/ros/yolov5_tc/yolov5_tc/inference_node.py

@@ -2,6 +2,8 @@
 from sensor_msgs.msg import Image
 from perception_msgs.msg import ObjectList
 from rclpy.node import Node
+from cv_bridge import CvBridge
+import torch
 
 # constants
 INPUT_TOPIC = "~/input"
@@ -32,17 +34,18 @@ def setup(self):
                                                qos_profile=1)
         self.get_logger().info(f"Publishing to '{self.publisher.topic_name}'")
 
+        self.cv_bridge = CvBridge()
+        self.model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+
     def topicCallback(self, msg: Image):
 
         # echo received message
         self.get_logger().info(f"Received image with stamp: '{msg.header.stamp}'")
 
-        # skip if the published topic has no subscription
-        # if self.publisher.get_subscription_count() == 0:
-        #     return
-
-        # TODO Inference
+        cv_image = self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
 
+        result = self.model(cv_image)
+        print(result)
 
         # create and publish a new message
         object_list = ObjectList()