OBB support

requirements.txt

@@ -1,4 +1,4 @@
 opencv-python==4.8.1.78
 typing-extensions>=4.4.0
-ultralytics==8.2.45
+ultralytics==8.2.62
 lap==0.4.0

yolov8_bringup/package.xml

@@ -2,7 +2,7 @@
 <?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
 <package format="3">
   <name>yolov8_bringup</name>
-  <version>0.0.0</version>
+  <version>3.2.1</version>
   <description>YOLOv8 for ROS 2</description>
   <maintainer email="[email protected]">Miguel Ángel González Santamarta</maintainer>
   <license>GPL-3</license>

yolov8_msgs/package.xml

@@ -2,7 +2,7 @@
 <?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
 <package format="3">
   <name>yolov8_msgs</name>
-  <version>0.0.0</version>
+  <version>3.2.1</version>
   <description>Msgs for YOLOv8</description>
   <maintainer email="[email protected]">Miguel Ángel González Santamarta</maintainer>
   <license>GPL-3</license>

yolov8_ros/package.xml

@@ -2,7 +2,7 @@
 <?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
 <package format="3">
   <name>yolov8_ros</name>
-  <version>0.0.0</version>
+  <version>3.2.1</version>
   <description>YOLOv8 for ROS 2</description>
   <maintainer email="[email protected]">Miguel Ángel González Santamarta</maintainer>
   <license>GPL-3</license>

yolov8_ros/yolov8_ros/debug_node.py

@@ -124,8 +124,30 @@ def draw_box(self, cv_image: np.ndarray, detection: Detection, color: Tuple[int]
         max_pt = (round(box_msg.center.position.x + box_msg.size.x / 2.0),
                   round(box_msg.center.position.y + box_msg.size.y / 2.0))
 
-        # draw box
-        cv2.rectangle(cv_image, min_pt, max_pt, color, 2)
+        # define the four corners of the rectangle
+        rect_pts = np.array([
+            [min_pt[0], min_pt[1]],
+            [max_pt[0], min_pt[1]],
+            [max_pt[0], max_pt[1]],
+            [min_pt[0], max_pt[1]]
+        ])
+
+        # calculate the rotation matrix
+        rotation_matrix = cv2.getRotationMatrix2D(
+            (box_msg.center.position.x, box_msg.center.position.y),
+            -np.rad2deg(box_msg.center.theta),
+            1.0
+        )
+
+        # rotate the corners of the rectangle
+        rect_pts = np.int0(cv2.transform(
+            np.array([rect_pts]), rotation_matrix)[0])
+
+        # Draw the rotated rectangle
+        for i in range(4):
+            pt1 = tuple(rect_pts[i])
+            pt2 = tuple(rect_pts[(i + 1) % 4])
+            cv2.line(cv_image, pt1, pt2, color, 2)
 
         # write text
         label = "{} ({}) ({:.3f})".format(label, str(track_id), score)

yolov8_ros/yolov8_ros/yolov8_node.py

@@ -147,35 +147,61 @@ def parse_hypothesis(self, results: Results) -> List[Dict]:
 
         hypothesis_list = []
 
-        box_data: Boxes
-        for box_data in results.boxes:
-            hypothesis = {
-                "class_id": int(box_data.cls),
-                "class_name": self.yolo.names[int(box_data.cls)],
-                "score": float(box_data.conf)
-            }
-            hypothesis_list.append(hypothesis)
+        if results.boxes:
+            box_data: Boxes
+            for box_data in results.boxes:
+                hypothesis = {
+                    "class_id": int(box_data.cls),
+                    "class_name": self.yolo.names[int(box_data.cls)],
+                    "score": float(box_data.conf)
+                }
+                hypothesis_list.append(hypothesis)
+
+        elif results.obb:
+            for i in range(results.obb.cls.shape[0]):
+                hypothesis = {
+                    "class_id": int(results.obb.cls[i]),
+                    "class_name": self.yolo.names[int(results.obb.cls[i])],
+                    "score": float(results.obb.conf[i])
+                }
+                hypothesis_list.append(hypothesis)
 
         return hypothesis_list
 
     def parse_boxes(self, results: Results) -> List[BoundingBox2D]:
 
         boxes_list = []
 
-        box_data: Boxes
-        for box_data in results.boxes:
+        if results.boxes:
+            box_data: Boxes
+            for box_data in results.boxes:
 
-            msg = BoundingBox2D()
+                msg = BoundingBox2D()
 
-            # get boxes values
-            box = box_data.xywh[0]
-            msg.center.position.x = float(box[0])
-            msg.center.position.y = float(box[1])
-            msg.size.x = float(box[2])
-            msg.size.y = float(box[3])
+                # get boxes values
+                box = box_data.xywh[0]
+                msg.center.position.x = float(box[0])
+                msg.center.position.y = float(box[1])
+                msg.size.x = float(box[2])
+                msg.size.y = float(box[3])
 
-            # append msg
-            boxes_list.append(msg)
+                # append msg
+                boxes_list.append(msg)
+
+        elif results.obb:
+            for i in range(results.obb.cls.shape[0]):
+                msg = BoundingBox2D()
+
+                # get boxes values
+                box = results.obb.xywhr[i]
+                msg.center.position.x = float(box[0])
+                msg.center.position.y = float(box[1])
+                msg.center.theta = float(box[4])
+                msg.size.x = float(box[2])
+                msg.size.y = float(box[3])
+
+                # append msg
+                boxes_list.append(msg)
 
         return boxes_list
 
@@ -246,7 +272,7 @@ def image_cb(self, msg: Image) -> None:
             )
             results: Results = results[0].cpu()
 
-            if results.boxes:
+            if results.boxes or results.obb:
                 hypothesis = self.parse_hypothesis(results)
                 boxes = self.parse_boxes(results)
 
@@ -263,7 +289,7 @@ def image_cb(self, msg: Image) -> None:
 
                 aux_msg = Detection()
 
-                if results.boxes:
+                if results.boxes or results.obb:
                     aux_msg.class_id = hypothesis[i]["class_id"]
                     aux_msg.class_name = hypothesis[i]["class_name"]
                     aux_msg.score = hypothesis[i]["score"]