diff --git a/rcdt_detection/launch/detect.launch.py b/rcdt_detection/launch/detect.launch.py index 3ee90e4..59a7a9c 100644 --- a/rcdt_detection/launch/detect.launch.py +++ b/rcdt_detection/launch/detect.launch.py @@ -10,21 +10,30 @@ from launch_ros.substitutions import FindPackageShare -def generate_launch_description(): - return LaunchDescription([ - Node( - package="rcdt_detection", - executable="object_detection.py", - ), - IncludeLaunchDescription( - PythonLaunchDescriptionSource([ - PathJoinSubstitution([ - FindPackageShare("realsense2_camera"), "launch", "rs_launch.py"]) - ]), - launch_arguments={ - "align_depth.enable": "true", - "enable_sync": "true", - "enable_rgbd": "true", - }.items(), - ), - ]) +def generate_launch_description() -> LaunchDescription: + return LaunchDescription( + [ + Node( + package="rcdt_detection", + executable="object_detection.py", + ), + IncludeLaunchDescription( + PythonLaunchDescriptionSource( + [ + PathJoinSubstitution( + [ + FindPackageShare("realsense2_camera"), + "launch", + "rs_launch.py", + ] + ) + ] + ), + launch_arguments={ + "align_depth.enable": "true", + "enable_sync": "true", + "enable_rgbd": "true", + }.items(), + ), + ] + ) diff --git a/rcdt_detection/nodes/object_detection.py b/rcdt_detection/nodes/object_detection.py index a2f2dec..9061353 100644 --- a/rcdt_detection/nodes/object_detection.py +++ b/rcdt_detection/nodes/object_detection.py @@ -37,6 +37,7 @@ @dataclasses.dataclass class ImageCoordinate: """Image coordinate where (0, 0) is at the top-left corner.""" + x: int y: int @@ -46,13 +47,17 @@ def __init__(self, is_object_of_interest: Callable): super().__init__("detection_service") self.is_object_of_interest: Callable = is_object_of_interest - self.segmentation_model: ultralytics.engine.model.Model = load_segmentation_model() + self.segmentation_model: ultralytics.engine.model.Model = ( + load_segmentation_model() + ) self.intrinsics: rs2.intrinsics | None = None self.rgbd_topic: str = "/camera/camera/rgbd" self.time_to_wait: float = 2.0 self.max_attempts: int = 3 - self.service = self.create_service(DetectObjects, "detect_objects", self.detection_callback) + self.service = self.create_service( + DetectObjects, "detect_objects", self.detection_callback + ) def detection_callback(self, _: None, response: PointList) -> PointList: """Gets RGBD image, detects objects, and returns them.""" @@ -65,7 +70,7 @@ def detection_callback(self, _: None, response: PointList) -> PointList: centroid_image_coordinates = process_rgb_image( message=message.rgb, - segmentation_model=self.segmentation_model, + segmentation_model=self.segmentation_model, is_object_of_interest=self.is_object_of_interest, ) ros_logger.info(f"{centroid_image_coordinates}") @@ -75,7 +80,7 @@ def detection_callback(self, _: None, response: PointList) -> PointList: world_coordinates = process_depth_image( message=message.depth, image_coordinates=centroid_image_coordinates, - intrinsics=self.intrinsics + intrinsics=self.intrinsics, ) ros_logger.info(f"{world_coordinates}") response.objects = PointList(points=world_coordinates) @@ -84,10 +89,10 @@ def detection_callback(self, _: None, response: PointList) -> PointList: def get_rgbd(self) -> RGBD: """Gets RGBD-message from the rgbd topic. - + Raises: LookupError: if after `self.max_attempts` no message could be received. - + """ attempts = 0 while attempts < self.max_attempts: @@ -98,11 +103,13 @@ def get_rgbd(self) -> RGBD: time_to_wait=self.time_to_wait, ) if success: - break + break attempts += 1 else: raise LookupError("unable to get message from RGBD topic") - ros_logger.info(f"Received RGBD message for frame '{message.header.frame_id}' stamped at {repr(message.header.stamp)}.") + ros_logger.info( + f"Received RGBD message for frame '{message.header.frame_id}' stamped at {repr(message.header.stamp)}." + ) return message @@ -122,7 +129,7 @@ def calculate_intrinsics(message: CameraInfo) -> rs2.intrinsics: elif message.distortion_model == "equidistant": intrinsics.model = rs2.distortion.kannala_brandt4 - intrinsics.coeffs = [i for i in message.d] + intrinsics.coeffs = list(message.d) return intrinsics @@ -137,7 +144,7 @@ def process_rgb_image( Object detection starts with segmenting the image. The result contains segmentation masks per object. For every segmented object it is determined whether it is an object of interest. If it is, the image coordinates of its centroid are calculated. - + """ rgb_image = ros_image_to_cv2_image(message) @@ -158,14 +165,20 @@ def process_rgb_image( return centroid_image_coordinates -def process_depth_image(message: Image, image_coordinates: Iterable[ImageCoordinate], intrinsics: rs2.intrinsics) -> list[Point]: +def process_depth_image( + message: Image, + image_coordinates: Iterable[ImageCoordinate], + intrinsics: rs2.intrinsics, +) -> list[Point]: """Calculate world coordinate relative to the camera of image coordinates.""" depth_image = ros_image_to_cv2_image(message) world_coordinates = [] for image_coordinate in image_coordinates: depth_value = depth_image[image_coordinate.y, image_coordinate.x] - x, y, z = rs2.rs2_deproject_pixel_to_point(intrinsics, [image_coordinate.x, image_coordinate.y], depth_value) + x, y, z = rs2.rs2_deproject_pixel_to_point( + intrinsics, [image_coordinate.x, image_coordinate.y], depth_value + ) world_coordinates.append(Point(x=x, y=y, z=z)) return world_coordinates @@ -202,4 +215,4 @@ def main(args: str = None) -> None: # rosdep update # rosdep install --from-paths src --ignore-src --reinstall -y # source install/setup.bash -# WARNING: The scripts f2py and numpy-config are installed in '/home/wp/.local/bin' which is not on PATH. \ No newline at end of file +# WARNING: The scripts f2py and numpy-config are installed in '/home/wp/.local/bin' which is not on PATH. diff --git a/rcdt_detection/rcdt_detection/__init__.py b/rcdt_detection/rcdt_detection/__init__.py index eace54f..cc29ca9 100644 --- a/rcdt_detection/rcdt_detection/__init__.py +++ b/rcdt_detection/rcdt_detection/__init__.py @@ -1,3 +1,3 @@ # SPDX-FileCopyrightText: Alliander N. V. # -# SPDX-License-Identifier: Apache-2.0 \ No newline at end of file +# SPDX-License-Identifier: Apache-2.0 diff --git a/rcdt_detection/rcdt_detection/image_manipulation.py b/rcdt_detection/rcdt_detection/image_manipulation.py index c645cfc..1c7a8a8 100644 --- a/rcdt_detection/rcdt_detection/image_manipulation.py +++ b/rcdt_detection/rcdt_detection/image_manipulation.py @@ -11,7 +11,9 @@ cv_bridge = CvBridge() -def ros_image_to_cv2_image(image_message: Image, desired_encoding: str = "passthrough") -> np.array: +def ros_image_to_cv2_image( + image_message: Image, desired_encoding: str = "passthrough" +) -> np.array: """Convert ROS image message to cv2 image.""" return cv_bridge.imgmsg_to_cv2(image_message, desired_encoding=desired_encoding) diff --git a/rcdt_detection/rcdt_detection/object_detectors.py b/rcdt_detection/rcdt_detection/object_detectors.py index 2b25c18..d493f0a 100644 --- a/rcdt_detection/rcdt_detection/object_detectors.py +++ b/rcdt_detection/rcdt_detection/object_detectors.py @@ -7,11 +7,11 @@ def is_brick(mask: np.array) -> bool: """Determine if given mask belongs to a brick using mean pixel thresholds.""" - b, g, r = np.nanmean(np.where(mask==0.0, np.nan, mask), axis=(0, 1)) - return (b > (r + 100.0)) and (b > (g + 10.0)) # (85.0 < r < 110.0) and (100.0 < g < 130.0) and (100.0 < b < 140.0) + b, g, r = np.nanmean(np.where(mask == 0.0, np.nan, mask), axis=(0, 1)) + return (b > (r + 100.0)) and (b > (g + 10.0)) def is_blue(mask: np.array) -> bool: """Determine if given mask belongs to a blue object using mean pixel thresholds.""" - b, g, r = np.nanmean(np.where(mask==0.0, np.nan, mask), axis=(0, 1)) + b, g, r = np.nanmean(np.where(mask == 0.0, np.nan, mask), axis=(0, 1)) return (b > (r - 5.0)) and (b > (g - 5.0)) diff --git a/rcdt_detection/rcdt_detection/segmentation.py b/rcdt_detection/rcdt_detection/segmentation.py index dd9b83f..6fc5ac3 100644 --- a/rcdt_detection/rcdt_detection/segmentation.py +++ b/rcdt_detection/rcdt_detection/segmentation.py @@ -8,12 +8,16 @@ SEGMENTATION_MODEL_PATH: str = "/home/wp/projects/realsense/models/FastSAM-x.pt" -def load_segmentation_model(model_path: str = SEGMENTATION_MODEL_PATH) -> ultralytics.engine.model.Model: +def load_segmentation_model( + model_path: str = SEGMENTATION_MODEL_PATH, +) -> ultralytics.engine.model.Model: """Load segmentation model from given path.""" return ultralytics.FastSAM(model_path) -def segment_image(model: ultralytics.engine.model.Model, image: np.array) -> ultralytics.engine.results.Results: +def segment_image( + model: ultralytics.engine.model.Model, image: np.array +) -> ultralytics.engine.results.Results: """Segment given image using given model.""" height, width, _ = image.shape return model(image, imgsz=(height, width))[0]