added detection image and bounding boxes to run on bag

ros2_ws/src/yolov3/detect_on_bag.py

@@ -41,7 +41,8 @@ from utils.plots import Annotator, colors, save_one_box
 from utils.torch_utils import select_device, time_sync
 from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox
 import numpy as np
-
+from yolov4_msg.msg import BoundingBox, BoundingBoxes
+from rclpy.clock import Clock
 
 @torch.no_grad()
 def run(weights=ROOT / 'yolov3.pt',  # model.pt path(s)
@@ -119,27 +120,40 @@ def run(weights=ROOT / 'yolov3.pt',  # model.pt path(s)
         p = Path(p)  # to Path
         save_path = '/home/parallels/predict.jpeg'
         annotator = Annotator(im0, line_width=line_thickness, example=str(names))
-        # import pdb; pdb.set_trace()
+        
+        time_stamp = Clock().now()
+        bounding_boxes_msg = BoundingBoxes()
+        bounding_boxes_msg.header.stamp = time_stamp.to_msg()
+        bounding_boxes_msg.header.frame_id = 'detection'
+        s = ''
         if len(det):
             # Rescale boxes from img_size to im0 size
             det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
-
             # Print results
             for c in det[:, -1].unique():
                 n = (det[:, -1] == c).sum()  # detections per class
-                s = f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
+                s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string
             
             # Write results
             for *xyxy, conf, cls in reversed(det):
+                bounding_box = BoundingBox()
                 c = int(cls)  # integer class
                 label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                 annotator.box_label(xyxy, label, color=colors(c, True))
+                bounding_box.probability = float(conf)
+                bounding_box.xmin = int(xyxy[0])
+                bounding_box.ymin = int(xyxy[1])
+                bounding_box.xmax = int(xyxy[2])
+                bounding_box.ymax = int(xyxy[3])
+                bounding_box.id = c
+                bounding_box.class_id = names[c]
+                bounding_boxes_msg.bounding_boxes.append(bounding_box)
         # Print time (inference-only)
         LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
         # Stream results
         im0 = annotator.result()
         # cv2.imwrite(save_path, im0)
-        return im0
+        return im0, bounding_boxes_msg
 
 class DetectOnBag(Node):
 
@@ -152,6 +166,7 @@ class DetectOnBag(Node):
             10)
         self.subscription  # prevent unused variable warning
         self.detected_img_pub = self.create_publisher(Image, '/detection_image', 10)
+        self.bboxes_pub = self.create_publisher(BoundingBoxes, '/bboxes', 10)
 
     def image_callback(self, msg):
         self.get_logger().info('Image')
@@ -159,8 +174,10 @@ class DetectOnBag(Node):
 
         bridge = CvBridge()
         cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding='passthrough')
-        detected_img = run('runs/train/exp14/weights/best.pt', cv_image)
+        detected_img, bboxes = run('src/yolov3/runs/train/exp14/weights/best.pt', cv_image)
         self.detected_img_pub.publish(self.numpy_array_to_image_msg(detected_img))
+        bboxes.image_header = msg.header
+        self.bboxes_pub.publish(bboxes)
     
     def numpy_array_to_image_msg(self, numpy_array):
         # Create a CvBridge object