diff --git a/modules/docker-compose.perception.yaml b/modules/docker-compose.perception.yaml index f095ff37..7df22ac2 100644 --- a/modules/docker-compose.perception.yaml +++ b/modules/docker-compose.perception.yaml @@ -33,7 +33,7 @@ services: volumes: - /mnt/wato-drive2/perception_models/yolov8m.pt:/perception_models/yolov8m.pt - /mnt/wato-drive2/perception_models/traffic_light.pt:/perception_models/traffic_light.pt - - /mnt/wato-drive2/perception_models/traffic_signs_v0.pt:/perception_models/traffic_signs_v1.pt + - /mnt/wato-drive2/perception_models/traffic_signs_v3.pt:/perception_models/traffic_signs.pt lidar_object_detection: build: diff --git a/src/perception/camera_object_detection/camera_object_detection/yolov8_detection.py b/src/perception/camera_object_detection/camera_object_detection/yolov8_detection.py index 05ceb431..67ddf741 100755 --- a/src/perception/camera_object_detection/camera_object_detection/yolov8_detection.py +++ b/src/perception/camera_object_detection/camera_object_detection/yolov8_detection.py @@ -24,18 +24,31 @@ import torch +class Model(): + def __init__(self, model_path, device): + self.model_path = model_path + self.model = AutoBackend(self.model_path, device=device, dnn=False, fp16=False) + self.names = self.model.module.names if hasattr(self.model, "module") else self.model.names + self.stride = int(self.model.stride) + + class CameraDetectionNode(Node): def __init__(self): torch.zeros(1).cuda() - super().__init__("camera_object_detection_node") + super().__init__("left_combined_detection_node") self.get_logger().info("Creating camera detection node...") self.declare_parameter("camera_topic", "/camera/right/image_color") self.declare_parameter("publish_vis_topic", "/annotated_img") self.declare_parameter("publish_detection_topic", "/detections") - self.declare_parameter("model_path", "/perception_models/yolov8m.pt") + self.declare_parameter("models.traffic_signs.model_path", + "/perception_models/traffic_signs.pt") + self.declare_parameter("models.traffic_light.model_path", + "/perception_models/traffic_light.pt") + self.declare_parameter("models.pretrained_yolov8m.model_path", + "/perception_models/yolov8m.pt") self.declare_parameter("image_size", 1024) self.declare_parameter("compressed", False) self.declare_parameter("crop_mode", "LetterBox") @@ -44,7 +57,11 @@ def __init__(self): self.camera_topic = self.get_parameter("camera_topic").value self.publish_vis_topic = self.get_parameter("publish_vis_topic").value self.publish_detection_topic = self.get_parameter("publish_detection_topic").value - self.model_path = self.get_parameter("model_path").value + self.model_paths = [ + self.get_parameter("models.traffic_signs.model_path").value, + self.get_parameter("models.traffic_light.model_path").value, + self.get_parameter("models.pretrained_yolov8m.model_path").value + ] self.image_size = self.get_parameter("image_size").value self.compressed = self.get_parameter("compressed").value self.crop_mode = self.get_parameter("crop_mode").value @@ -82,11 +99,7 @@ def __init__(self): self.cv_bridge = CvBridge() # load yolov8 model - self.model = AutoBackend(self.model_path, device=self.device, dnn=False, fp16=False) - - self.names = self.model.module.names if hasattr(self.model, "module") else self.model.names - - self.stride = int(self.model.stride) + self.models = self.load_models() # setup vis publishers self.vis_publisher = self.create_publisher(Image, self.publish_vis_topic, 10) @@ -98,9 +111,16 @@ def __init__(self): f"Successfully created node listening on camera topic: {self.camera_topic}..." ) - def crop_image(self, cv_image): + def load_models(self): + models = [] + for model_path in self.model_paths: + if model_path: + models.append(Model(model_path, self.device)) + return models + + def crop_image(self, cv_image, model): if self.crop_mode == "LetterBox": - img = LetterBox(self.image_size, stride=self.stride)(image=cv_image) + img = LetterBox(self.image_size, stride=model.stride)(image=cv_image) elif self.crop_mode == "CenterCrop": img = CenterCrop(self.image_size)(cv_image) else: @@ -144,7 +164,7 @@ def convert_bboxes_to_orig_frame(self, bbox): bbox[3] * height_scale, ] - def crop_and_convert_to_tensor(self, cv_image): + def crop_and_convert_to_tensor(self, cv_image, model): """ Preprocess the image by resizing, padding and rearranging the dimensions. @@ -154,7 +174,7 @@ def crop_and_convert_to_tensor(self, cv_image): Returns: torch.Tensor image for model input of shape (1,3,w,h) """ - img = self.crop_image(cv_image) + img = self.crop_image(cv_image, model) # Convert img = img.transpose(2, 0, 1) @@ -250,41 +270,42 @@ def image_callback(self, msg): self.get_logger().error(str(e)) return - # preprocess image and run through prediction - img = self.crop_and_convert_to_tensor(cv_image) - pred = self.model(img) - - # nms function used same as yolov8 detect.py - pred = non_max_suppression(pred) detections = [] - for i, det in enumerate(pred): # per image - if len(det): - # Write results - for *xyxy, conf, cls in reversed(det): - label = self.names[int(cls)] - - bbox = [ - xyxy[0], - xyxy[1], - xyxy[2] - xyxy[0], - xyxy[3] - xyxy[1], - ] - bbox = [b.item() for b in bbox] - bbox = self.convert_bboxes_to_orig_frame(bbox) - - detections.append( - { - "label": label, - "conf": conf.item(), - "bbox": bbox, - } - ) - self.get_logger().debug(f"{label}: {bbox}") + for model in self.models: + # preprocess image and run through prediction + img = self.crop_and_convert_to_tensor(cv_image, model) + pred = model.model(img) + + # nms function used same as yolov8 detect.py + pred = non_max_suppression(pred) + for i, det in enumerate(pred): # per image + if len(det): + # Write results + for *xyxy, conf, cls in reversed(det): + label = model.names[int(cls)] + + bbox = [ + xyxy[0], + xyxy[1], + xyxy[2] - xyxy[0], + xyxy[3] - xyxy[1], + ] + bbox = [b.item() for b in bbox] + bbox = self.convert_bboxes_to_orig_frame(bbox) + + detections.append( + { + "label": label, + "conf": conf.item(), + "bbox": bbox, + } + ) + self.get_logger().debug(f"{label}: {bbox}") annotator = Annotator( cv_image, line_width=self.line_thickness, - example=str(self.names), + # example=str(model.names), ) (detections, annotated_img) = self.postprocess_detections(detections, annotator) diff --git a/src/perception/camera_object_detection/config/combined_config.yaml b/src/perception/camera_object_detection/config/combined_config.yaml new file mode 100644 index 00000000..422658cf --- /dev/null +++ b/src/perception/camera_object_detection/config/combined_config.yaml @@ -0,0 +1,56 @@ +left_combined_detection_node: + ros__parameters: + camera_topic: /camera/left/image_color + publish_vis_topic: /camera/left/combined_detection_viz + publish_detection_topic: /camera/left/combined_detection + models: + pretrained_yolov8m: + name: yolov8m + model_path: /perception_models/yolov8m.pt + traffic_light: + name: traffic_light + model_path: /perception_models/traffic_light.pt + traffic_signs: + name: traffic_signs + model_path: /perception_models/traffic_signs.pt + crop_mode: CenterCrop + image_size: 1024 + save_detections: false + +center_combined_detection_node: + ros__parameters: + camera_topic: /camera/center/image_color + publish_vis_topic: /camera/center/combined_detection_viz + publish_detection_topic: /camera/center/combined_detection + models: + pretrained_yolov8m: + name: yolov8m + model_path: /perception_models/yolov8m.pt + traffic_light: + name: traffic_light + model_path: /perception_models/traffic_light.pt + traffic_signs: + name: traffic_signs + model_path: /perception_models/traffic_signs.pt + crop_mode: CenterCrop + image_size: 1024 + save_detections: false + +right_combined_detection_node: + ros__parameters: + camera_topic: /camera/right/image_color + publish_vis_topic: /camera/right/combined_detection_viz + publish_detection_topic: /camera/right/combined_detection + models: + pretrained_yolov8m: + name: yolov8m + model_path: /perception_models/yolov8m.pt + traffic_light: + name: traffic_light + model_path: /perception_models/traffic_light.pt + traffic_signs: + name: traffic_signs + model_path: /perception_models/traffic_signs.pt + crop_mode: CenterCrop + image_size: 1024 + save_detections: false \ No newline at end of file diff --git a/src/perception/camera_object_detection/config/traffic_signs_config.yaml b/src/perception/camera_object_detection/config/traffic_signs_config.yaml index 8a143a77..ac312d1a 100644 --- a/src/perception/camera_object_detection/config/traffic_signs_config.yaml +++ b/src/perception/camera_object_detection/config/traffic_signs_config.yaml @@ -3,7 +3,7 @@ traffic_signs_node: camera_topic: /camera/left/image_color publish_vis_topic: /traffic_signs_viz publish_detection_topic: /traffic_signs - model_path: /perception_models/traffic_signs_v1.pt + model_path: /perception_models/traffic_signs.pt crop_mode: CenterCrop image_size: 1024 save_detections: false diff --git a/src/perception/camera_object_detection/launch/eve.launch.py b/src/perception/camera_object_detection/launch/eve.launch.py index 9ff4992e..b840698b 100755 --- a/src/perception/camera_object_detection/launch/eve.launch.py +++ b/src/perception/camera_object_detection/launch/eve.launch.py @@ -3,55 +3,43 @@ from launch.conditions import LaunchConfigurationEquals from launch.actions import DeclareLaunchArgument, IncludeLaunchDescription from launch.launch_description_sources import PythonLaunchDescriptionSource +from launch_ros.actions import Node from ament_index_python.packages import get_package_share_directory import os def generate_launch_description(): - launch_traffic_light = LaunchConfiguration("launch_traffic_light", default=True) - launch_traffic_light_arg = DeclareLaunchArgument( - "launch_traffic_light", - default_value=launch_traffic_light, - description="Launch traffic light detection", - ) - launch_traffic_signs = LaunchConfiguration("launch_traffic_signs", default=True) - launch_traffic_signs_arg = DeclareLaunchArgument( - "launch_traffic_signs", - default_value=launch_traffic_signs, - description="Launch traffic signs detection", - ) - - launch_args = [launch_traffic_light_arg, launch_traffic_signs_arg] - - camera_object_detection_launch_include_dir = os.path.join( - get_package_share_directory("camera_object_detection"), "launch", "include" + config = os.path.join( + get_package_share_directory("camera_object_detection"), + "config", + "combined_config.yaml" ) - pretrained_yolov8_launch = IncludeLaunchDescription( - PythonLaunchDescriptionSource( - [camera_object_detection_launch_include_dir, "/pretrained_yolov8.launch.py"] - ), + left_combined_detection_node = Node( + package="camera_object_detection", + executable="camera_object_detection_node", + name="left_combined_detection_node", + parameters=[config], ) - traffic_light_launch = IncludeLaunchDescription( - PythonLaunchDescriptionSource( - [camera_object_detection_launch_include_dir, "/traffic_light.launch.py"] - ), - condition=LaunchConfigurationEquals("launch_traffic_light", "True"), + center_combined_detection_node = Node( + package="camera_object_detection", + executable="camera_object_detection_node", + name="center_combined_detection_node", + parameters=[config], ) - traffic_signs_launch = IncludeLaunchDescription( - PythonLaunchDescriptionSource( - [camera_object_detection_launch_include_dir, "/traffic_signs.launch.py"] - ), - condition=LaunchConfigurationEquals("launch_traffic_signs", "True"), + right_combined_detection_node = Node( + package="camera_object_detection", + executable="camera_object_detection_node", + name="right_combined_detection_node", + parameters=[config], ) return LaunchDescription( - launch_args - + [ - pretrained_yolov8_launch, - traffic_light_launch, - traffic_signs_launch, + [ + left_combined_detection_node, + center_combined_detection_node, + right_combined_detection_node ] ) diff --git a/src/samples/python/aggregator/setup.py b/src/samples/python/aggregator/setup.py index b0afb9f6..f77c1804 100755 --- a/src/samples/python/aggregator/setup.py +++ b/src/samples/python/aggregator/setup.py @@ -14,7 +14,7 @@ # Include our package.xml file (os.path.join('share', package_name), ['package.xml']), # Include all launch files. - (os.path.join('share', package_name, 'launch'), \ + (os.path.join('share', package_name, 'launch'), glob(os.path.join('launch', '*.launch.py'))), ], install_requires=['setuptools'],