Skip to content

Commit 8a2d743

Browse files
committed
yolo_encoding param added
1 parent a52adf7 commit 8a2d743

File tree

3 files changed

+19
-5
lines changed

3 files changed

+19
-5
lines changed

README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ These are the parameters from the [yolo.launch.py](./yolo_bringup/launch/yolo.la
125125
- **model**: YOLO model (default: yolov8m.pt)
126126
- **tracker**: tracker file (default: bytetrack.yaml)
127127
- **device**: GPU/CUDA (default: cuda:0)
128+
- **yolo_encoding**: Encoding to convert input image before using YOLO (default: bgr8)
128129
- **enable**: whether to start YOLO enabled (default: True)
129130
- **threshold**: detection threshold (default: 0.5)
130131
- **iou**: intersection Over Union (IoU) threshold for Non-Maximum Suppression (NMS) (default: 0.7)

yolo_bringup/launch/yolo.launch.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,13 @@ def run_yolo(context: LaunchContext, use_tracking, use_3d):
5757
description="Device to use (GPU/CPU)",
5858
)
5959

60+
yolo_encoding = LaunchConfiguration("yolo_encoding")
61+
yolo_encoding_cmd = DeclareLaunchArgument(
62+
"yolo_encoding",
63+
default_value="bgr8",
64+
description="Encoding of the input image topic",
65+
)
66+
6067
enable = LaunchConfiguration("enable")
6168
enable_cmd = DeclareLaunchArgument(
6269
"enable",
@@ -229,6 +236,7 @@ def run_yolo(context: LaunchContext, use_tracking, use_3d):
229236
"model_type": model_type,
230237
"model": model,
231238
"device": device,
239+
"yolo_encoding": yolo_encoding,
232240
"enable": enable,
233241
"threshold": threshold,
234242
"iou": iou,
@@ -295,6 +303,7 @@ def run_yolo(context: LaunchContext, use_tracking, use_3d):
295303
model_cmd,
296304
tracker_cmd,
297305
device_cmd,
306+
yolo_encoding_cmd,
298307
enable_cmd,
299308
threshold_cmd,
300309
iou_cmd,

yolo_ros/yolo_ros/yolo_node.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ def __init__(self) -> None:
5555
self.declare_parameter("model_type", "YOLO")
5656
self.declare_parameter("model", "yolov8m.pt")
5757
self.declare_parameter("device", "cuda:0")
58+
self.declare_parameter("yolo_encoding", "bgr8")
59+
self.declare_parameter("enable", True)
60+
self.declare_parameter("image_reliability", QoSReliabilityPolicy.BEST_EFFORT)
5861

5962
self.declare_parameter("threshold", 0.5)
6063
self.declare_parameter("iou", 0.5)
@@ -66,9 +69,6 @@ def __init__(self) -> None:
6669
self.declare_parameter("agnostic_nms", False)
6770
self.declare_parameter("retina_masks", False)
6871

69-
self.declare_parameter("enable", True)
70-
self.declare_parameter("image_reliability", QoSReliabilityPolicy.BEST_EFFORT)
71-
7272
self.type_to_model = {"YOLO": YOLO, "World": YOLOWorld}
7373

7474
def on_configure(self, state: LifecycleState) -> TransitionCallbackReturn:
@@ -80,6 +80,9 @@ def on_configure(self, state: LifecycleState) -> TransitionCallbackReturn:
8080
)
8181
self.model = self.get_parameter("model").get_parameter_value().string_value
8282
self.device = self.get_parameter("device").get_parameter_value().string_value
83+
self.yolo_encoding = (
84+
self.get_parameter("yolo_encoding").get_parameter_value().string_value
85+
)
8386

8487
# inference params
8588
self.threshold = (
@@ -326,8 +329,9 @@ def image_cb(self, msg: Image) -> None:
326329
if self.enable:
327330

328331
# convert image + predict
329-
cv_image = self.cv_bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
330-
cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
332+
cv_image = self.cv_bridge.imgmsg_to_cv2(
333+
msg, desired_encoding=self.yolo_encoding
334+
)
331335
results = self.yolo.predict(
332336
source=cv_image,
333337
verbose=False,

0 commit comments

Comments
 (0)