Skip to content

Commit efed1df

Browse files
mgonzs13mgonzs13
committed
OBB support
1 parent 19da2b4 commit efed1df

File tree

6 files changed

+75
-27
lines changed

6 files changed

+75
-27
lines changed

‎requirements.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
opencv-python==4.8.1.78
22
typing-extensions>=4.4.0
3-
ultralytics==8.2.45
3+
ultralytics==8.2.62
44
lap==0.4.0

‎yolov8_bringup/package.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
33
<package format="3">
44
<name>yolov8_bringup</name>
5-
<version>0.0.0</version>
5+
<version>3.2.1</version>
66
<description>YOLOv8 for ROS 2</description>
77
<maintainer email="mgons@unileon.es">Miguel Ángel González Santamarta</maintainer>
88
<license>GPL-3</license>

‎yolov8_msgs/package.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
33
<package format="3">
44
<name>yolov8_msgs</name>
5-
<version>0.0.0</version>
5+
<version>3.2.1</version>
66
<description>Msgs for YOLOv8</description>
77
<maintainer email="mgons@unileon.es">Miguel Ángel González Santamarta</maintainer>
88
<license>GPL-3</license>

‎yolov8_ros/package.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
<?xml-model href="http://download.ros.org/schema/package_format3.xsd" schematypens="http://www.w3.org/2001/XMLSchema"?>
33
<package format="3">
44
<name>yolov8_ros</name>
5-
<version>0.0.0</version>
5+
<version>3.2.1</version>
66
<description>YOLOv8 for ROS 2</description>
77
<maintainer email="mgons@unileon.es">Miguel Ángel González Santamarta</maintainer>
88
<license>GPL-3</license>

‎yolov8_ros/yolov8_ros/debug_node.py

+24-2
Original file line numberDiff line numberDiff line change
@@ -124,8 +124,30 @@ def draw_box(self, cv_image: np.ndarray, detection: Detection, color: Tuple[int]
124124
max_pt = (round(box_msg.center.position.x + box_msg.size.x / 2.0),
125125
round(box_msg.center.position.y + box_msg.size.y / 2.0))
126126

127-
# draw box
128-
cv2.rectangle(cv_image, min_pt, max_pt, color, 2)
127+
# define the four corners of the rectangle
128+
rect_pts = np.array([
129+
[min_pt[0], min_pt[1]],
130+
[max_pt[0], min_pt[1]],
131+
[max_pt[0], max_pt[1]],
132+
[min_pt[0], max_pt[1]]
133+
])
134+
135+
# calculate the rotation matrix
136+
rotation_matrix = cv2.getRotationMatrix2D(
137+
(box_msg.center.position.x, box_msg.center.position.y),
138+
-np.rad2deg(box_msg.center.theta),
139+
1.0
140+
)
141+
142+
# rotate the corners of the rectangle
143+
rect_pts = np.int0(cv2.transform(
144+
np.array([rect_pts]), rotation_matrix)[0])
145+
146+
# Draw the rotated rectangle
147+
for i in range(4):
148+
pt1 = tuple(rect_pts[i])
149+
pt2 = tuple(rect_pts[(i + 1) % 4])
150+
cv2.line(cv_image, pt1, pt2, color, 2)
129151

130152
# write text
131153
label = "{} ({}) ({:.3f})".format(label, str(track_id), score)

‎yolov8_ros/yolov8_ros/yolov8_node.py

+47-21
Original file line numberDiff line numberDiff line change
@@ -147,35 +147,61 @@ def parse_hypothesis(self, results: Results) -> List[Dict]:
147147

148148
hypothesis_list = []
149149

150-
box_data: Boxes
151-
for box_data in results.boxes:
152-
hypothesis = {
153-
"class_id": int(box_data.cls),
154-
"class_name": self.yolo.names[int(box_data.cls)],
155-
"score": float(box_data.conf)
156-
}
157-
hypothesis_list.append(hypothesis)
150+
if results.boxes:
151+
box_data: Boxes
152+
for box_data in results.boxes:
153+
hypothesis = {
154+
"class_id": int(box_data.cls),
155+
"class_name": self.yolo.names[int(box_data.cls)],
156+
"score": float(box_data.conf)
157+
}
158+
hypothesis_list.append(hypothesis)
159+
160+
elif results.obb:
161+
for i in range(results.obb.cls.shape[0]):
162+
hypothesis = {
163+
"class_id": int(results.obb.cls[i]),
164+
"class_name": self.yolo.names[int(results.obb.cls[i])],
165+
"score": float(results.obb.conf[i])
166+
}
167+
hypothesis_list.append(hypothesis)
158168

159169
return hypothesis_list
160170

161171
def parse_boxes(self, results: Results) -> List[BoundingBox2D]:
162172

163173
boxes_list = []
164174

165-
box_data: Boxes
166-
for box_data in results.boxes:
175+
if results.boxes:
176+
box_data: Boxes
177+
for box_data in results.boxes:
167178

168-
msg = BoundingBox2D()
179+
msg = BoundingBox2D()
169180

170-
# get boxes values
171-
box = box_data.xywh[0]
172-
msg.center.position.x = float(box[0])
173-
msg.center.position.y = float(box[1])
174-
msg.size.x = float(box[2])
175-
msg.size.y = float(box[3])
181+
# get boxes values
182+
box = box_data.xywh[0]
183+
msg.center.position.x = float(box[0])
184+
msg.center.position.y = float(box[1])
185+
msg.size.x = float(box[2])
186+
msg.size.y = float(box[3])
176187

177-
# append msg
178-
boxes_list.append(msg)
188+
# append msg
189+
boxes_list.append(msg)
190+
191+
elif results.obb:
192+
for i in range(results.obb.cls.shape[0]):
193+
msg = BoundingBox2D()
194+
195+
# get boxes values
196+
box = results.obb.xywhr[i]
197+
msg.center.position.x = float(box[0])
198+
msg.center.position.y = float(box[1])
199+
msg.center.theta = float(box[4])
200+
msg.size.x = float(box[2])
201+
msg.size.y = float(box[3])
202+
203+
# append msg
204+
boxes_list.append(msg)
179205

180206
return boxes_list
181207

@@ -246,7 +272,7 @@ def image_cb(self, msg: Image) -> None:
246272
)
247273
results: Results = results[0].cpu()
248274

249-
if results.boxes:
275+
if results.boxes or results.obb:
250276
hypothesis = self.parse_hypothesis(results)
251277
boxes = self.parse_boxes(results)
252278

@@ -263,7 +289,7 @@ def image_cb(self, msg: Image) -> None:
263289

264290
aux_msg = Detection()
265291

266-
if results.boxes:
292+
if results.boxes or results.obb:
267293
aux_msg.class_id = hypothesis[i]["class_id"]
268294
aux_msg.class_name = hypothesis[i]["class_name"]
269295
aux_msg.score = hypothesis[i]["score"]

0 commit comments

Comments
 (0)